agent-brain-rag 1.2.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

agent_brain_server/providers/embedding/cohere.py

@@ -0,0 +1,163 @@
+"""Cohere embedding provider implementation."""
+
+import logging
+from typing import TYPE_CHECKING
+
+import cohere
+
+from agent_brain_server.providers.base import BaseEmbeddingProvider
+from agent_brain_server.providers.exceptions import AuthenticationError, ProviderError
+
+if TYPE_CHECKING:
+    from agent_brain_server.config.provider_config import EmbeddingConfig
+
+logger = logging.getLogger(__name__)
+
+# Model dimension mappings for Cohere embedding models
+COHERE_MODEL_DIMENSIONS: dict[str, int] = {
+    "embed-english-v3.0": 1024,
+    "embed-english-light-v3.0": 384,
+    "embed-multilingual-v3.0": 1024,
+    "embed-multilingual-light-v3.0": 384,
+    "embed-english-v2.0": 4096,
+    "embed-english-light-v2.0": 1024,
+    "embed-multilingual-v2.0": 768,
+}
+
+DEFAULT_COHERE_DIMENSIONS = 1024
+
+
+class CohereEmbeddingProvider(BaseEmbeddingProvider):
+    """Cohere embedding provider using Cohere's embedding models.
+
+    Supports:
+    - embed-english-v3.0 (1024 dimensions, best for English)
+    - embed-english-light-v3.0 (384 dimensions, faster)
+    - embed-multilingual-v3.0 (1024 dimensions, 100+ languages)
+    - embed-multilingual-light-v3.0 (384 dimensions, faster multilingual)
+
+    Cohere embeddings support different input types for optimal performance:
+    - search_document: For indexing documents to be searched
+    - search_query: For search queries
+    - classification: For classification tasks
+    - clustering: For clustering tasks
+    """
+
+    def __init__(self, config: "EmbeddingConfig") -> None:
+        """Initialize Cohere embedding provider.
+
+        Args:
+            config: Embedding configuration
+
+        Raises:
+            AuthenticationError: If API key is not available
+        """
+        api_key = config.get_api_key()
+        if not api_key:
+            raise AuthenticationError(
+                f"Missing API key. Set {config.api_key_env} environment variable.",
+                self.provider_name,
+            )
+
+        batch_size = config.params.get("batch_size", 96)  # Cohere limit
+        super().__init__(model=config.model, batch_size=batch_size)
+
+        self._client = cohere.AsyncClientV2(api_key=api_key)
+        self._input_type = config.params.get("input_type", "search_document")
+        self._truncate = config.params.get("truncate", "END")
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name."""
+        return "Cohere"
+
+    def get_dimensions(self) -> int:
+        """Get embedding dimensions for current model.
+
+        Returns:
+            Number of dimensions in embedding vector
+        """
+        return COHERE_MODEL_DIMENSIONS.get(self._model, DEFAULT_COHERE_DIMENSIONS)
+
+    async def embed_text(self, text: str) -> list[float]:
+        """Generate embedding for single text.
+
+        Args:
+            text: Text to embed
+
+        Returns:
+            Embedding vector as list of floats
+
+        Raises:
+            ProviderError: If embedding generation fails
+        """
+        try:
+            response = await self._client.embed(
+                texts=[text],
+                model=self._model,
+                input_type=self._input_type,
+                truncate=self._truncate,
+            )
+            embeddings = response.embeddings.float_
+            if embeddings is None:
+                raise ProviderError(
+                    "No embeddings returned from Cohere",
+                    self.provider_name,
+                )
+            return list(embeddings[0])
+        except Exception as e:
+            raise ProviderError(
+                f"Failed to generate embedding: {e}",
+                self.provider_name,
+                cause=e,
+            ) from e
+
+    async def _embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for a batch of texts.
+
+        Args:
+            texts: List of texts to embed
+
+        Returns:
+            List of embedding vectors
+
+        Raises:
+            ProviderError: If embedding generation fails
+        """
+        try:
+            response = await self._client.embed(
+                texts=texts,
+                model=self._model,
+                input_type=self._input_type,
+                truncate=self._truncate,
+            )
+            embeddings = response.embeddings.float_
+            if embeddings is None:
+                raise ProviderError(
+                    "No embeddings returned from Cohere",
+                    self.provider_name,
+                )
+            return [list(emb) for emb in embeddings]
+        except Exception as e:
+            raise ProviderError(
+                f"Failed to generate batch embeddings: {e}",
+                self.provider_name,
+                cause=e,
+            ) from e
+
+    def set_input_type(self, input_type: str) -> None:
+        """Set the input type for embeddings.
+
+        Args:
+            input_type: One of 'search_document', 'search_query',
+                       'classification', or 'clustering'
+        """
+        valid_types = [
+            "search_document",
+            "search_query",
+            "classification",
+            "clustering",
+        ]
+        if input_type not in valid_types:
+            raise ValueError(f"Invalid input_type. Must be one of: {valid_types}")
+        self._input_type = input_type

agent_brain_server/providers/embedding/ollama.py

@@ -0,0 +1,150 @@
+"""Ollama embedding provider implementation."""
+
+import logging
+from typing import TYPE_CHECKING
+
+from openai import AsyncOpenAI
+
+from agent_brain_server.providers.base import BaseEmbeddingProvider
+from agent_brain_server.providers.exceptions import (
+    OllamaConnectionError,
+    ProviderError,
+)
+
+if TYPE_CHECKING:
+    from agent_brain_server.config.provider_config import EmbeddingConfig
+
+logger = logging.getLogger(__name__)
+
+# Model dimension mappings for common Ollama embedding models
+OLLAMA_MODEL_DIMENSIONS: dict[str, int] = {
+    "nomic-embed-text": 768,
+    "mxbai-embed-large": 1024,
+    "all-minilm": 384,
+    "snowflake-arctic-embed": 1024,
+    "bge-m3": 1024,
+    "bge-large": 1024,
+}
+
+DEFAULT_OLLAMA_DIMENSIONS = 768
+
+
+class OllamaEmbeddingProvider(BaseEmbeddingProvider):
+    """Ollama embedding provider using local models.
+
+    Uses OpenAI-compatible API endpoint provided by Ollama.
+
+    Supports:
+    - nomic-embed-text (768 dimensions, general purpose)
+    - mxbai-embed-large (1024 dimensions, multilingual)
+    - all-minilm (384 dimensions, lightweight)
+    - snowflake-arctic-embed (1024 dimensions, high quality)
+    - And any other embedding model available in Ollama
+    """
+
+    def __init__(self, config: "EmbeddingConfig") -> None:
+        """Initialize Ollama embedding provider.
+
+        Args:
+            config: Embedding configuration
+
+        Note:
+            Ollama does not require an API key as it runs locally.
+        """
+        batch_size = config.params.get("batch_size", 100)
+        super().__init__(model=config.model, batch_size=batch_size)
+
+        # Ollama uses OpenAI-compatible API
+        base_url = config.get_base_url() or "http://localhost:11434/v1"
+        self._base_url = base_url
+        self._client = AsyncOpenAI(
+            api_key="ollama",  # Ollama doesn't need real key
+            base_url=base_url,
+        )
+
+        # Optional parameters
+        self._num_ctx = config.params.get("num_ctx", 2048)
+        self._num_threads = config.params.get("num_threads")
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name."""
+        return "Ollama"
+
+    def get_dimensions(self) -> int:
+        """Get embedding dimensions for current model.
+
+        Returns:
+            Number of dimensions in embedding vector
+        """
+        return OLLAMA_MODEL_DIMENSIONS.get(self._model, DEFAULT_OLLAMA_DIMENSIONS)
+
+    async def embed_text(self, text: str) -> list[float]:
+        """Generate embedding for single text.
+
+        Args:
+            text: Text to embed
+
+        Returns:
+            Embedding vector as list of floats
+
+        Raises:
+            OllamaConnectionError: If Ollama is not running
+            ProviderError: If embedding generation fails
+        """
+        try:
+            response = await self._client.embeddings.create(
+                model=self._model,
+                input=text,
+            )
+            return response.data[0].embedding
+        except Exception as e:
+            if "connection" in str(e).lower() or "refused" in str(e).lower():
+                raise OllamaConnectionError(self._base_url, cause=e) from e
+            raise ProviderError(
+                f"Failed to generate embedding: {e}",
+                self.provider_name,
+                cause=e,
+            ) from e
+
+    async def _embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for a batch of texts.
+
+        Args:
+            texts: List of texts to embed
+
+        Returns:
+            List of embedding vectors
+
+        Raises:
+            OllamaConnectionError: If Ollama is not running
+            ProviderError: If embedding generation fails
+        """
+        try:
+            response = await self._client.embeddings.create(
+                model=self._model,
+                input=texts,
+            )
+            return [item.embedding for item in response.data]
+        except Exception as e:
+            if "connection" in str(e).lower() or "refused" in str(e).lower():
+                raise OllamaConnectionError(self._base_url, cause=e) from e
+            raise ProviderError(
+                f"Failed to generate batch embeddings: {e}",
+                self.provider_name,
+                cause=e,
+            ) from e
+
+    async def health_check(self) -> bool:
+        """Check if Ollama is running and accessible.
+
+        Returns:
+            True if Ollama is healthy, False otherwise
+        """
+        try:
+            # Try to list models to verify connection
+            await self._client.models.list()
+            return True
+        except Exception as e:
+            logger.warning(f"Ollama health check failed: {e}")
+            return False

agent_brain_server/providers/embedding/openai.py

@@ -0,0 +1,118 @@
+"""OpenAI embedding provider implementation."""
+
+import logging
+from typing import TYPE_CHECKING
+
+from openai import AsyncOpenAI
+
+from agent_brain_server.providers.base import BaseEmbeddingProvider
+from agent_brain_server.providers.exceptions import AuthenticationError, ProviderError
+
+if TYPE_CHECKING:
+    from agent_brain_server.config.provider_config import EmbeddingConfig
+
+logger = logging.getLogger(__name__)
+
+# Model dimension mappings for OpenAI embedding models
+OPENAI_MODEL_DIMENSIONS: dict[str, int] = {
+    "text-embedding-3-large": 3072,
+    "text-embedding-3-small": 1536,
+    "text-embedding-ada-002": 1536,
+}
+
+
+class OpenAIEmbeddingProvider(BaseEmbeddingProvider):
+    """OpenAI embedding provider using text-embedding models.
+
+    Supports:
+    - text-embedding-3-large (3072 dimensions, highest quality)
+    - text-embedding-3-small (1536 dimensions, faster)
+    - text-embedding-ada-002 (1536 dimensions, legacy)
+    """
+
+    def __init__(self, config: "EmbeddingConfig") -> None:
+        """Initialize OpenAI embedding provider.
+
+        Args:
+            config: Embedding configuration
+
+        Raises:
+            AuthenticationError: If API key is not available
+        """
+        api_key = config.get_api_key()
+        if not api_key:
+            raise AuthenticationError(
+                f"Missing API key. Set {config.api_key_env} environment variable.",
+                self.provider_name,
+            )
+
+        batch_size = config.params.get("batch_size", 100)
+        super().__init__(model=config.model, batch_size=batch_size)
+
+        self._client = AsyncOpenAI(api_key=api_key)
+        self._dimensions_override = config.params.get("dimensions")
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name."""
+        return "OpenAI"
+
+    def get_dimensions(self) -> int:
+        """Get embedding dimensions for current model.
+
+        Returns:
+            Number of dimensions in embedding vector
+        """
+        if self._dimensions_override:
+            return int(self._dimensions_override)
+        return OPENAI_MODEL_DIMENSIONS.get(self._model, 3072)
+
+    async def embed_text(self, text: str) -> list[float]:
+        """Generate embedding for single text.
+
+        Args:
+            text: Text to embed
+
+        Returns:
+            Embedding vector as list of floats
+
+        Raises:
+            ProviderError: If embedding generation fails
+        """
+        try:
+            response = await self._client.embeddings.create(
+                model=self._model,
+                input=text,
+            )
+            return response.data[0].embedding
+        except Exception as e:
+            raise ProviderError(
+                f"Failed to generate embedding: {e}",
+                self.provider_name,
+                cause=e,
+            ) from e
+
+    async def _embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for a batch of texts.
+
+        Args:
+            texts: List of texts to embed
+
+        Returns:
+            List of embedding vectors
+
+        Raises:
+            ProviderError: If embedding generation fails
+        """
+        try:
+            response = await self._client.embeddings.create(
+                model=self._model,
+                input=texts,
+            )
+            return [item.embedding for item in response.data]
+        except Exception as e:
+            raise ProviderError(
+                f"Failed to generate batch embeddings: {e}",
+                self.provider_name,
+                cause=e,
+            ) from e

agent_brain_server/providers/exceptions.py

@@ -0,0 +1,95 @@
+"""Exception hierarchy for provider errors."""
+
+from typing import Optional
+
+
+class ProviderError(Exception):
+    """Base exception for provider errors."""
+
+    def __init__(
+        self, message: str, provider: str, cause: Optional[Exception] = None
+    ) -> None:
+        self.provider = provider
+        self.cause = cause
+        super().__init__(f"[{provider}] {message}")
+
+
+class ConfigurationError(ProviderError):
+    """Raised when provider configuration is invalid."""
+
+    pass
+
+
+class AuthenticationError(ProviderError):
+    """Raised when API key is missing or invalid."""
+
+    pass
+
+
+class ProviderNotFoundError(ProviderError):
+    """Raised when requested provider type is not registered."""
+
+    pass
+
+
+class ProviderMismatchError(ProviderError):
+    """Raised when current provider doesn't match indexed data."""
+
+    def __init__(
+        self,
+        current_provider: str,
+        current_model: str,
+        indexed_provider: str,
+        indexed_model: str,
+    ) -> None:
+        message = (
+            f"Provider mismatch: index was created with "
+            f"{indexed_provider}/{indexed_model}, "
+            f"but current config uses {current_provider}/{current_model}. "
+            f"Re-index with --force to update."
+        )
+        super().__init__(message, current_provider)
+        self.current_model = current_model
+        self.indexed_provider = indexed_provider
+        self.indexed_model = indexed_model
+
+
+class RateLimitError(ProviderError):
+    """Raised when provider rate limit is hit."""
+
+    def __init__(self, provider: str, retry_after: Optional[int] = None) -> None:
+        self.retry_after = retry_after
+        message = "Rate limit exceeded"
+        if retry_after:
+            message += f", retry after {retry_after}s"
+        super().__init__(message, provider)
+
+
+class ModelNotFoundError(ProviderError):
+    """Raised when specified model is not available."""
+
+    def __init__(
+        self, provider: str, model: str, available_models: Optional[list[str]] = None
+    ) -> None:
+        self.model = model
+        self.available_models = available_models or []
+        if available_models:
+            message = (
+                f"Model '{model}' not found. "
+                f"Available: {', '.join(available_models[:5])}"
+            )
+        else:
+            message = f"Model '{model}' not found"
+        super().__init__(message, provider)
+
+
+class OllamaConnectionError(ProviderError):
+    """Raised when Ollama is not running or unreachable."""
+
+    def __init__(self, base_url: str, cause: Optional[Exception] = None) -> None:
+        message = (
+            f"Cannot connect to Ollama at {base_url}. "
+            "Ensure Ollama is running with 'ollama serve' command."
+        )
+        super().__init__(message, "ollama", cause)
+        self.base_url = base_url

agent_brain_server/providers/factory.py

@@ -0,0 +1,157 @@
+"""Provider factory and registry for dynamic provider instantiation."""
+
+import logging
+from typing import TYPE_CHECKING, Any, cast
+
+from agent_brain_server.providers.exceptions import ProviderNotFoundError
+
+if TYPE_CHECKING:
+    from agent_brain_server.config.provider_config import (
+        EmbeddingConfig,
+        SummarizationConfig,
+    )
+    from agent_brain_server.providers.base import (
+        EmbeddingProvider,
+        SummarizationProvider,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+class ProviderRegistry:
+    """Registry for provider implementations.
+
+    Allows dynamic registration of providers and lazy instantiation.
+    Implements singleton pattern for provider instance caching.
+    """
+
+    _embedding_providers: dict[str, type[Any]] = {}
+    _summarization_providers: dict[str, type[Any]] = {}
+    _instances: dict[str, Any] = {}
+
+    @classmethod
+    def register_embedding_provider(
+        cls,
+        provider_type: str,
+        provider_class: type["EmbeddingProvider"],
+    ) -> None:
+        """Register an embedding provider class.
+
+        Args:
+            provider_type: Provider identifier (e.g., 'openai', 'ollama')
+            provider_class: Provider class implementing EmbeddingProvider protocol
+        """
+        cls._embedding_providers[provider_type] = provider_class
+        logger.debug(f"Registered embedding provider: {provider_type}")
+
+    @classmethod
+    def register_summarization_provider(
+        cls,
+        provider_type: str,
+        provider_class: type["SummarizationProvider"],
+    ) -> None:
+        """Register a summarization provider class.
+
+        Args:
+            provider_type: Provider identifier (e.g., 'anthropic', 'openai')
+            provider_class: Provider class implementing SummarizationProvider protocol
+        """
+        cls._summarization_providers[provider_type] = provider_class
+        logger.debug(f"Registered summarization provider: {provider_type}")
+
+    @classmethod
+    def get_embedding_provider(cls, config: "EmbeddingConfig") -> "EmbeddingProvider":
+        """Get or create embedding provider instance.
+
+        Args:
+            config: Embedding provider configuration
+
+        Returns:
+            Configured EmbeddingProvider instance
+
+        Raises:
+            ProviderNotFoundError: If provider type is not registered
+        """
+        # Get provider type as string value
+        provider_type = (
+            config.provider.value
+            if hasattr(config.provider, "value")
+            else str(config.provider)
+        )
+        cache_key = f"embed:{provider_type}:{config.model}"
+
+        if cache_key not in cls._instances:
+            provider_class = cls._embedding_providers.get(provider_type)
+            if not provider_class:
+                available = list(cls._embedding_providers.keys())
+                raise ProviderNotFoundError(
+                    f"Unknown embedding provider: {provider_type}. "
+                    f"Available: {', '.join(available)}",
+                    provider_type,
+                )
+            cls._instances[cache_key] = provider_class(config)
+            logger.info(
+                f"Created {provider_type} embedding provider with model {config.model}"
+            )
+
+        from agent_brain_server.providers.base import EmbeddingProvider
+
+        return cast(EmbeddingProvider, cls._instances[cache_key])
+
+    @classmethod
+    def get_summarization_provider(
+        cls, config: "SummarizationConfig"
+    ) -> "SummarizationProvider":
+        """Get or create summarization provider instance.
+
+        Args:
+            config: Summarization provider configuration
+
+        Returns:
+            Configured SummarizationProvider instance
+
+        Raises:
+            ProviderNotFoundError: If provider type is not registered
+        """
+        # Get provider type as string value
+        provider_type = (
+            config.provider.value
+            if hasattr(config.provider, "value")
+            else str(config.provider)
+        )
+        cache_key = f"summ:{provider_type}:{config.model}"
+
+        if cache_key not in cls._instances:
+            provider_class = cls._summarization_providers.get(provider_type)
+            if not provider_class:
+                available = list(cls._summarization_providers.keys())
+                raise ProviderNotFoundError(
+                    f"Unknown summarization provider: {provider_type}. "
+                    f"Available: {', '.join(available)}",
+                    provider_type,
+                )
+            cls._instances[cache_key] = provider_class(config)
+            logger.info(
+                f"Created {provider_type} summarization provider "
+                f"with model {config.model}"
+            )
+
+        from agent_brain_server.providers.base import SummarizationProvider
+
+        return cast(SummarizationProvider, cls._instances[cache_key])
+
+    @classmethod
+    def clear_cache(cls) -> None:
+        """Clear provider instance cache (for testing)."""
+        cls._instances.clear()
+        logger.debug("Cleared provider instance cache")
+
+    @classmethod
+    def get_available_embedding_providers(cls) -> list[str]:
+        """Get list of registered embedding provider types."""
+        return list(cls._embedding_providers.keys())
+
+    @classmethod
+    def get_available_summarization_providers(cls) -> list[str]:
+        """Get list of registered summarization provider types."""
+        return list(cls._summarization_providers.keys())