latest

routers/llm/private/ollama/ollama_wrapper.py

@@ -4,14 +4,16 @@ from langchain_core.runnables.base import Runnable
 from langchain.prompts.base import StringPromptValue
 
 class OllamaWrapper(Runnable):
-    def __init__(self, host: str):
+    def __init__(self, host: str, model: str = "llama3.2:latest"):
         self.client = Client(host=host)
+        self.model = model
 
     def invoke(self, prompt: Any, config: Dict[str, Any] = None, **kwargs: Any) -> str:
         if isinstance(prompt, StringPromptValue):
             prompt = prompt.to_string()
 
-        model_name = kwargs.get("model", "llama3")
+        # Use the model from constructor, but allow override via kwargs
+        model_name = kwargs.get("model", self.model)
         options = {
             "temperature": kwargs.get("temperature"),
             "top_p": kwargs.get("top_p"),