Feat/ollama integration

.gitignore

@@ -133,3 +133,6 @@ dist
 # Rollup build dir
 build
 .DS_Store
+
+# Wrangler temp directory
+.wrangler 

src/globals.ts

@@ -5,7 +5,9 @@ export const HEADER_KEYS: Record<string, string> = {
   RETRIES: `x-${POWERED_BY}-retry-count`,
   PROVIDER: `x-${POWERED_BY}-provider`,
   TRACE_ID: `x-${POWERED_BY}-trace-id`,
-  CACHE: `x-${POWERED_BY}-cache`
+  CACHE: `x-${POWERED_BY}-cache`,
+  FORWARD_HEADERS: `x-${POWERED_BY}-forward-headers`,
+  CUSTOM_HOST: `x-${POWERED_BY}-custom-host`
 }
 
 export const RESPONSE_HEADER_KEYS: Record<string, string> = {
@@ -33,6 +35,7 @@ export const MISTRAL_AI: string = "mistral-ai";
 export const DEEPINFRA: string = "deepinfra";
 export const STABILITY_AI: string = "stability-ai";
 export const NOMIC: string = "nomic";
+export const OLLAMA: string = "ollama";
 
 export const providersWithStreamingSupport = [OPEN_AI, AZURE_OPEN_AI, ANTHROPIC, COHERE];
 export const allowedProxyProviders = [OPEN_AI, COHERE, AZURE_OPEN_AI, ANTHROPIC];

src/handlers/handlerUtils.ts

@@ -1,5 +1,5 @@
 import { Context } from "hono";
-import { AZURE_OPEN_AI, CONTENT_TYPES, GOOGLE, HEADER_KEYS, PALM, POWERED_BY, RESPONSE_HEADER_KEYS, RETRY_STATUS_CODES, STABILITY_AI } from "../globals";
+import { ANTHROPIC, AZURE_OPEN_AI, CONTENT_TYPES, GOOGLE, HEADER_KEYS, OLLAMA, PALM, POWERED_BY, RESPONSE_HEADER_KEYS, RETRY_STATUS_CODES, STABILITY_AI } from "../globals";
 import Providers from "../providers";
 import { ProviderAPIConfig, endpointStrings } from "../providers/types";
 import transformToProviderRequest from "../services/transformToProviderRequest";
@@ -19,13 +19,23 @@ import { OpenAICompleteJSONToStreamResponseTransform } from "../providers/openai
  * @param {string} method - The HTTP method for the request.
  * @returns {RequestInit} - The fetch options for the request.
  */
-export function constructRequest(headers: any, provider: string = "", method: string = "POST") {
+export function constructRequest(providerConfigMappedHeaders: any, provider: string, method: string, forwardHeaders: string[], requestHeaders: Record<string, string>) {
   let baseHeaders: any = {
     "content-type": "application/json"
   };
 
+  let headers: Record<string, string> = {
+    ...providerConfigMappedHeaders
+  };
+
+  const forwardHeadersMap: Record<string, string> = {};
+
+  forwardHeaders.forEach((h: string) => {
+    if (requestHeaders[h]) forwardHeadersMap[h] = requestHeaders[h];
+  })
+
   // Add any headers that the model might need
-  headers = {...baseHeaders, ...headers}
+  headers = {...baseHeaders, ...headers, ...forwardHeadersMap}
 
   let fetchOptions: RequestInit = {
     method,
@@ -122,7 +132,8 @@ export const fetchProviderOptionsFromConfig = (config: Config | ShortConfig): Op
       virtualKey: camelCaseConfig.virtualKey, 
       apiKey: camelCaseConfig.apiKey,
       cache: camelCaseConfig.cache,
-      retry: camelCaseConfig.retry
+      retry: camelCaseConfig.retry,
+      customHost: camelCaseConfig.customHost
       }];
       if (camelCaseConfig.resourceName) providerOptions[0].resourceName = camelCaseConfig.resourceName;
       if (camelCaseConfig.deploymentId) providerOptions[0].deploymentId = camelCaseConfig.deploymentId;
@@ -162,40 +173,48 @@ export async function tryPostProxy(c: Context, providerOption:Options, inputPara
   const apiConfig: ProviderAPIConfig = Providers[provider].api;
   let fetchOptions;
   let url = providerOption.urlToFetch as string;
-
+  
   let baseUrl:string, endpoint:string;
+
+  const forwardHeaders: string[] = [];
+  baseUrl = requestHeaders[HEADER_KEYS.CUSTOM_HOST] || providerOption.customHost || "";
+
   if (provider === AZURE_OPEN_AI && apiConfig.getBaseURL && apiConfig.getEndpoint) {
     // Construct the base object for the request
     if(!!providerOption.apiKey) {
-      fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey, "apiKey"), provider, method);
+      fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey, "apiKey"), provider, method, forwardHeaders, requestHeaders);
     } else {
-      fetchOptions = constructRequest(apiConfig.headers(providerOption.adAuth, "adAuth"), provider, method);
+      fetchOptions = constructRequest(apiConfig.headers(providerOption.adAuth, "adAuth"), provider, method, forwardHeaders, requestHeaders);
     }
-    baseUrl = apiConfig.getBaseURL(providerOption.resourceName, providerOption.deploymentId);
+    baseUrl = baseUrl || apiConfig.getBaseURL(providerOption.resourceName, providerOption.deploymentId);
     endpoint = apiConfig.getEndpoint(fn, providerOption.apiVersion, url);
     url = `${baseUrl}${endpoint}`;
   } else if (provider === PALM && apiConfig.baseURL && apiConfig.getEndpoint) {
-    fetchOptions = constructRequest(apiConfig.headers(), provider, method);
-    baseUrl = apiConfig.baseURL;
+    fetchOptions = constructRequest(apiConfig.headers(), provider, method, forwardHeaders, requestHeaders);
+    baseUrl = baseUrl || apiConfig.baseURL;
     endpoint = apiConfig.getEndpoint(fn, providerOption.apiKey, params?.model);
     url = `${baseUrl}${endpoint}`;
-  } else if (provider === "anthropic" && apiConfig.baseURL) {
+  } else if (provider === ANTHROPIC && apiConfig.baseURL) {
     // Construct the base object for the POST request
-    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey, fn), provider);
-    baseUrl = apiConfig.baseURL;
+    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey, fn), provider, "POST", forwardHeaders, requestHeaders);
+    baseUrl = baseUrl || apiConfig.baseURL;
     endpoint = apiConfig[fn] || "";
+    url = `${baseUrl}${endpoint}`;
   } else if (provider === GOOGLE && apiConfig.baseURL && apiConfig.getEndpoint) {
-    fetchOptions = constructRequest(apiConfig.headers(), provider);
-    baseUrl = apiConfig.baseURL;
+    fetchOptions = constructRequest(apiConfig.headers(), provider, "POST", forwardHeaders, requestHeaders);
+    baseUrl = baseUrl || apiConfig.baseURL;
     endpoint = apiConfig.getEndpoint(fn, providerOption.apiKey, params.model, params.stream);
+    url = `${baseUrl}${endpoint}`;
   } else if (provider === STABILITY_AI && apiConfig.baseURL && apiConfig.getEndpoint) {
-    fetchOptions = constructRequest(apiConfig.headers(), provider);
-    baseUrl = apiConfig.baseURL;
+    fetchOptions = constructRequest(apiConfig.headers(), provider, "POST", forwardHeaders, requestHeaders);
+    baseUrl = baseUrl || apiConfig.baseURL;
     endpoint = apiConfig.getEndpoint(fn, params.model, url);
+    url = `${baseUrl}${endpoint}`;
   } else {
     // Construct the base object for the request
-    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey), provider, method);
+    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey), provider, method, forwardHeaders, requestHeaders);
   }
+
   if (method === "POST") {
     fetchOptions.body = JSON.stringify(params)
   }
@@ -314,37 +333,45 @@ export async function tryPost(c: Context, providerOption:Options, inputParams: P
 
 
   let baseUrl:string, endpoint:string, fetchOptions;
+  const forwardHeaders = requestHeaders[HEADER_KEYS.FORWARD_HEADERS]?.split(",").map(h => h.trim()) || providerOption.forwardHeaders || [];
+  baseUrl = requestHeaders[HEADER_KEYS.CUSTOM_HOST] || providerOption.customHost || "";
+
   if (provider === AZURE_OPEN_AI && apiConfig.getBaseURL && apiConfig.getEndpoint) {
     // Construct the base object for the POST request
     if(!!providerOption.apiKey) {
-      fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey, "apiKey"), provider);
+      fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey, "apiKey"), provider, "POST", forwardHeaders, requestHeaders);
     } else {
-      fetchOptions = constructRequest(apiConfig.headers(providerOption.adAuth, "adAuth"), provider);
+      fetchOptions = constructRequest(apiConfig.headers(providerOption.adAuth, "adAuth"), provider, "POST", forwardHeaders, requestHeaders);
     }
-    baseUrl = apiConfig.getBaseURL(providerOption.resourceName, providerOption.deploymentId);
+    baseUrl = baseUrl || apiConfig.getBaseURL(providerOption.resourceName, providerOption.deploymentId);
     endpoint = apiConfig.getEndpoint(fn, providerOption.apiVersion);
   } else if (provider === PALM && apiConfig.baseURL && apiConfig.getEndpoint) {
-    fetchOptions = constructRequest(apiConfig.headers(), provider);
-    baseUrl = apiConfig.baseURL;
+    fetchOptions = constructRequest(apiConfig.headers(), provider, "POST", forwardHeaders, requestHeaders);
+    baseUrl = baseUrl || apiConfig.baseURL;
     endpoint = apiConfig.getEndpoint(fn, providerOption.apiKey, providerOption.overrideParams?.model || params?.model);
-  } else if (provider === "anthropic" && apiConfig.baseURL) {
+  } else if (provider === ANTHROPIC && apiConfig.baseURL) {
     // Construct the base object for the POST request
-    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey, fn), provider);
-    baseUrl = apiConfig.baseURL;
+    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey, fn), provider, "POST", forwardHeaders, requestHeaders);
+    baseUrl = baseUrl || apiConfig.baseURL;
     endpoint = apiConfig[fn] || "";
   } else if (provider === GOOGLE && apiConfig.baseURL && apiConfig.getEndpoint) {
-    fetchOptions = constructRequest(apiConfig.headers(), provider);
-    baseUrl = apiConfig.baseURL;
+    fetchOptions = constructRequest(apiConfig.headers(), provider, "POST", forwardHeaders, requestHeaders);
+    baseUrl = baseUrl || apiConfig.baseURL;
     endpoint = apiConfig.getEndpoint(fn, providerOption.apiKey, transformedRequestBody.model, params.stream);
   } else if (provider === STABILITY_AI && apiConfig.baseURL && apiConfig.getEndpoint) {
-    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey), provider);
-    baseUrl = apiConfig.baseURL;
+    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey), provider, "POST", forwardHeaders, requestHeaders);
+    baseUrl = baseUrl || apiConfig.baseURL;
     endpoint = apiConfig.getEndpoint(fn, params.model);
-  } else {
+  } else if (provider === OLLAMA && apiConfig.getEndpoint) {    
+    fetchOptions = constructRequest(apiConfig.headers(), provider, "POST", forwardHeaders, requestHeaders);
+    baseUrl = baseUrl;
+    endpoint = apiConfig.getEndpoint(fn, providerOption.apiKey, transformedRequestBody.model, params.stream);
+  } 
+  else {
     // Construct the base object for the POST request
-    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey), provider);
+    fetchOptions = constructRequest(apiConfig.headers(providerOption.apiKey), provider, "POST", forwardHeaders, requestHeaders);
 
-    baseUrl = apiConfig.baseURL || "";
+    baseUrl = baseUrl || apiConfig.baseURL || "";
     endpoint = apiConfig[fn] || "";
   }
 
@@ -358,7 +385,7 @@ export async function tryPost(c: Context, providerOption:Options, inputParams: P
 
   providerOption.retry = {
     attempts: providerOption.retry?.attempts ?? 0,
-    onStatusCodes: providerOption.retry?.onStatusCodes ?? []
+    onStatusCodes: providerOption.retry?.onStatusCodes ?? RETRY_STATUS_CODES
   }
 
   const [getFromCacheFunction, cacheIdentifier, requestOptions] = [
@@ -543,7 +570,7 @@ export async function tryTargetsRecursively(
     const strategyMode = currentTarget.strategy?.mode;
 
     // start: merge inherited config with current target config (preference given to current)
-    const currentInheritedConfig = {
+    const currentInheritedConfig: Record<string, any> = {
       overrideParams : {
         ...inheritedConfig.overrideParams,
         ...currentTarget.overrideParams
@@ -553,11 +580,29 @@ export async function tryTargetsRecursively(
       requestTimeout: null
     }
 
+    if (currentTarget.forwardHeaders) {
+      currentInheritedConfig.forwardHeaders = [...currentTarget.forwardHeaders];
+    } else if (inheritedConfig.forwardHeaders) {
+      currentInheritedConfig.forwardHeaders = [...inheritedConfig.forwardHeaders];
+      currentTarget.forwardHeaders = [...inheritedConfig.forwardHeaders];
+    }
+
+    if (currentTarget.customHost) {
+      currentInheritedConfig.customHost = currentTarget.customHost
+    } else if (inheritedConfig.customHost) {
+      currentInheritedConfig.customHost = inheritedConfig.customHost;
+      currentTarget.customHost = inheritedConfig.customHost;
+    }
+
     if (currentTarget.requestTimeout) {
       currentInheritedConfig.requestTimeout = currentTarget.requestTimeout
     } else if (inheritedConfig.requestTimeout) {
       currentInheritedConfig.requestTimeout = inheritedConfig.requestTimeout;
+      currentTarget.requestTimeout = inheritedConfig.requestTimeout;
     }
+
+
+
     currentTarget.overrideParams = {
         ...currentInheritedConfig.overrideParams
     }

src/handlers/proxyGetHandler.ts

@@ -10,11 +10,16 @@ function proxyProvider(proxyModeHeader:string, providerHeader: string) {
   return proxyProvider;
 }
 
-function getProxyPath(requestURL:string, proxyProvider:string, proxyEndpointPath:string) {
+function getProxyPath(requestURL:string, proxyProvider:string, proxyEndpointPath:string, customHost: string) {
   let reqURL = new URL(requestURL);
   let reqPath = reqURL.pathname;
   const reqQuery = reqURL.search;
   reqPath = reqPath.replace(proxyEndpointPath, "");
+
+  if (customHost) {
+    return `${customHost}${reqPath}${reqQuery}`
+  }
+
   const providerBasePath = Providers[proxyProvider].api.baseURL;
   if (proxyProvider === AZURE_OPEN_AI) {
     return `https:/${reqPath}${reqQuery}`;
@@ -55,7 +60,9 @@ export async function proxyGetHandler(c: Context): Promise<Response> {
       proxyPath: c.req.url.indexOf("/v1/proxy") > -1 ? "/v1/proxy" : "/v1"
     }
 
-    let urlToFetch = getProxyPath(c.req.url, store.proxyProvider, store.proxyPath);
+    const customHost = requestHeaders[HEADER_KEYS.CUSTOM_HOST] || "";
+
+    let urlToFetch = getProxyPath(c.req.url, store.proxyProvider, store.proxyPath, customHost);
 
     let fetchOptions = {
         headers: headersToSend(requestHeaders, store.customHeadersToAvoid),

src/handlers/proxyHandler.ts

@@ -1,7 +1,7 @@
 import { Context } from "hono";
 import { retryRequest } from "./retryHandler";
 import Providers from "../providers";
-import { ANTHROPIC, MAX_RETRIES, HEADER_KEYS, RETRY_STATUS_CODES, POWERED_BY, RESPONSE_HEADER_KEYS, AZURE_OPEN_AI, CONTENT_TYPES } from "../globals";
+import { ANTHROPIC, MAX_RETRIES, HEADER_KEYS, RETRY_STATUS_CODES, POWERED_BY, RESPONSE_HEADER_KEYS, AZURE_OPEN_AI, CONTENT_TYPES, OLLAMA } from "../globals";
 import { fetchProviderOptionsFromConfig, responseHandler, tryProvidersInSequence, updateResponseHeaders } from "./handlerUtils";
 import { convertKeysToCamelCase, getStreamingMode } from "../utils";
 import { Config, ShortConfig } from "../types/requestBody";
@@ -12,15 +12,24 @@ function proxyProvider(proxyModeHeader:string, providerHeader: string) {
   return proxyProvider;
 }
 
-function getProxyPath(requestURL:string, proxyProvider:string, proxyEndpointPath:string) {
+function getProxyPath(requestURL:string, proxyProvider:string, proxyEndpointPath:string, customHost: string) {
   let reqURL = new URL(requestURL);
   let reqPath = reqURL.pathname;
   const reqQuery = reqURL.search;
   reqPath = reqPath.replace(proxyEndpointPath, "");
+
+  if (customHost) {
+    return `${customHost}${reqPath}${reqQuery}`
+  }
+
   const providerBasePath = Providers[proxyProvider].api.baseURL;
   if (proxyProvider === AZURE_OPEN_AI) {
     return `https:/${reqPath}${reqQuery}`;
   }
+
+  if (proxyProvider === OLLAMA) {
+    return `https:/${reqPath}`;
+  }
   let proxyPath = `${providerBasePath}${reqPath}${reqQuery}`;
 
   // Fix specific for Anthropic SDK calls. Is this needed? - Yes
@@ -90,7 +99,8 @@ export async function proxyHandler(c: Context): Promise<Response> {
       }
     };
 
-    let urlToFetch = getProxyPath(c.req.url, store.proxyProvider, store.proxyPath);
+    const customHost = requestHeaders[HEADER_KEYS.CUSTOM_HOST] || requestConfig?.customHost || "";
+    let urlToFetch = getProxyPath(c.req.url, store.proxyProvider, store.proxyPath, customHost);
     store.isStreamingMode = getStreamingMode(store.reqBody, store.proxyProvider, urlToFetch)
 
     if (requestConfig &&

src/handlers/streamHandler.ts

@@ -1,4 +1,4 @@
-import { AZURE_OPEN_AI, CONTENT_TYPES, COHERE, GOOGLE, REQUEST_TIMEOUT_STATUS_CODE } from "../globals";
+import { AZURE_OPEN_AI, CONTENT_TYPES, COHERE, GOOGLE, REQUEST_TIMEOUT_STATUS_CODE, OLLAMA } from "../globals";
 import { OpenAIChatCompleteResponse } from "../providers/openai/chatComplete";
 import { OpenAICompleteResponse } from "../providers/openai/complete";
 import { getStreamModeSplitPattern } from "../utils";

src/middlewares/requestValidator/index.ts

@@ -15,6 +15,7 @@ import {
     DEEPINFRA,
     STABILITY_AI,
     NOMIC,
+    OLLAMA,
 } from "../../globals";
 import { configSchema } from "./schema/config";
 
@@ -63,7 +64,7 @@ export const requestValidator = (c: Context, next: any) => {
     }
     if (
         requestHeaders[`x-${POWERED_BY}-provider`] &&
-        ![OPEN_AI, AZURE_OPEN_AI, COHERE, ANTHROPIC, ANYSCALE, PALM, TOGETHER_AI, GOOGLE, MISTRAL_AI, PERPLEXITY_AI, DEEPINFRA, NOMIC, STABILITY_AI].includes(
+        ![OPEN_AI, AZURE_OPEN_AI, COHERE, ANTHROPIC, ANYSCALE, PALM, TOGETHER_AI, GOOGLE, MISTRAL_AI, PERPLEXITY_AI, DEEPINFRA, NOMIC, STABILITY_AI, OLLAMA].includes(
             requestHeaders[`x-${POWERED_BY}-provider`]
         )
     ) {
@@ -81,6 +82,22 @@ export const requestValidator = (c: Context, next: any) => {
         );
     }
 
+    const customHostHeader = requestHeaders[`x-${POWERED_BY}-custom-host`];
+    if (customHostHeader && customHostHeader.indexOf("api.portkey") > -1) {
+        return new Response(
+            JSON.stringify({
+                status: "failure",
+                message: `Invalid custom host`,
+            }),
+            {
+                status: 400,
+                headers: {
+                    "content-type": "application/json",
+                },
+            }
+        );;
+    }
+
 
     if (requestHeaders[`x-${POWERED_BY}-config`]) {
         try {