fithia2/libs/llm/client.py

"""Async Ollama HTTP client with retry logic."""
from __future__ import annotations

import json
import time
from typing import Any

import httpx

from libs.common.logging import get_logger
from libs.common.retries import RetryableError, with_retry
from libs.llm.exceptions import LLMError, LLMTimeoutError

logger = get_logger(__name__)

_CHAT_PATH = "/api/chat"
_TAGS_PATH = "/api/tags"


class OllamaClient:
    """Async client for Ollama REST API."""

    def __init__(self, base_url: str, model: str, timeout: float = 60.0) -> None:
        self._base_url = base_url.rstrip("/")
        self.model = model
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None

    async def __aenter__(self) -> OllamaClient:
        self._client = httpx.AsyncClient(base_url=self._base_url, timeout=self._timeout)
        return self

    async def __aexit__(self, *_: object) -> None:
        if self._client is not None:
            await self._client.aclose()
        self._client = None

    @with_retry(max_attempts=2, min_wait=0.5, max_wait=10.0, multiplier=2.0)
    async def chat(
        self,
        messages: list[dict[str, str]],
        response_format: str = "json",
        temperature: float = 0.0,
    ) -> tuple[dict[str, Any], dict[str, int], int]:
        """Call Ollama /api/chat and return (parsed_json, token_usage, elapsed_ms).

        Args:
            messages: List of {role, content} message dicts.
            response_format: "json" forces JSON output mode.
            temperature: Sampling temperature (0.0 = deterministic).

        Returns:
            Tuple of (parsed response dict, token usage dict, elapsed_ms).

        Raises:
            LLMTimeoutError: On request timeout.
            LLMError: On non-retryable HTTP errors.
            RetryableError: On 5xx server errors.
        """
        if self._client is None:
            raise LLMError("OllamaClient must be used as an async context manager")

        payload: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "stream": False,
            "options": {"temperature": temperature},
        }
        if response_format == "json":
            payload["format"] = "json"

        t0 = time.monotonic()
        try:
            response = await self._client.post(_CHAT_PATH, json=payload)
        except httpx.TimeoutException as exc:
            raise LLMTimeoutError(
                f"Ollama request timed out after {self._timeout}s",
                source="ollama",
                context={"model": self.model},
            ) from exc
        except httpx.ConnectError as exc:
            raise RetryableError(
                f"Cannot connect to Ollama at {self._base_url}",
                source="ollama",
                context={"model": self.model},
            ) from exc

        elapsed_ms = int((time.monotonic() - t0) * 1000)

        if response.status_code >= 500:
            raise RetryableError(
                f"Ollama server error {response.status_code}",
                source="ollama",
                context={"status": response.status_code, "body": response.text[:200]},
            )
        if response.status_code >= 400:
            raise LLMError(
                f"Ollama client error {response.status_code}: {response.text[:200]}",
                source="ollama",
                context={"status": response.status_code},
            )

        data = response.json()
        raw_content: str = data.get("message", {}).get("content", "")
        token_usage = {
            "prompt_tokens": data.get("prompt_eval_count", 0),
            "completion_tokens": data.get("eval_count", 0),
        }

        try:
            parsed = json.loads(raw_content)
        except json.JSONDecodeError as exc:
            raise LLMError(
                "Ollama returned non-JSON content",
                source="ollama",
                context={"raw": raw_content[:300]},
            ) from exc

        logger.debug(
            "ollama_chat_ok",
            model=self.model,
            elapsed_ms=elapsed_ms,
            **token_usage,
        )
        return parsed, token_usage, elapsed_ms

    async def health_check(self) -> bool:
        """Return True if Ollama is reachable and the model is available."""
        if self._client is None:
            raise LLMError("OllamaClient must be used as an async context manager")
        try:
            response = await self._client.get(_TAGS_PATH)
            if response.status_code != 200:
                return False
            data = response.json()
            models = [m.get("name", "") for m in data.get("models", [])]
            return any(self.model in name for name in models)
        except Exception:
            return False


def make_ollama_client() -> OllamaClient:
    """Factory that reads config from settings."""
    from libs.common.config import get_settings

    s = get_settings()
    return OllamaClient(
        base_url=s.ollama_url,
        model=s.ollama_model,
        timeout=float(s.ollama_timeout),
    )