refactor: ollama provider (#866)

yomorun · Sep 6, 2024 · 02356ec · 02356ec
1 parent c91a388
commit 02356ec
Show file tree

Hide file tree

Showing 5 changed files with 121 additions and 338 deletions.
diff --git a/cli/serve.go b/cli/serve.go
@@ -150,7 +150,7 @@ func registerAIProvider(aiConfig *ai.Config) error {
 				provider["model"],
 			))
 		case "ollama":
-			providerpkg.RegisterProvider(ollama.NewProvider(provider["api_endpoint"]))
+			providerpkg.RegisterProvider(ollama.NewProvider(provider["api_endpoint"], provider["model"]))
 		case "gemini":
 			providerpkg.RegisterProvider(gemini.NewProvider(provider["api_key"]))
 		case "githubmodels":

diff --git a/example/10-ai/zipper.yaml b/example/10-ai/zipper.yaml
@@ -2,15 +2,15 @@ name: ai-zipper
 host: 0.0.0.0
 port: 9000
 
-auth:
-  type: token
-  token: Happy New Year
+# auth:
+#   type: token
+#   token: Happy New Year
 
 bridge:
   ai:
     server:
       addr: localhost:8000
-      provider: anthropic
+      provider: ollama
 
     providers:
       azopenai:
@@ -33,6 +33,10 @@ bridge:
         deployment_id:
         api_version:
 
+      ollama:
+        api_endpoint: http://localhost:11434/v1
+        model: llama3.1
+
       cerebras:
         api_key:
         model:

diff --git a/pkg/bridge/ai/provider/ollama/Readme.md b/pkg/bridge/ai/provider/ollama/Readme.md
@@ -8,12 +8,10 @@ Follow the Ollama doc:
 
 <https://github.com/ollama/ollama?tab=readme-ov-file#ollama>
 
-## 2. Run the Mistral model
-
-Notice that only the Mistral v0.3+ models are supported currently.
+## 2. Run the model
 
 ```sh
-ollama run mistral:7b
+ollama run llama3.1
 ```
 
 ## 3. Start YoMo Zipper
@@ -34,7 +32,7 @@ bridge:
       provider: ollama
     providers:
       ollama:
-        api_endpoint: "http://localhost:11434/"
+        api_endpoint: "http://localhost:11434/v1"
 ```
 
 ```sh