gobby 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

gobby/search/backends/__init__.py

@@ -0,0 +1,159 @@
+"""Search backend abstractions.
+
+This module provides the protocol and implementations for search backends
+used by UnifiedSearcher:
+
+- AsyncSearchBackend: Protocol for async search backends
+- TFIDFBackend: TF-IDF based search (always available)
+- EmbeddingBackend: Embedding-based search (requires API)
+
+Usage:
+    from gobby.search.backends import AsyncSearchBackend, TFIDFBackend
+
+    backend: AsyncSearchBackend = TFIDFBackend()
+    await backend.fit_async([("id1", "content1"), ("id2", "content2")])
+    results = await backend.search_async("query", top_k=10)
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+# Re-export sync TFIDFSearcher for backwards compatibility
+from gobby.search.tfidf import TFIDFSearcher
+
+__all__ = [
+    "AsyncSearchBackend",
+    "TFIDFBackend",
+    "EmbeddingBackend",
+    "TFIDFSearcher",
+]
+
+
+@runtime_checkable
+class AsyncSearchBackend(Protocol):
+    """Protocol for async search backends.
+
+    All search backends must implement this interface. The protocol
+    uses async methods to support embedding-based backends that need
+    to call external APIs.
+
+    Methods:
+        fit_async: Build/rebuild the search index
+        search_async: Find relevant items for a query
+        needs_refit: Check if index needs rebuilding
+        get_stats: Get backend statistics
+        clear: Clear the search index
+    """
+
+    async def fit_async(self, items: list[tuple[str, str]]) -> None:
+        """Build or rebuild the search index.
+
+        Args:
+            items: List of (item_id, content) tuples to index
+        """
+        ...
+
+    async def search_async(
+        self,
+        query: str,
+        top_k: int = 10,
+    ) -> list[tuple[str, float]]:
+        """Search for items matching the query.
+
+        Args:
+            query: Search query text
+            top_k: Maximum number of results to return
+
+        Returns:
+            List of (item_id, similarity_score) tuples, sorted by
+            relevance (highest similarity first).
+        """
+        ...
+
+    def needs_refit(self) -> bool:
+        """Check if the search index needs rebuilding.
+
+        Returns:
+            True if fit_async() should be called before search_async()
+        """
+        ...
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get statistics about the search index.
+
+        Returns:
+            Dict with backend-specific statistics
+        """
+        ...
+
+    def clear(self) -> None:
+        """Clear the search index."""
+        ...
+
+
+class TFIDFBackend:
+    """Async wrapper around TFIDFSearcher.
+
+    Provides the AsyncSearchBackend interface for TF-IDF search.
+    This is a thin wrapper that delegates to the sync TFIDFSearcher.
+    """
+
+    def __init__(
+        self,
+        ngram_range: tuple[int, int] = (1, 2),
+        max_features: int = 10000,
+        min_df: int = 1,
+        stop_words: str | None = "english",
+        refit_threshold: int = 10,
+    ):
+        """Initialize TF-IDF backend.
+
+        Args:
+            ngram_range: Min/max n-gram sizes for tokenization
+            max_features: Maximum vocabulary size
+            min_df: Minimum document frequency for inclusion
+            stop_words: Language for stop words (None to disable)
+            refit_threshold: Number of updates before automatic refit
+        """
+        self._searcher = TFIDFSearcher(
+            ngram_range=ngram_range,
+            max_features=max_features,
+            min_df=min_df,
+            stop_words=stop_words,
+            refit_threshold=refit_threshold,
+        )
+
+    async def fit_async(self, items: list[tuple[str, str]]) -> None:
+        """Build or rebuild the search index."""
+        self._searcher.fit(items)
+
+    async def search_async(
+        self,
+        query: str,
+        top_k: int = 10,
+    ) -> list[tuple[str, float]]:
+        """Search for items matching the query."""
+        return self._searcher.search(query, top_k)
+
+    def needs_refit(self) -> bool:
+        """Check if the search index needs rebuilding."""
+        return self._searcher.needs_refit()
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get statistics about the search index."""
+        stats = self._searcher.get_stats()
+        stats["backend_type"] = "tfidf"
+        return stats
+
+    def clear(self) -> None:
+        """Clear the search index."""
+        self._searcher.clear()
+
+    def mark_update(self) -> None:
+        """Mark that an item update occurred."""
+        self._searcher.mark_update()
+
+
+# Import EmbeddingBackend - needs to be at end to avoid circular imports
+from gobby.search.backends.embedding import EmbeddingBackend  # noqa: E402

gobby/search/backends/embedding.py

@@ -0,0 +1,225 @@
+"""Embedding-based search backend.
+
+This module provides embedding-based semantic search using cosine similarity.
+It stores embeddings in memory and uses LiteLLM for embedding generation.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from gobby.search.models import SearchConfig
+
+logger = logging.getLogger(__name__)
+
+
+def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
+    """Compute cosine similarity between two vectors.
+
+    Args:
+        vec1: First vector
+        vec2: Second vector
+
+    Returns:
+        Cosine similarity score between -1 and 1
+    """
+    if len(vec1) != len(vec2):
+        return 0.0
+
+    dot_product = sum(a * b for a, b in zip(vec1, vec2, strict=True))
+    norm1 = math.sqrt(sum(a * a for a in vec1))
+    norm2 = math.sqrt(sum(b * b for b in vec2))
+
+    if norm1 == 0 or norm2 == 0:
+        return 0.0
+
+    return dot_product / (norm1 * norm2)
+
+
+class EmbeddingBackend:
+    """Embedding-based search backend using LiteLLM.
+
+    This backend generates embeddings for indexed items and uses
+    cosine similarity for search. Embeddings are stored in memory.
+
+    Supports all providers supported by LiteLLM:
+    - OpenAI (text-embedding-3-small)
+    - Ollama (openai/nomic-embed-text with api_base)
+    - Azure, Gemini, Mistral, etc.
+
+    Example:
+        backend = EmbeddingBackend(
+            model="text-embedding-3-small",
+            api_key="sk-..."
+        )
+        await backend.fit_async([("id1", "hello"), ("id2", "world")])
+        results = await backend.search_async("greeting", top_k=5)
+    """
+
+    def __init__(
+        self,
+        model: str = "text-embedding-3-small",
+        api_base: str | None = None,
+        api_key: str | None = None,
+    ):
+        """Initialize embedding backend.
+
+        Args:
+            model: LiteLLM model string
+            api_base: Optional API base URL for custom endpoints
+            api_key: Optional API key (uses env var if not set)
+        """
+        self._model = model
+        self._api_base = api_base
+        self._api_key = api_key
+
+        # Item storage
+        self._item_ids: list[str] = []
+        self._item_embeddings: list[list[float]] = []
+        self._item_contents: dict[str, str] = {}  # For reindexing
+        self._fitted = False
+
+    @classmethod
+    def from_config(cls, config: SearchConfig) -> EmbeddingBackend:
+        """Create an EmbeddingBackend from a SearchConfig.
+
+        Args:
+            config: SearchConfig with model and API settings
+
+        Returns:
+            Configured EmbeddingBackend instance
+        """
+        return cls(
+            model=config.embedding_model,
+            api_base=config.embedding_api_base,
+            api_key=config.embedding_api_key,
+        )
+
+    async def fit_async(self, items: list[tuple[str, str]]) -> None:
+        """Build or rebuild the search index.
+
+        Generates embeddings for all items and stores them in memory.
+
+        Args:
+            items: List of (item_id, content) tuples to index
+
+        Raises:
+            RuntimeError: If embedding generation fails
+        """
+        if not items:
+            self._item_ids = []
+            self._item_embeddings = []
+            self._item_contents = {}
+            self._fitted = False
+            logger.debug("Embedding index cleared (no items)")
+            return
+
+        from gobby.search.embeddings import generate_embeddings
+
+        # Store contents for potential reindexing
+        self._item_ids = [item_id for item_id, _ in items]
+        self._item_contents = dict(items)
+        contents = [content for _, content in items]
+
+        # Generate embeddings in batch
+        try:
+            self._item_embeddings = await generate_embeddings(
+                texts=contents,
+                model=self._model,
+                api_base=self._api_base,
+                api_key=self._api_key,
+            )
+            self._fitted = True
+            logger.info(f"Embedding index built with {len(items)} items")
+        except Exception as e:
+            # Clear stale state to prevent inconsistent data
+            self._item_ids = []
+            self._item_contents = {}
+            self._item_embeddings = []
+            self._fitted = False
+            logger.error(f"Failed to build embedding index: {e}")
+            raise
+
+    async def search_async(
+        self,
+        query: str,
+        top_k: int = 10,
+    ) -> list[tuple[str, float]]:
+        """Search for items matching the query.
+
+        Generates an embedding for the query and finds items with
+        highest cosine similarity.
+
+        Args:
+            query: Search query text
+            top_k: Maximum number of results to return
+
+        Returns:
+            List of (item_id, similarity_score) tuples, sorted by
+            similarity descending.
+
+        Raises:
+            RuntimeError: If embedding generation fails
+        """
+        if not self._fitted or not self._item_embeddings:
+            return []
+
+        from gobby.search.embeddings import generate_embedding
+
+        # Generate query embedding
+        try:
+            query_embedding = await generate_embedding(
+                text=query,
+                model=self._model,
+                api_base=self._api_base,
+                api_key=self._api_key,
+            )
+        except Exception as e:
+            logger.error(f"Failed to embed query: {e}")
+            raise
+
+        # Compute similarities
+        similarities: list[tuple[str, float]] = []
+        for item_id, item_embedding in zip(self._item_ids, self._item_embeddings, strict=True):
+            similarity = _cosine_similarity(query_embedding, item_embedding)
+            if similarity > 0:
+                similarities.append((item_id, similarity))
+
+        # Sort by similarity descending
+        similarities.sort(key=lambda x: x[1], reverse=True)
+
+        return similarities[:top_k]
+
+    def needs_refit(self) -> bool:
+        """Check if the search index needs rebuilding."""
+        return not self._fitted
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get statistics about the search index."""
+        return {
+            "backend_type": "embedding",
+            "fitted": self._fitted,
+            "item_count": len(self._item_ids),
+            "model": self._model,
+            "has_api_base": self._api_base is not None,
+        }
+
+    def clear(self) -> None:
+        """Clear the search index."""
+        self._item_ids = []
+        self._item_embeddings = []
+        self._item_contents = {}
+        self._fitted = False
+
+    def get_item_contents(self) -> dict[str, str]:
+        """Get stored item contents.
+
+        Useful for reindexing into a different backend (e.g., TF-IDF fallback).
+
+        Returns:
+            Dict mapping item_id to content
+        """
+        return self._item_contents.copy()

gobby/search/embeddings.py

@@ -0,0 +1,238 @@
+"""LiteLLM-based embedding generation.
+
+This module provides a unified interface for generating embeddings using
+LiteLLM, which supports multiple providers through a single API:
+
+| Provider   | Model Format                    | Config                          |
+|------------|--------------------------------|--------------------------------|
+| OpenAI     | text-embedding-3-small         | OPENAI_API_KEY                  |
+| Ollama     | openai/nomic-embed-text        | api_base=http://localhost:11434/v1 |
+| Azure      | azure/azure-embedding-model    | api_base, api_key, api_version  |
+| Vertex AI  | vertex_ai/text-embedding-004   | GCP credentials                 |
+| Gemini     | gemini/text-embedding-004      | GEMINI_API_KEY                  |
+| Mistral    | mistral/mistral-embed          | MISTRAL_API_KEY                 |
+
+Example usage:
+    from gobby.search.embeddings import generate_embeddings, is_embedding_available
+
+    if is_embedding_available("text-embedding-3-small"):
+        embeddings = await generate_embeddings(
+            texts=["hello world", "foo bar"],
+            model="text-embedding-3-small"
+        )
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from gobby.search.models import SearchConfig
+
+logger = logging.getLogger(__name__)
+
+
+async def generate_embeddings(
+    texts: list[str],
+    model: str = "text-embedding-3-small",
+    api_base: str | None = None,
+    api_key: str | None = None,
+) -> list[list[float]]:
+    """Generate embeddings using LiteLLM.
+
+    Supports OpenAI, Ollama, Azure, Gemini, Mistral and other providers
+    through LiteLLM's unified API.
+
+    Args:
+        texts: List of texts to embed
+        model: LiteLLM model string (e.g., "text-embedding-3-small",
+               "openai/nomic-embed-text" for Ollama)
+        api_base: Optional API base URL for custom endpoints (e.g., Ollama)
+        api_key: Optional API key (uses environment variable if not set)
+
+    Returns:
+        List of embedding vectors (one per input text). Returns an empty
+        list if the input texts list is empty.
+
+    Raises:
+        RuntimeError: If LiteLLM is not installed or embedding fails
+    """
+    if not texts:
+        return []
+
+    try:
+        import litellm
+        from litellm.exceptions import (
+            AuthenticationError,
+            ContextWindowExceededError,
+            NotFoundError,
+            RateLimitError,
+        )
+    except ImportError as e:
+        raise RuntimeError("litellm package not installed. Run: uv add litellm") from e
+
+    # Build kwargs for LiteLLM
+    kwargs: dict[str, str | list[str]] = {
+        "model": model,
+        "input": texts,
+    }
+
+    if api_key:
+        kwargs["api_key"] = api_key
+
+    if api_base:
+        kwargs["api_base"] = api_base
+
+    try:
+        response = await litellm.aembedding(**kwargs)
+        embeddings: list[list[float]] = [item["embedding"] for item in response.data]
+        logger.debug(f"Generated {len(embeddings)} embeddings via LiteLLM ({model})")
+        return embeddings
+    except AuthenticationError as e:
+        logger.error(f"LiteLLM authentication failed: {e}")
+        raise RuntimeError(f"Authentication failed: {e}") from e
+    except NotFoundError as e:
+        logger.error(f"LiteLLM model not found: {e}")
+        raise RuntimeError(f"Model not found: {e}") from e
+    except RateLimitError as e:
+        logger.error(f"LiteLLM rate limit exceeded: {e}")
+        raise RuntimeError(f"Rate limit exceeded: {e}") from e
+    except ContextWindowExceededError as e:
+        logger.error(f"LiteLLM context window exceeded: {e}")
+        raise RuntimeError(f"Context window exceeded: {e}") from e
+    except Exception as e:
+        logger.error(f"Failed to generate embeddings with LiteLLM: {e}")
+        raise RuntimeError(f"Embedding generation failed: {e}") from e
+
+
+async def generate_embedding(
+    text: str,
+    model: str = "text-embedding-3-small",
+    api_base: str | None = None,
+    api_key: str | None = None,
+) -> list[float]:
+    """Generate embedding for a single text.
+
+    Convenience wrapper around generate_embeddings for single texts.
+
+    Args:
+        text: Text to embed
+        model: LiteLLM model string
+        api_base: Optional API base URL
+        api_key: Optional API key
+
+    Returns:
+        Embedding vector as list of floats
+
+    Raises:
+        RuntimeError: If embedding generation fails
+    """
+    embeddings = await generate_embeddings(
+        texts=[text],
+        model=model,
+        api_base=api_base,
+        api_key=api_key,
+    )
+    if not embeddings:
+        raise RuntimeError(
+            f"Embedding API returned empty result for model={model}, "
+            f"api_base={api_base}, api_key={'[set]' if api_key else '[not set]'}"
+        )
+    return embeddings[0]
+
+
+def is_embedding_available(
+    model: str = "text-embedding-3-small",
+    api_key: str | None = None,
+    api_base: str | None = None,
+) -> bool:
+    """Check if embedding is available for the given model.
+
+    For local models (Ollama), assumes availability if api_base is set.
+    For cloud models, requires an API key.
+
+    Args:
+        model: LiteLLM model string
+        api_key: Optional explicit API key
+        api_base: Optional API base URL
+
+    Returns:
+        True if embeddings can be generated, False otherwise
+    """
+    # Local models with api_base (Ollama, custom endpoints) are assumed available
+    if api_base:
+        return True
+
+    # Check for Ollama-style models that use local endpoints
+    if model.startswith("ollama/"):
+        # Native Ollama models - assume available locally
+        # In practice, we'll catch connection errors at runtime
+        return True
+
+    # openai/ prefix models require OpenAI API key
+    if model.startswith("openai/"):
+        effective_key = api_key or os.environ.get("OPENAI_API_KEY")
+        return effective_key is not None and len(effective_key) > 0
+
+    # Cloud models need API key
+    effective_key = api_key
+
+    # Check environment variables based on model prefix
+    if not effective_key:
+        if model.startswith("gemini/"):
+            effective_key = os.environ.get("GEMINI_API_KEY")
+        elif model.startswith("mistral/"):
+            effective_key = os.environ.get("MISTRAL_API_KEY")
+        elif model.startswith("azure/"):
+            effective_key = os.environ.get("AZURE_API_KEY")
+        elif model.startswith("vertex_ai/"):
+            # Vertex AI uses GCP credentials, check for project
+            effective_key = os.environ.get("VERTEXAI_PROJECT")
+        else:
+            # Default to OpenAI
+            effective_key = os.environ.get("OPENAI_API_KEY")
+
+    return effective_key is not None and len(effective_key) > 0
+
+
+def is_embedding_available_for_config(config: SearchConfig) -> bool:
+    """Check if embedding is available for a SearchConfig.
+
+    Convenience wrapper that extracts config values.
+
+    Args:
+        config: SearchConfig to check
+
+    Returns:
+        True if embeddings can be generated, False otherwise
+    """
+    return is_embedding_available(
+        model=config.embedding_model,
+        api_key=config.embedding_api_key,
+        api_base=config.embedding_api_base,
+    )
+
+
+async def generate_embeddings_for_config(
+    texts: list[str],
+    config: SearchConfig,
+) -> list[list[float]]:
+    """Generate embeddings using a SearchConfig.
+
+    Convenience wrapper that extracts config values.
+
+    Args:
+        texts: List of texts to embed
+        config: SearchConfig with model and API settings
+
+    Returns:
+        List of embedding vectors
+    """
+    return await generate_embeddings(
+        texts=texts,
+        model=config.embedding_model,
+        api_base=config.embedding_api_base,
+        api_key=config.embedding_api_key,
+    )