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