From 144e3126fe61cadefa38257b09a2f4576715ad77 Mon Sep 17 00:00:00 2001
From: Yoshiki Miura <miurap400@gmail.com>
Date: Sat, 26 Oct 2024 22:36:37 +0900
Subject: [PATCH] Refactor environment variables and provider configuration

---
 .env.local.example    | 3 +++
 app/actions.tsx       | 2 +-
 lib/utils/registry.ts | 8 +++++---
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/.env.local.example b/.env.local.example
index b04a969f..9bc26308 100644
--- a/.env.local.example
+++ b/.env.local.example
@@ -47,6 +47,9 @@ SEARXNG_SAFESEARCH=0 # Safe search setting: 0 (off), 1 (moderate), 2 (strict)
 # Groq API key retrieved here: https://console.groq.com/keys
 # GROQ_API_KEY=[YOUR_GROQ_API_KEY]
 
+# Ollama Base URL
+# OLLAMA_BASE_URL=http://localhost:11434
+
 # Azure OpenAI API key retrieved here: https://oai.azure.com/resource/deployments/
 # AZURE_API_KEY=
 # The resource name is used in the assembled URL: https://{resourceName}.openai.azure.com/openai/deployments/{modelId}{path}.
diff --git a/app/actions.tsx b/app/actions.tsx
index b706c835..d110e773 100644
--- a/app/actions.tsx
+++ b/app/actions.tsx
@@ -78,7 +78,7 @@ async function submit(
   // Check if provider is enabled
   if (!isProviderEnabled(providerId)) {
     throw new Error(
-      `Provider ${providerId} is not available (API key not configured)`
+      `Provider ${providerId} is not available (API key not configured or base URL not set)`
     )
   }
 
diff --git a/lib/utils/registry.ts b/lib/utils/registry.ts
index f2f6329c..b41696b5 100644
--- a/lib/utils/registry.ts
+++ b/lib/utils/registry.ts
@@ -3,7 +3,7 @@ import { openai, createOpenAI } from '@ai-sdk/openai'
 import { anthropic } from '@ai-sdk/anthropic'
 import { google } from '@ai-sdk/google'
 import { createAzure } from '@ai-sdk/azure'
-import { ollama } from 'ollama-ai-provider'
+import { createOllama } from 'ollama-ai-provider'
 
 export const registry = createProviderRegistry({
   openai,
@@ -13,7 +13,9 @@ export const registry = createProviderRegistry({
     apiKey: process.env.GROQ_API_KEY,
     baseURL: 'https://api.groq.com/openai/v1'
   }),
-  ollama,
+  ollama: createOllama({
+    baseURL: `${process.env.OLLAMA_BASE_URL}/api`
+  }),
   azure: createAzure({
     apiKey: process.env.AZURE_API_KEY,
     resourceName: process.env.AZURE_RESOURCE_NAME
@@ -35,7 +37,7 @@ export function isProviderEnabled(providerId: string): boolean {
     case 'groq':
       return !!process.env.GROQ_API_KEY
     case 'ollama':
-      return true
+      return !!process.env.OLLAMA_BASE_URL
     case 'azure':
       return !!process.env.AZURE_API_KEY && !!process.env.AZURE_RESOURCE_NAME
     default: