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

agent_brain_server/providers/__init__.py

@@ -0,0 +1,64 @@
+"""Pluggable model providers for Agent Brain.
+
+This package provides abstractions for embedding and summarization providers,
+allowing configuration-driven selection between OpenAI, Ollama, Cohere (embeddings)
+and Anthropic, OpenAI, Gemini, Grok, Ollama (summarization) providers.
+"""
+
+from agent_brain_server.providers.base import (
+    BaseEmbeddingProvider,
+    BaseSummarizationProvider,
+    EmbeddingProvider,
+    EmbeddingProviderType,
+    SummarizationProvider,
+    SummarizationProviderType,
+)
+from agent_brain_server.providers.exceptions import (
+    AuthenticationError,
+    ConfigurationError,
+    ModelNotFoundError,
+    ProviderError,
+    ProviderMismatchError,
+    ProviderNotFoundError,
+    RateLimitError,
+)
+from agent_brain_server.providers.factory import ProviderRegistry
+
+__all__ = [
+    # Protocols
+    "EmbeddingProvider",
+    "SummarizationProvider",
+    # Base classes
+    "BaseEmbeddingProvider",
+    "BaseSummarizationProvider",
+    # Enums
+    "EmbeddingProviderType",
+    "SummarizationProviderType",
+    # Factory
+    "ProviderRegistry",
+    # Exceptions
+    "ProviderError",
+    "ConfigurationError",
+    "AuthenticationError",
+    "ProviderNotFoundError",
+    "ProviderMismatchError",
+    "RateLimitError",
+    "ModelNotFoundError",
+]
+
+
+def _register_providers() -> None:
+    """Register all built-in providers with the registry."""
+    # Import providers to trigger registration
+    from agent_brain_server.providers import (
+        embedding,  # noqa: F401
+        summarization,  # noqa: F401
+    )
+
+    # Silence unused import warnings
+    _ = embedding
+    _ = summarization
+
+
+# Auto-register providers on import
+_register_providers()

agent_brain_server/providers/base.py

@@ -0,0 +1,251 @@
+"""Base protocols and classes for pluggable providers."""
+
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Awaitable, Callable
+from enum import Enum
+from typing import Optional, Protocol, runtime_checkable
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingProviderType(str, Enum):
+    """Supported embedding providers."""
+
+    OPENAI = "openai"
+    OLLAMA = "ollama"
+    COHERE = "cohere"
+
+
+class SummarizationProviderType(str, Enum):
+    """Supported summarization providers."""
+
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    GEMINI = "gemini"
+    GROK = "grok"
+    OLLAMA = "ollama"
+
+
+@runtime_checkable
+class EmbeddingProvider(Protocol):
+    """Protocol for embedding providers.
+
+    All embedding providers must implement this interface to be usable
+    by the Agent Brain indexing and query systems.
+    """
+
+    async def embed_text(self, text: str) -> list[float]:
+        """Generate embedding for a single text.
+
+        Args:
+            text: Text to embed.
+
+        Returns:
+            Embedding vector as list of floats.
+
+        Raises:
+            ProviderError: If embedding generation fails.
+        """
+        ...
+
+    async def embed_texts(
+        self,
+        texts: list[str],
+        progress_callback: Optional[Callable[[int, int], Awaitable[None]]] = None,
+    ) -> list[list[float]]:
+        """Generate embeddings for multiple texts.
+
+        Args:
+            texts: List of texts to embed.
+            progress_callback: Optional callback(processed, total) for progress.
+
+        Returns:
+            List of embedding vectors, one per input text.
+
+        Raises:
+            ProviderError: If embedding generation fails.
+        """
+        ...
+
+    def get_dimensions(self) -> int:
+        """Get the embedding vector dimensions for the current model.
+
+        Returns:
+            Number of dimensions in the embedding vector.
+        """
+        ...
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name for logging."""
+        ...
+
+    @property
+    def model_name(self) -> str:
+        """Model identifier being used."""
+        ...
+
+
+@runtime_checkable
+class SummarizationProvider(Protocol):
+    """Protocol for summarization/LLM providers.
+
+    All summarization providers must implement this interface to be usable
+    by the Agent Brain code summarization system.
+    """
+
+    async def summarize(self, text: str) -> str:
+        """Generate a summary of the given text.
+
+        Args:
+            text: Text to summarize (typically source code).
+
+        Returns:
+            Natural language summary of the text.
+
+        Raises:
+            ProviderError: If summarization fails.
+        """
+        ...
+
+    async def generate(self, prompt: str) -> str:
+        """Generate text based on a prompt (generic LLM call).
+
+        Args:
+            prompt: The prompt to send to the LLM.
+
+        Returns:
+            Generated text response.
+
+        Raises:
+            ProviderError: If generation fails.
+        """
+        ...
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name for logging."""
+        ...
+
+    @property
+    def model_name(self) -> str:
+        """Model identifier being used."""
+        ...
+
+
+class BaseEmbeddingProvider(ABC):
+    """Base class for embedding providers with common functionality."""
+
+    def __init__(self, model: str, batch_size: int = 100) -> None:
+        self._model = model
+        self._batch_size = batch_size
+        logger.info(
+            f"Initialized {self.provider_name} embedding provider with model {model}"
+        )
+
+    @property
+    def model_name(self) -> str:
+        """Model identifier being used."""
+        return self._model
+
+    async def embed_texts(
+        self,
+        texts: list[str],
+        progress_callback: Optional[Callable[[int, int], Awaitable[None]]] = None,
+    ) -> list[list[float]]:
+        """Default batch implementation using _embed_batch."""
+        if not texts:
+            return []
+
+        all_embeddings: list[list[float]] = []
+
+        for i in range(0, len(texts), self._batch_size):
+            batch = texts[i : i + self._batch_size]
+            batch_embeddings = await self._embed_batch(batch)
+            all_embeddings.extend(batch_embeddings)
+
+            if progress_callback:
+                await progress_callback(
+                    min(i + self._batch_size, len(texts)),
+                    len(texts),
+                )
+
+            logger.debug(
+                f"Generated embeddings for batch {i // self._batch_size + 1} "
+                f"({len(batch)} texts)"
+            )
+
+        return all_embeddings
+
+    @abstractmethod
+    async def _embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Provider-specific batch embedding implementation."""
+        ...
+
+    @abstractmethod
+    async def embed_text(self, text: str) -> list[float]:
+        """Provider-specific single text embedding."""
+        ...
+
+    @abstractmethod
+    def get_dimensions(self) -> int:
+        """Provider-specific dimension lookup."""
+        ...
+
+    @property
+    @abstractmethod
+    def provider_name(self) -> str:
+        """Human-readable provider name for logging."""
+        ...
+
+
+class BaseSummarizationProvider(ABC):
+    """Base class for summarization providers with common functionality."""
+
+    DEFAULT_PROMPT_TEMPLATE = (
+        "You are an expert software engineer analyzing source code. "
+        "Provide a concise 1-2 sentence summary of what this code does. "
+        "Focus on the functionality, purpose, and behavior. "
+        "Be specific about inputs, outputs, and side effects. "
+        "Ignore implementation details and focus on what the code accomplishes.\n\n"
+        "Code to summarize:\n{code}\n\n"
+        "Summary:"
+    )
+
+    def __init__(
+        self,
+        model: str,
+        max_tokens: int = 300,
+        temperature: float = 0.1,
+        prompt_template: Optional[str] = None,
+    ) -> None:
+        self._model = model
+        self._max_tokens = max_tokens
+        self._temperature = temperature
+        self._prompt_template = prompt_template or self.DEFAULT_PROMPT_TEMPLATE
+        logger.info(
+            f"Initialized {self.provider_name} summarization provider "
+            f"with model {model}"
+        )
+
+    @property
+    def model_name(self) -> str:
+        """Model identifier being used."""
+        return self._model
+
+    async def summarize(self, text: str) -> str:
+        """Generate summary using the prompt template."""
+        prompt = self._prompt_template.format(code=text)
+        return await self.generate(prompt)
+
+    @abstractmethod
+    async def generate(self, prompt: str) -> str:
+        """Provider-specific text generation."""
+        ...
+
+    @property
+    @abstractmethod
+    def provider_name(self) -> str:
+        """Human-readable provider name for logging."""
+        ...

agent_brain_server/providers/embedding/__init__.py

@@ -0,0 +1,23 @@
+"""Embedding providers for Agent Brain.
+
+This module provides embedding implementations for:
+- OpenAI (text-embedding-3-large, text-embedding-3-small, text-embedding-ada-002)
+- Ollama (nomic-embed-text, mxbai-embed-large, etc.)
+- Cohere (embed-english-v3, embed-multilingual-v3, etc.)
+"""
+
+from agent_brain_server.providers.embedding.cohere import CohereEmbeddingProvider
+from agent_brain_server.providers.embedding.ollama import OllamaEmbeddingProvider
+from agent_brain_server.providers.embedding.openai import OpenAIEmbeddingProvider
+from agent_brain_server.providers.factory import ProviderRegistry
+
+# Register embedding providers
+ProviderRegistry.register_embedding_provider("openai", OpenAIEmbeddingProvider)
+ProviderRegistry.register_embedding_provider("ollama", OllamaEmbeddingProvider)
+ProviderRegistry.register_embedding_provider("cohere", CohereEmbeddingProvider)
+
+__all__ = [
+    "OpenAIEmbeddingProvider",
+    "OllamaEmbeddingProvider",
+    "CohereEmbeddingProvider",
+]

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