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(