tatlock/src/core/embeddings.py

"""
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