Fixed Store embedding model url

antoninoLorenzo · Jul 16, 2024 · c9385c8 · c9385c8
1 parent b97c091
commit c9385c8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git 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
@@ -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(