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

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())

agent_brain_server/providers/summarization/__init__.py

@@ -0,0 +1,41 @@
+"""Summarization providers for Agent Brain.
+
+This module provides summarization/LLM implementations for:
+- Anthropic (Claude 4.5 Haiku, Sonnet, Opus)
+- OpenAI (GPT-5, GPT-5 Mini)
+- Gemini (gemini-3-flash, gemini-3-pro)
+- Grok (grok-4, grok-4-fast)
+- Ollama (llama4:scout, mistral-small3.2, qwen3-coder, gemma3)
+"""
+
+from agent_brain_server.providers.factory import ProviderRegistry
+from agent_brain_server.providers.summarization.anthropic import (
+    AnthropicSummarizationProvider,
+)
+from agent_brain_server.providers.summarization.gemini import (
+    GeminiSummarizationProvider,
+)
+from agent_brain_server.providers.summarization.grok import GrokSummarizationProvider
+from agent_brain_server.providers.summarization.ollama import (
+    OllamaSummarizationProvider,
+)
+from agent_brain_server.providers.summarization.openai import (
+    OpenAISummarizationProvider,
+)
+
+# Register summarization providers
+ProviderRegistry.register_summarization_provider(
+    "anthropic", AnthropicSummarizationProvider
+)
+ProviderRegistry.register_summarization_provider("openai", OpenAISummarizationProvider)
+ProviderRegistry.register_summarization_provider("gemini", GeminiSummarizationProvider)
+ProviderRegistry.register_summarization_provider("grok", GrokSummarizationProvider)
+ProviderRegistry.register_summarization_provider("ollama", OllamaSummarizationProvider)
+
+__all__ = [
+    "AnthropicSummarizationProvider",
+    "OpenAISummarizationProvider",
+    "GeminiSummarizationProvider",
+    "GrokSummarizationProvider",
+    "OllamaSummarizationProvider",
+]

agent_brain_server/providers/summarization/anthropic.py

@@ -0,0 +1,87 @@
+"""Anthropic (Claude) summarization provider implementation."""
+
+import logging
+from typing import TYPE_CHECKING
+
+from anthropic import AsyncAnthropic
+
+from agent_brain_server.providers.base import BaseSummarizationProvider
+from agent_brain_server.providers.exceptions import AuthenticationError, ProviderError
+
+if TYPE_CHECKING:
+    from agent_brain_server.config.provider_config import SummarizationConfig
+
+logger = logging.getLogger(__name__)
+
+
+class AnthropicSummarizationProvider(BaseSummarizationProvider):
+    """Anthropic (Claude) summarization provider.
+
+    Supports:
+    - claude-haiku-4-5-20251001 (fast, cost-effective)
+    - claude-sonnet-4-5-20250514 (balanced)
+    - claude-opus-4-5-20251101 (highest quality)
+    - And other Claude models
+    """
+
+    def __init__(self, config: "SummarizationConfig") -> None:
+        """Initialize Anthropic summarization provider.
+
+        Args:
+            config: Summarization 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,
+            )
+
+        max_tokens = config.params.get("max_tokens", 300)
+        temperature = config.params.get("temperature", 0.1)
+        prompt_template = config.params.get("prompt_template")
+
+        super().__init__(
+            model=config.model,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            prompt_template=prompt_template,
+        )
+
+        self._client = AsyncAnthropic(api_key=api_key)
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name."""
+        return "Anthropic"
+
+    async def generate(self, prompt: str) -> str:
+        """Generate text based on prompt using Claude.
+
+        Args:
+            prompt: The prompt to send to Claude
+
+        Returns:
+            Generated text response
+
+        Raises:
+            ProviderError: If generation fails
+        """
+        try:
+            response = await self._client.messages.create(
+                model=self._model,
+                max_tokens=self._max_tokens,
+                temperature=self._temperature,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            # Extract text from response
+            return response.content[0].text  # type: ignore[union-attr]
+        except Exception as e:
+            raise ProviderError(
+                f"Failed to generate text: {e}",
+                self.provider_name,
+                cause=e,
+            ) from e

agent_brain_server/providers/summarization/gemini.py

@@ -0,0 +1,96 @@
+"""Google Gemini summarization provider implementation."""
+
+import logging
+from typing import TYPE_CHECKING
+
+import google.generativeai as genai
+
+from agent_brain_server.providers.base import BaseSummarizationProvider
+from agent_brain_server.providers.exceptions import AuthenticationError, ProviderError
+
+if TYPE_CHECKING:
+    from agent_brain_server.config.provider_config import SummarizationConfig
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiSummarizationProvider(BaseSummarizationProvider):
+    """Google Gemini summarization provider.
+
+    Supports:
+    - gemini-3-flash (fast, cost-effective)
+    - gemini-3-pro (highest quality)
+    - And other Gemini models
+    """
+
+    def __init__(self, config: "SummarizationConfig") -> None:
+        """Initialize Gemini summarization provider.
+
+        Args:
+            config: Summarization 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,
+            )
+
+        max_tokens = config.params.get(
+            "max_output_tokens", config.params.get("max_tokens", 300)
+        )
+        temperature = config.params.get("temperature", 0.1)
+        prompt_template = config.params.get("prompt_template")
+        top_p = config.params.get("top_p", 0.95)
+
+        super().__init__(
+            model=config.model,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            prompt_template=prompt_template,
+        )
+
+        # Configure Gemini with API key
+        genai.configure(api_key=api_key)  # type: ignore[attr-defined]
+
+        # Create model with generation config
+        generation_config = genai.types.GenerationConfig(
+            max_output_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+        )
+        self._model_instance = genai.GenerativeModel(  # type: ignore[attr-defined]
+            model_name=config.model,
+            generation_config=generation_config,
+        )
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name."""
+        return "Gemini"
+
+    async def generate(self, prompt: str) -> str:
+        """Generate text based on prompt using Gemini.
+
+        Args:
+            prompt: The prompt to send to Gemini
+
+        Returns:
+            Generated text response
+
+        Raises:
+            ProviderError: If generation fails
+        """
+        try:
+            # Use async generation
+            response = await self._model_instance.generate_content_async(prompt)
+            return str(response.text)
+        except Exception as e:
+            raise ProviderError(
+                f"Failed to generate text: {e}",
+                self.provider_name,
+                cause=e,
+            ) from e

agent_brain_server/providers/summarization/grok.py

@@ -0,0 +1,95 @@
+"""xAI Grok summarization provider implementation."""
+
+import logging
+from typing import TYPE_CHECKING
+
+from openai import AsyncOpenAI
+
+from agent_brain_server.providers.base import BaseSummarizationProvider
+from agent_brain_server.providers.exceptions import AuthenticationError, ProviderError
+
+if TYPE_CHECKING:
+    from agent_brain_server.config.provider_config import SummarizationConfig
+
+logger = logging.getLogger(__name__)
+
+
+class GrokSummarizationProvider(BaseSummarizationProvider):
+    """xAI Grok summarization provider.
+
+    Uses OpenAI-compatible API at https://api.x.ai/v1
+
+    Supports:
+    - grok-4 (most capable, with reasoning)
+    - grok-4-fast (fast variant)
+    - grok-3 (previous generation)
+    - And other Grok models
+    """
+
+    def __init__(self, config: "SummarizationConfig") -> None:
+        """Initialize Grok summarization provider.
+
+        Args:
+            config: Summarization 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,
+            )
+
+        max_tokens = config.params.get("max_tokens", 300)
+        temperature = config.params.get("temperature", 0.1)
+        prompt_template = config.params.get("prompt_template")
+
+        super().__init__(
+            model=config.model,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            prompt_template=prompt_template,
+        )
+
+        # Grok uses OpenAI-compatible API
+        base_url = config.get_base_url() or "https://api.x.ai/v1"
+        self._client = AsyncOpenAI(
+            api_key=api_key,
+            base_url=base_url,
+        )
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name."""
+        return "Grok"
+
+    async def generate(self, prompt: str) -> str:
+        """Generate text based on prompt using Grok.
+
+        Args:
+            prompt: The prompt to send to Grok
+
+        Returns:
+            Generated text response
+
+        Raises:
+            ProviderError: If generation fails
+        """
+        try:
+            response = await self._client.chat.completions.create(
+                model=self._model,
+                max_tokens=self._max_tokens,
+                temperature=self._temperature,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            # Extract text from response
+            content = response.choices[0].message.content
+            return content if content else ""
+        except Exception as e:
+            raise ProviderError(
+                f"Failed to generate text: {e}",
+                self.provider_name,
+                cause=e,
+            ) from e