From c9385c8f389e67792db6b0bf57cb244698081a4d Mon Sep 17 00:00:00 2001
From: Antonino Lorenzo <94693967+antoninoLorenzo@users.noreply.github.com>
Date: Tue, 16 Jul 2024 11:38:36 +0200
Subject: [PATCH] Fixed Store embedding model url

---
 src/agent/knowledge/routing.py | 2 +-
 src/agent/knowledge/store.py   | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/agent/knowledge/routing.py b/src/agent/knowledge/routing.py
index 680da22..be5ee0f 100644
--- a/src/agent/knowledge/routing.py
+++ b/src/agent/knowledge/routing.py
@@ -61,7 +61,7 @@ class LLMRouter(Router):
     Using a local model is not the best choice for performance, HuggingFace
     Inference API could be used in future"""
 
-    def __init__(self, model: str = 'gemma:2b'):
+    def __init__(self, model: str = 'gemma2:9b'):
         self.llm = LLM(model)
         self.system_prompt = PROMPTS[model]['routing']['system']
         self.user_prompt = PROMPTS[model]['routing']['user']
diff --git a/src/agent/knowledge/store.py b/src/agent/knowledge/store.py
index c6312ef..bfbb454 100644
--- a/src/agent/knowledge/store.py
+++ b/src/agent/knowledge/store.py
@@ -17,6 +17,7 @@ class Store:
 
     def __init__(self,
                  url: str = 'http://localhost:6333',
+                 embedding_url: str = 'http://localhost:11434',
                  in_memory: bool = False,
                  router: Router = None
                  ):
@@ -38,7 +39,7 @@ def __init__(self,
                 coll = {}
             self._collections: Dict[str: Collection] = coll
 
-        self._encoder = ollama.embeddings
+        self._encoder = ollama.Client(host=embedding_url).embeddings
         self._embedding_model: str = 'nomic-embed-text'
         self._embedding_size: int = len(
             self._encoder(