diff --git a/src/core/dependencies.py b/src/core/dependencies.py new file mode 100644 index 0000000..e69de29 diff --git a/src/ollama/__init__.py b/src/ollama/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/ollama/client.py b/src/ollama/client.py new file mode 100644 index 0000000..c4ed984 --- /dev/null +++ b/src/ollama/client.py @@ -0,0 +1,181 @@ +""" +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 diff --git a/src/ollama/schemas.py b/src/ollama/schemas.py new file mode 100644 index 0000000..78010ee --- /dev/null +++ b/src/ollama/schemas.py @@ -0,0 +1,42 @@ +""" +Ollama API schemas. +Internal models for Ollama API communication. +""" +from typing import Any + +from src.core.models import CustomBaseModel + + +class OllamaMessage(CustomBaseModel): + """Message format for Ollama API.""" + role: str + content: str + + +class OllamaChatRequest(CustomBaseModel): + """Chat request to Ollama API.""" + model: str + messages: list[OllamaMessage] + stream: bool = False + options: dict[str, Any] | None = None + + +class OllamaChatResponse(CustomBaseModel): + """Chat response from Ollama API.""" + model: str + created_at: str + message: OllamaMessage + done: bool + + +class OllamaModelInfo(CustomBaseModel): + """Model information from Ollama.""" + name: str + modified_at: str + size: int + digest: str + + +class OllamaModelsResponse(CustomBaseModel): + """Response from Ollama models list endpoint.""" + models: list[OllamaModelInfo]