Implement async Ollama client
Add production-ready async HTTP client for Ollama API communication with proper error handling and dependency injection. Ollama Client (src/ollama/client.py): - Async context manager for connection lifecycle - Non-streaming chat endpoint - Streaming chat endpoint with async generator - Model listing endpoint - Health check endpoint - Timeout configuration per request - Comprehensive error handling with custom exceptions - FastAPI dependency injection support Ollama Schemas (src/ollama/schemas.py): - OllamaMessage: Chat message format - OllamaChatRequest: Request with model, messages, options - OllamaChatResponse: Complete chat response - OllamaModelInfo: Model metadata - OllamaModelsResponse: Model list response Features: - Async/await throughout for non-blocking I/O - Connection pooling via httpx.AsyncClient - Configurable timeouts (default: 120s) - Proper exception mapping (connection errors, timeouts) - Ready for integration (currently not connected to routes) Following Best Practices: - Async context manager pattern - Dependency injection for FastAPI routes - Separation of concerns (client vs schemas) - Type hints throughout - Comprehensive logging Status: Ready for integration (mock responses used in routes currently) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Claude
@@ -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
|
||||
@@ -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]
|
||||
Reference in New Issue
Block a user