Added llm logic and config

1 files changed, 42 insertions, 0 deletions

diff --git a/llm.py b/llm.py
new file mode 100644
index 0000000..524f173
--- /dev/null
+++ b/llm.py
@@ -0,0 +1,42 @@
+import json
+from collections.abc import Iterator
+
+import requests
+
+class LLM:
+    def __init__(self, url: str) -> None:
+        self.url = url
+
+    def chat(self, messages: list[dict[str, str]], max_tokens: int = 256) -> Iterator[str]:
+        response = requests.post(
+            self.url,
+            json = {
+                "messages": messages,
+                "max_tokens": max_tokens,
+                "temperature": 0.7,
+                "stream": True,
+            },
+            stream=True,
+        )
+        response.raise_for_status()
+
+        for line in response.iter_lines():
+            if not line:
+                continue
+            
+            if not line.startswith(b"data: "):
+                continue
+
+            data = line[6:]
+
+            if data == b"[DONE]":
+                break
+
+            chunk = json.loads(data)
+
+            delta = chunk["choices"][0]["delta"]
+            content = delta.get("content")
+
+            if content:
+                yield content
+