fithia2/libs/llm/client.py

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

import asyncio
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"
# Per-chunk read timeout for streaming; model can take minutes to start tokens on CPU.
_STREAM_CHUNK_TIMEOUT = 600.0


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

    Uses synchronous httpx inside asyncio.to_thread() to fully isolate Ollama
    HTTP calls from the asyncio event loop, preventing hangs when called from
    within an asyncpg / SQLAlchemy session context on Python 3.13.
    """

    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

    async def __aenter__(self) -> OllamaClient:
        return self

    async def __aexit__(self, *_: object) -> None:
        pass

    def _sync_call(
        self,
        messages: list[dict[str, str]],
        response_format: str,
        temperature: float,
    ) -> tuple[dict[str, Any], dict[str, int], int]:
        """Synchronous HTTP call to Ollama /api/chat.

        Runs inside asyncio.to_thread() so it does not block the event loop.

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

        Raises:
            LLMTimeoutError: On request timeout.
            LLMError: On non-retryable HTTP errors or bad JSON.
            RetryableError: On 5xx server errors or connection failures.
        """
        payload: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "stream": True,
            # Disable Qwen3/Qwen3.5 extended thinking mode — thinking tokens add
            # 30-60s latency with no benefit for structured JSON extraction tasks.
            "think": False,
            "options": {
                "temperature": temperature,
                # Limit KV cache to actual prompt size; default 262144 for Qwen3.5
                # allocates 20 GB of VRAM which slows cold-start significantly.
                "num_ctx": 8192,
            },
        }
        if response_format == "json":
            payload["format"] = "json"

        stream_timeout = httpx.Timeout(
            connect=10.0,
            read=_STREAM_CHUNK_TIMEOUT,
            write=30.0,
            pool=10.0,
        )

        t0 = time.monotonic()
        try:
            chunks: list[str] = []
            prompt_tokens = 0
            completion_tokens = 0
            with httpx.Client(base_url=self._base_url, timeout=stream_timeout) as http:
                with http.stream("POST", _CHAT_PATH, json=payload) as response:
                    if response.status_code >= 500:
                        body = response.read()
                        raise RetryableError(
                            f"Ollama server error {response.status_code}",
                            source="ollama",
                            context={"status": response.status_code, "body": body[:200].decode()},
                        )
                    if response.status_code >= 400:
                        body = response.read()
                        raise LLMError(
                            f"Ollama client error {response.status_code}: {body[:200].decode()}",
                            source="ollama",
                            context={"status": response.status_code},
                        )
                    for line in response.iter_lines():
                        if not line.strip():
                            continue
                        try:
                            chunk = json.loads(line)
                        except json.JSONDecodeError:
                            continue
                        content = chunk.get("message", {}).get("content", "")
                        if content:
                            chunks.append(content)
                        if chunk.get("done"):
                            prompt_tokens = chunk.get("prompt_eval_count", 0)
                            completion_tokens = chunk.get("eval_count", 0)
                            break
        except (LLMTimeoutError, LLMError, RetryableError):
            raise
        except httpx.TimeoutException as exc:
            raise LLMTimeoutError(
                f"Ollama request timed out after {_STREAM_CHUNK_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)
        raw_content = "".join(chunks)
        token_usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
        }

        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

    @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.
        """
        try:
            return await asyncio.to_thread(
                self._sync_call, messages, response_format, temperature
            )
        except (LLMTimeoutError, LLMError, RetryableError):
            raise
        except Exception as exc:
            raise LLMError(
                f"Unexpected error calling Ollama: {exc}",
                source="ollama",
                context={"model": self.model},
            ) from exc

    async def health_check(self) -> bool:
        """Return True if Ollama is reachable and the model is available."""

        def _sync_health() -> bool:
            try:
                with httpx.Client(base_url=self._base_url, timeout=10.0) as http:
                    response = http.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

        return await asyncio.to_thread(_sync_health)


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),
    )