src/edurag_math_bot/ollama_client.py

from __future__ import annotations

from typing import Any

import requests


class OllamaClient:
    def __init__(self, base_url: str, timeout: int = 120) -> None:
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    def _post(self, endpoint: str, payload: dict[str, Any]) -> dict[str, Any]:
        try:
            response = requests.post(
                f"{self.base_url}{endpoint}",
                json=payload,
                timeout=self.timeout,
            )
            response.raise_for_status()
            return response.json()
        except requests.RequestException as exc:
            raise RuntimeError(
                "Could not connect to Ollama. Make sure `ollama serve` is running "
                f"and the base URL `{self.base_url}` is correct."
            ) from exc

    def embed_texts(self, texts: list[str], model: str) -> list[list[float]]:
        try:
            response = self._post("/api/embed", {"model": model, "input": texts})
            if "embeddings" in response:
                return response["embeddings"]
            if "embedding" in response:
                return [response["embedding"]]
        except requests.RequestException:
            pass

        embeddings: list[list[float]] = []
        for text in texts:
            response = self._post("/api/embeddings", {"model": model, "prompt": text})
            embeddings.append(response["embedding"])
        return embeddings

    def chat(
        self,
        model: str,
        messages: list[dict[str, str]],
        temperature: float = 0.2,
    ) -> str:
        response = self._post(
            "/api/chat",
            {
                "model": model,
                "messages": messages,
                "options": {"temperature": temperature},
                "stream": False,
            },
        )
        message = response.get("message", {})
        return message.get("content", "").strip()