From fa8c372b2d517331b605d774dfbdf0dadfe68a78 Mon Sep 17 00:00:00 2001
From: CP500 <34076958+CP500@users.noreply.github.com>
Date: Mon, 19 Feb 2024 16:38:40 +1000
Subject: [PATCH] Create ollama.py
Adding Ollama support.
Howto use:
from ollama import OllamaLLM
rl = RouteLayer(encoder = encoder, routes=routes, llm=llm)
---
semantic_router/llms/ollama.py | 37 ++++++++++++++++++++++++++++++++++
1 file changed, 37 insertions(+)
create mode 100644 semantic_router/llms/ollama.py
diff --git a/semantic_router/llms/ollama.py b/semantic_router/llms/ollama.py
new file mode 100644
index 00000000..5b9dfc7a
--- /dev/null
+++ b/semantic_router/llms/ollama.py
@@ -0,0 +1,37 @@
+import os
+from typing import List, Optional
+import requests
+import json
+
+from semantic_router.llms import BaseLLM
+from semantic_router.schema import Message
+from semantic_router.utils.logger import logger
+
+
+class OllamaLLM(BaseLLM):
+ max_tokens: Optional[int] = 200
+
+
+ def _call_(self, messages: List[Message]) -> str:
+
+ try:
+
+ payload = {
+ "model": self.name,
+ "messages": [m.to_openai() for m in messages],
+ "options":{
+ "temperature":0.0,
+ "num_predict":self.max_tokens
+ },
+ "format":"json",
+ "stream":False
+ }
+
+ response = requests.post("http://localhost:11434/api/chat", json=payload)
+
+ output = response.json()["message"]["content"]
+
+ return output
+ except Exception as e:
+ logger.error(f"LLM error: {e}")
+ raise Exception(f"LLM error: {e}") from e
\ No newline at end of file
--
GitLab