"""Ollama LLM client with robust error handling.""" import logging import os from typing import Any import requests from requests.exceptions import HTTPError, RequestException, Timeout from ..config import settings from .exceptions import LLMAPIError, LLMConfigurationError logger = logging.getLogger(__name__) class OllamaClient: """ Client for interacting with Ollama API. Ollama runs locally and provides an OpenAI-compatible API. Example: >>> client = OllamaClient(model="llama3.2") >>> messages = [{"role": "user", "content": "Hello!"}] >>> response = client.complete(messages) >>> print(response) """ def __init__( self, base_url: str | None = None, model: str | None = None, timeout: int | None = None, temperature: float | None = None, ): """ Initialize Ollama client. Args: base_url: Ollama API base URL (defaults to http://localhost:11434) model: Model name to use (e.g., "llama3.2", "mistral", "codellama") timeout: Request timeout in seconds (defaults to settings) temperature: Temperature for generation (defaults to settings) Raises: LLMConfigurationError: If configuration is invalid """ self.base_url = base_url or os.getenv( "OLLAMA_BASE_URL", "http://localhost:11434" ) self.model = model or os.getenv("OLLAMA_MODEL", "llama3.2") self.timeout = timeout or settings.request_timeout self.temperature = ( temperature if temperature is not None else settings.temperature ) if not self.base_url: raise LLMConfigurationError( "Ollama base URL is required. Set OLLAMA_BASE_URL environment variable." ) if not self.model: raise LLMConfigurationError( "Ollama model is required. Set OLLAMA_MODEL environment variable." ) logger.info(f"Ollama client initialized with model: {self.model}") def complete(self, messages: list[dict[str, Any]]) -> str: """ Generate a completion from the LLM. Args: messages: List of message dicts with 'role' and 'content' keys Returns: Generated text response Raises: LLMAPIError: If API request fails ValueError: If messages format is invalid """ # Validate messages format if not messages: raise ValueError("Messages list cannot be empty") for msg in messages: if not isinstance(msg, dict): raise ValueError(f"Each message must be a dict, got {type(msg)}") if "role" not in msg or "content" not in msg: raise ValueError( f"Each message must have 'role' and 'content' keys, got {msg.keys()}" ) if msg["role"] not in ("system", "user", "assistant"): raise ValueError(f"Invalid role: {msg['role']}") url = f"{self.base_url}/api/chat" payload = { "model": self.model, "messages": messages, "stream": False, "options": { "temperature": self.temperature, }, } try: logger.debug(f"Sending request to {url} with {len(messages)} messages") response = requests.post(url, json=payload, timeout=self.timeout) response.raise_for_status() data = response.json() # Validate response structure if "message" not in data: raise LLMAPIError("Invalid API response: missing 'message'") if "content" not in data["message"]: raise LLMAPIError("Invalid API response: missing 'content' in message") content = data["message"]["content"] logger.debug(f"Received response with {len(content)} characters") return content except Timeout as e: logger.error(f"Request timeout after {self.timeout}s: {e}") raise LLMAPIError(f"Request timeout after {self.timeout} seconds") from e except HTTPError as e: logger.error(f"HTTP error from Ollama API: {e}") if e.response is not None: try: error_data = e.response.json() error_msg = error_data.get("error", str(e)) except Exception: error_msg = str(e) raise LLMAPIError(f"Ollama API error: {error_msg}") from e raise LLMAPIError(f"HTTP error: {e}") from e except RequestException as e: logger.error(f"Request failed: {e}") raise LLMAPIError(f"Failed to connect to Ollama API: {e}") from e except (KeyError, IndexError, TypeError) as e: logger.error(f"Failed to parse API response: {e}") raise LLMAPIError(f"Invalid API response format: {e}") from e def list_models(self) -> list[str]: """ List available models in Ollama. Returns: List of model names """ url = f"{self.base_url}/api/tags" try: response = requests.get(url, timeout=self.timeout) response.raise_for_status() data = response.json() models = [model["name"] for model in data.get("models", [])] logger.info(f"Found {len(models)} models: {models}") return models except Exception as e: logger.error(f"Failed to list models: {e}") return [] def is_available(self) -> bool: """ Check if Ollama is running and accessible. Returns: True if Ollama is available, False otherwise """ try: url = f"{self.base_url}/api/tags" response = requests.get(url, timeout=5) return response.status_code == 200 except Exception: return False