init

assistant_bot/ollama.py

@@ -0,0 +1,99 @@
+import asyncio
+import json
+import urllib.error
+import urllib.request
+from typing import Any
+
+from .ai import AIClientError
+
+
+class OllamaError(AIClientError):
+    pass
+
+
+class OllamaClient:
+    provider = "local"
+    display_name = "Ollama"
+
+    def __init__(self, base_url: str) -> None:
+        self.base_url = base_url.rstrip("/")
+
+    def normalize_model(self, model: str) -> str:
+        return model.strip()
+
+    async def list_models(self) -> list[str]:
+        return await asyncio.to_thread(self._list_models)
+
+    async def chat(
+        self,
+        model: str,
+        messages: list[dict[str, str]],
+        json_mode: bool = False,
+    ) -> str:
+        return await asyncio.to_thread(self._chat, model, messages, json_mode)
+
+    def _list_models(self) -> list[str]:
+        data = self._request_json("GET", "/api/tags", None, timeout=15)
+        models = data.get("models", [])
+        return sorted(model["name"] for model in models if model.get("name"))
+
+    def _chat(
+        self,
+        model: str,
+        messages: list[dict[str, str]],
+        json_mode: bool,
+    ) -> str:
+        payload = {
+            "model": model,
+            "messages": messages,
+            "stream": False,
+        }
+        if json_mode:
+            payload["format"] = "json"
+
+        data = self._request_json("POST", "/api/chat", payload, timeout=180)
+        content = data.get("message", {}).get("content")
+        if isinstance(content, str) and content.strip():
+            return content.strip()
+
+        fallback = data.get("response")
+        if isinstance(fallback, str) and fallback.strip():
+            return fallback.strip()
+
+        raise OllamaError("Ollama returned an empty response.")
+
+    def _request_json(
+        self,
+        method: str,
+        path: str,
+        payload: dict[str, Any] | None,
+        timeout: int,
+    ) -> dict[str, Any]:
+        url = f"{self.base_url}{path}"
+        data = json.dumps(payload).encode("utf-8") if payload is not None else None
+        request = urllib.request.Request(
+            url,
+            data=data,
+            method=method,
+            headers={"Content-Type": "application/json"},
+        )
+
+        try:
+            with urllib.request.urlopen(request, timeout=timeout) as response:
+                raw = response.read().decode("utf-8")
+        except urllib.error.HTTPError as exc:
+            body = exc.read().decode("utf-8", errors="replace")[:400]
+            raise OllamaError(f"Ollama HTTP {exc.code}: {body}") from exc
+        except urllib.error.URLError as exc:
+            raise OllamaError(f"Ollama is unavailable at {self.base_url}: {exc.reason}") from exc
+        except TimeoutError as exc:
+            raise OllamaError("Ollama request timed out.") from exc
+
+        try:
+            parsed = json.loads(raw)
+        except json.JSONDecodeError as exc:
+            raise OllamaError("Ollama returned invalid JSON.") from exc
+
+        if not isinstance(parsed, dict):
+            raise OllamaError("Ollama returned an unexpected response.")
+        return parsed