""" Ollama client for embeddings generation. Provides async embedding operations via Ollama API: - Text embedding generation - Batch embedding support - Health checks Adapted from library-desk patterns. """ from typing import Optional import httpx from .config import config from .logging_config import get_logger logger = get_logger(__name__) class OllamaEmbeddingClient: """ Ollama API client for embeddings. Uses the Ollama embeddings endpoint to generate vector representations of text using the nomic-embed-text model (768 dimensions). Usage: client = OllamaEmbeddingClient() embedding = await client.embed("Hello world") await client.close() Or with context manager: async with OllamaEmbeddingClient() as client: embedding = await client.embed("Hello world") """ def __init__( self, base_url: str | None = None, model: str | None = None, timeout: float = 120.0, ): """ Initialize Ollama embedding client. Args: base_url: Ollama server URL (defaults to config.OLLAMA_HOST) model: Embedding model name (defaults to config.OLLAMA_EMBEDDING_MODEL) timeout: Request timeout in seconds (embeddings can be slow) """ self.base_url = (base_url or str(config.OLLAMA_HOST)).rstrip("/") self.model = model or config.OLLAMA_EMBEDDING_MODEL self.embeddings_url = f"{self.base_url}/api/embeddings" self.tags_url = f"{self.base_url}/api/tags" self._client: httpx.AsyncClient | None = None self._timeout = timeout logger.info( "ollama_embedding_client_initialized", base_url=self.base_url, model=self.model, ) async def _get_client(self) -> httpx.AsyncClient: """Get or create HTTP client.""" if self._client is None: self._client = httpx.AsyncClient(timeout=self._timeout) return self._client async def __aenter__(self) -> "OllamaEmbeddingClient": """Async context manager entry.""" await self._get_client() return self async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: """Async context manager exit.""" await self.close() async def close(self) -> None: """Close HTTP client.""" if self._client is not None: await self._client.aclose() self._client = None async def embed(self, text: str) -> list[float] | None: """ Generate embedding for single text. Args: text: Text to embed Returns: Embedding vector (768-dimensional for nomic-embed-text) or None on failure Example: >>> embedding = await client.embed("Hello world") >>> len(embedding) 768 """ try: client = await self._get_client() payload = { "model": self.model, "prompt": text, } response = await client.post(self.embeddings_url, json=payload) response.raise_for_status() data = response.json() embedding = data.get("embedding") if not embedding: logger.error("ollama_embed_no_embedding", response_data=data) return None return embedding except httpx.HTTPStatusError as e: logger.error( "ollama_embed_http_error", status_code=e.response.status_code, detail=e.response.text, ) return None except Exception as e: logger.error("ollama_embed_failed", error=str(e), exc_info=True) return None async def embed_batch( self, texts: list[str], show_progress: bool = False, ) -> list[list[float] | None]: """ Generate embeddings for multiple texts. Note: Ollama doesn't support native batch embeddings, so this sequentially calls embed() for each text. Args: texts: List of texts to embed show_progress: Log progress for large batches Returns: List of embedding vectors (same order as input) None entries for texts that failed to embed Example: >>> texts = ["Hello", "World", "Test"] >>> embeddings = await client.embed_batch(texts) >>> len(embeddings) 3 """ embeddings = [] for i, text in enumerate(texts): if show_progress and i % 10 == 0: logger.info( "ollama_embed_batch_progress", current=i, total=len(texts), ) embedding = await self.embed(text) embeddings.append(embedding) if show_progress: logger.info( "ollama_embed_batch_complete", successful=sum(1 for e in embeddings if e is not None), total=len(texts), ) return embeddings async def embed_batch_filtered( self, texts: list[str], show_progress: bool = False, ) -> list[list[float]]: """ Generate embeddings for multiple texts, filtering out failures. Args: texts: List of texts to embed show_progress: Log progress for large batches Returns: List of successful embedding vectors (may be shorter than input) Example: >>> embeddings = await client.embed_batch_filtered(texts) >>> all(e is not None for e in embeddings) True """ all_embeddings = await self.embed_batch(texts, show_progress) return [e for e in all_embeddings if e is not None] async def get_embedding_dimension(self) -> int | None: """ Get embedding dimension for current model. Returns: Embedding dimension (e.g., 768 for nomic-embed-text) or None on failure Example: >>> dim = await client.get_embedding_dimension() >>> dim 768 """ test_embedding = await self.embed("test") if test_embedding: return len(test_embedding) return None async def health_check(self) -> bool: """ Check if Ollama server is reachable and model is available. Returns: True if healthy, False otherwise """ try: client = await self._get_client() response = await client.get(self.tags_url, timeout=5.0) response.raise_for_status() data = response.json() models = data.get("models", []) # Check if our embedding model is available model_found = False for m in models: name = m.get("name", "") if name == self.model or name.startswith(f"{self.model}:"): model_found = True break if not model_found: logger.warning( "ollama_embedding_model_not_found", model=self.model, available=[m.get("name") for m in models], ) return False return True except Exception as e: logger.error("ollama_embedding_health_check_failed", error=str(e)) return False # Global client instance (lazy initialization) _embedding_client: OllamaEmbeddingClient | None = None def get_embedding_client() -> OllamaEmbeddingClient: """ Get global embedding client instance. Returns: OllamaEmbeddingClient instance """ global _embedding_client if _embedding_client is None: _embedding_client = OllamaEmbeddingClient() return _embedding_client