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

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

agent_brain_server/providers/summarization/ollama.py

@@ -0,0 +1,114 @@
+"""Ollama 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 (
+    OllamaConnectionError,
+    ProviderError,
+)
+
+if TYPE_CHECKING:
+    from agent_brain_server.config.provider_config import SummarizationConfig
+
+logger = logging.getLogger(__name__)
+
+
+class OllamaSummarizationProvider(BaseSummarizationProvider):
+    """Ollama summarization provider using local models.
+
+    Uses OpenAI-compatible API endpoint provided by Ollama.
+
+    Supports:
+    - llama4:scout (Meta's Llama 4 Scout - lightweight, fast)
+    - mistral-small3.2 (Mistral Small 3.2 - balanced)
+    - qwen3-coder (Alibaba Qwen 3 Coder - code-focused)
+    - gemma3 (Google Gemma 3 - efficient)
+    - deepseek-coder-v3 (DeepSeek Coder V3)
+    - And any other chat model available in Ollama
+    """
+
+    def __init__(self, config: "SummarizationConfig") -> None:
+        """Initialize Ollama summarization provider.
+
+        Args:
+            config: Summarization configuration
+
+        Note:
+            Ollama does not require an API key as it runs locally.
+        """
+        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,
+        )
+
+        # 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", 4096)
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name."""
+        return "Ollama"
+
+    async def generate(self, prompt: str) -> str:
+        """Generate text based on prompt using Ollama.
+
+        Args:
+            prompt: The prompt to send to Ollama
+
+        Returns:
+            Generated text response
+
+        Raises:
+            OllamaConnectionError: If Ollama is not running
+            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:
+            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 text: {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/summarization/openai.py

@@ -0,0 +1,87 @@
+"""OpenAI (GPT) 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 OpenAISummarizationProvider(BaseSummarizationProvider):
+    """OpenAI (GPT) summarization provider.
+
+    Supports:
+    - gpt-5 (most capable)
+    - gpt-5-mini (fast, cost-effective)
+    - And other OpenAI chat models
+    """
+
+    def __init__(self, config: "SummarizationConfig") -> None:
+        """Initialize OpenAI 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 = AsyncOpenAI(api_key=api_key)
+
+    @property
+    def provider_name(self) -> str:
+        """Human-readable provider name."""
+        return "OpenAI"
+
+    async def generate(self, prompt: str) -> str:
+        """Generate text based on prompt using GPT.
+
+        Args:
+            prompt: The prompt to send to GPT
+
+        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

agent_brain_server/services/indexing_service.py

@@ -10,6 +10,7 @@ from typing import Any, Callable, Optional, Union
 
 from llama_index.core.schema import TextNode
 
+from agent_brain_server.config import settings
 from agent_brain_server.indexing import (
     BM25IndexManager,
     ContextAwareChunker,
@@ -18,6 +19,10 @@ from agent_brain_server.indexing import (
     get_bm25_manager,
 )
 from agent_brain_server.indexing.chunking import CodeChunk, CodeChunker, TextChunk
+from agent_brain_server.indexing.graph_index import (
+    GraphIndexManager,
+    get_graph_index_manager,
+)
 from agent_brain_server.models import IndexingState, IndexingStatusEnum, IndexRequest
 from agent_brain_server.storage import VectorStoreManager, get_vector_store
 
@@ -43,6 +48,7 @@ class IndexingService:
         chunker: Optional[ContextAwareChunker] = None,
         embedding_generator: Optional[EmbeddingGenerator] = None,
         bm25_manager: Optional[BM25IndexManager] = None,
+        graph_index_manager: Optional[GraphIndexManager] = None,
     ):
         """
         Initialize the indexing service.
@@ -53,12 +59,14 @@ class IndexingService:
             chunker: Text chunker instance.
             embedding_generator: Embedding generator instance.
             bm25_manager: BM25 index manager instance.
+            graph_index_manager: Graph index manager instance (Feature 113).
         """
         self.vector_store = vector_store or get_vector_store()
         self.document_loader = document_loader or DocumentLoader()
         self.chunker = chunker or ContextAwareChunker()
         self.embedding_generator = embedding_generator or EmbeddingGenerator()
         self.bm25_manager = bm25_manager or get_bm25_manager()
+        self.graph_index_manager = graph_index_manager or get_graph_index_manager()
 
         # Internal state
         self._state = IndexingState(
@@ -382,6 +390,21 @@ class IndexingService:
             ]
             self.bm25_manager.build_index(nodes)
 
+            # Step 6: Build graph index if enabled (Feature 113)
+            if settings.ENABLE_GRAPH_INDEX:
+                if progress_callback:
+                    await progress_callback(97, 100, "Building graph index...")
+
+                def graph_progress(current: int, total: int, message: str) -> None:
+                    # Synchronous callback wrapper
+                    logger.debug(f"Graph indexing: {message}")
+
+                triplet_count = self.graph_index_manager.build_from_documents(
+                    chunks,
+                    progress_callback=graph_progress,
+                )
+                logger.info(f"Graph index built with {triplet_count} triplets")
+
             # Mark as completed
             self._state.status = IndexingStatusEnum.COMPLETED
             self._state.completed_at = datetime.now(timezone.utc)
@@ -424,6 +447,9 @@ class IndexingService:
         total_code_chunks = self._total_code_chunks
         supported_languages = sorted(self._supported_languages)
 
+        # Get graph index status (Feature 113)
+        graph_status = self.graph_index_manager.get_status()
+
         return {
             "status": self._state.status.value,
             "is_indexing": self._state.is_indexing,
@@ -446,6 +472,14 @@ class IndexingService:
             ),
             "error": self._state.error,
             "indexed_folders": sorted(self._indexed_folders),
+            # Graph index status (Feature 113)
+            "graph_index": {
+                "enabled": graph_status.enabled,
+                "initialized": graph_status.initialized,
+                "entity_count": graph_status.entity_count,
+                "relationship_count": graph_status.relationship_count,
+                "store_type": graph_status.store_type,
+            },
         }
 
     async def reset(self) -> None:
@@ -453,6 +487,8 @@ class IndexingService:
         async with self._lock:
             await self.vector_store.reset()
             self.bm25_manager.reset()
+            # Clear graph index (Feature 113)
+            self.graph_index_manager.clear()
             self._state = IndexingState(
                 current_job_id="",
                 folder_path="",
@@ -461,6 +497,9 @@ class IndexingService:
                 error=None,
             )
             self._indexed_folders.clear()
+            self._total_doc_chunks = 0
+            self._total_code_chunks = 0
+            self._supported_languages.clear()
             logger.info("Indexing service reset")