tatlock/src/ollama/client.py

"""
Ollama HTTP client.
Handles all communication with the Ollama service.
"""
import logging
from typing import Any, AsyncGenerator

import httpx
from httpx import ConnectError, TimeoutException

from src.core.config import config
from src.core.exceptions import OllamaConnectionError, OllamaTimeoutError
from src.ollama.schemas import (
    OllamaChatRequest,
    OllamaChatResponse,
    OllamaModelsResponse,
)

logger = logging.getLogger(__name__)


class OllamaClient:
    """
    Async client for Ollama API.
    
    Follows best practice of using async for I/O operations.
    """
    
    def __init__(self, base_url: str | None = None, timeout: int | None = None):
        """
        Initialize Ollama client.
        
        Args:
            base_url: Ollama server URL (defaults to config)
            timeout: Request timeout in seconds (defaults to config)
        """
        self.base_url = base_url or str(config.OLLAMA_HOST)
        self.timeout = timeout or config.OLLAMA_TIMEOUT
        self._client: httpx.AsyncClient | None = None
    
    async def __aenter__(self) -> "OllamaClient":
        """Async context manager entry."""
        self._client = httpx.AsyncClient(
            base_url=self.base_url,
            timeout=self.timeout,
        )
        return self
    
    async def __aexit__(self, *args: Any) -> None:
        """Async context manager exit."""
        if self._client:
            await self._client.aclose()
    
    async def chat(
        self,
        request: OllamaChatRequest,
    ) -> OllamaChatResponse:
        """
        Send chat request to Ollama (non-streaming).
        
        Args:
            request: Chat request with model and messages
            
        Returns:
            Complete chat response
            
        Raises:
            OllamaConnectionError: Cannot connect to Ollama
            OllamaTimeoutError: Request timed out
        """
        if not self._client:
            raise RuntimeError("Client not initialized. Use async with context.")
        
        try:
            response = await self._client.post(
                "/api/chat",
                json=request.model_dump(),
            )
            response.raise_for_status()
            return OllamaChatResponse(**response.json())
            
        except ConnectError as e:
            logger.error(f"Cannot connect to Ollama at {self.base_url}: {e}")
            raise OllamaConnectionError() from e
            
        except TimeoutException as e:
            logger.error(f"Ollama request timed out after {self.timeout}s: {e}")
            raise OllamaTimeoutError() from e
    
    async def chat_stream(
        self,
        request: OllamaChatRequest,
    ) -> AsyncGenerator[dict[str, Any], None]:
        """
        Send streaming chat request to Ollama.
        
        Args:
            request: Chat request with stream=True
            
        Yields:
            Streaming response chunks
            
        Raises:
            OllamaConnectionError: Cannot connect to Ollama
            OllamaTimeoutError: Request timed out
        """
        if not self._client:
            raise RuntimeError("Client not initialized. Use async with context.")
        
        # Ensure streaming is enabled
        request.stream = True
        
        try:
            async with self._client.stream(
                "POST",
                "/api/chat",
                json=request.model_dump(),
            ) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    if line.strip():
                        import json
                        yield json.loads(line)
                        
        except ConnectError as e:
            logger.error(f"Cannot connect to Ollama at {self.base_url}: {e}")
            raise OllamaConnectionError() from e
            
        except TimeoutException as e:
            logger.error(f"Ollama request timed out after {self.timeout}s: {e}")
            raise OllamaTimeoutError() from e
    
    async def list_models(self) -> OllamaModelsResponse:
        """
        List available models from Ollama.
        
        Returns:
            List of available models
            
        Raises:
            OllamaConnectionError: Cannot connect to Ollama
        """
        if not self._client:
            raise RuntimeError("Client not initialized. Use async with context.")
        
        try:
            response = await self._client.get("/api/tags")
            response.raise_for_status()
            return OllamaModelsResponse(**response.json())
            
        except ConnectError as e:
            logger.error(f"Cannot connect to Ollama at {self.base_url}: {e}")
            raise OllamaConnectionError() from e
    
    async def health_check(self) -> bool:
        """
        Check if Ollama service is healthy.
        
        Returns:
            True if healthy, False otherwise
        """
        if not self._client:
            raise RuntimeError("Client not initialized. Use async with context.")
        
        try:
            response = await self._client.get("/")
            return response.status_code == 200
        except Exception as e:
            logger.warning(f"Ollama health check failed: {e}")
            return False


# Dependency for FastAPI routes
async def get_ollama_client() -> AsyncGenerator[OllamaClient, None]:
    """
    FastAPI dependency to provide Ollama client.
    
    Follows best practice of dependency injection.
    """
    async with OllamaClient() as client:
        yield client