Merge pull request stackblitz-labs#64 from noobydp/main

Further changes to support OLLAMA_API_BASE_URL with ollama models
arsenikstiger · Oct 24, 2024 · b21b145 · b21b145
2 parents 5f6a595 + f1b8c02
commit b21b145
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/.gitignore b/.gitignore
@@ -12,7 +12,7 @@ dist-ssr
 *.local
 
 .vscode/*
-!.vscode/launch.json
+.vscode/launch.json
 !.vscode/extensions.json
 .idea
 .DS_Store

diff --git a/app/lib/.server/llm/api-key.ts b/app/lib/.server/llm/api-key.ts
@@ -34,6 +34,8 @@ export function getBaseURL(cloudflareEnv: Env, provider: string) {
   switch (provider) {
     case 'OpenAILike':
       return env.OPENAI_LIKE_API_BASE_URL || cloudflareEnv.OPENAI_LIKE_API_BASE_URL;
+    case 'Ollama':
+        return env.OLLAMA_API_BASE_URL || cloudflareEnv.OLLAMA_API_BASE_URL;
     default:
       return "";
   }

diff --git a/app/lib/.server/llm/model.ts b/app/lib/.server/llm/model.ts
@@ -57,6 +57,12 @@ export function getGroqModel(apiKey: string, model: string) {
   return openai(model);
 }
 
+export function getOllamaModel(baseURL: string, model: string) {
+  let Ollama = ollama(model);
+  Ollama.config.baseURL = `${baseURL}/api`;
+  return Ollama;
+}
+
 export function getDeepseekModel(apiKey: string, model: string){
   const openai = createOpenAI({
     baseURL: 'https://api.deepseek.com/beta',
@@ -65,9 +71,6 @@ export function getDeepseekModel(apiKey: string, model: string){
 
   return openai(model);
 }
-export function getOllamaModel(model: string) {
-  return ollama(model);
-}
 
 export function getOpenRouterModel(apiKey: string, model: string) {
   const openRouter = createOpenRouter({
@@ -99,6 +102,6 @@ export function getModel(provider: string, model: string, env: Env) {
     case 'Mistral':
       return  getMistralModel(apiKey, model);
     default:
-      return getOllamaModel(model);
+      return getOllamaModel(baseURL, model);
   }
 }