Initial commit

routers/llm/private/ollama/ollama_wrapper.py

@@ -0,0 +1,28 @@
+from typing import Any, Dict
+from ollama import Client
+from langchain_core.runnables.base import Runnable
+from langchain.prompts.base import StringPromptValue
+
+class OllamaWrapper(Runnable):
+    def __init__(self, host: str):
+        self.client = Client(host=host)
+
+    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")
+        options = {
+            "temperature": kwargs.get("temperature"),
+            "top_p": kwargs.get("top_p"),
+            "max_tokens": kwargs.get("max_tokens"),
+        }
+        messages = [{"role": "user", "content": prompt}]
+        response = self.client.chat(model=model_name, messages=messages, options=options)
+        if response and "message" in response and "content" in response["message"]:
+            return response["message"]["content"]
+        else:
+            raise ValueError("Invalid response structure from model")
+
+    async def ainvoke(self, prompt: Any, config: Dict[str, Any] = None, **kwargs: Any) -> str:
+        return self.invoke(prompt, config, **kwargs)