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

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/runtime.py

@@ -1,4 +1,4 @@
-"""Runtime state management for doc-serve instances."""
+"""Runtime state management for Agent Brain instances."""
 
 import json
 import logging
@@ -15,7 +15,7 @@ logger = logging.getLogger(__name__)
 
 
 class RuntimeState(BaseModel):
-    """Runtime state for a doc-serve instance."""
+    """Runtime state for an Agent Brain instance."""
 
     schema_version: str = "1.0"
     mode: str = "project"  # "project" or "shared"

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

agent_brain_server/services/query_service.py

@@ -7,8 +7,13 @@ from typing import Any, Optional
 from llama_index.core.retrievers import BaseRetriever
 from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
 
+from agent_brain_server.config import settings
 from agent_brain_server.indexing import EmbeddingGenerator, get_embedding_generator
 from agent_brain_server.indexing.bm25_index import BM25IndexManager, get_bm25_manager
+from agent_brain_server.indexing.graph_index import (
+    GraphIndexManager,
+    get_graph_index_manager,
+)
 from agent_brain_server.models import (
     QueryMode,
     QueryRequest,
@@ -69,6 +74,7 @@ class QueryService:
         vector_store: Optional[VectorStoreManager] = None,
         embedding_generator: Optional[EmbeddingGenerator] = None,
         bm25_manager: Optional[BM25IndexManager] = None,
+        graph_index_manager: Optional[GraphIndexManager] = None,
     ):
         """
         Initialize the query service.
@@ -77,10 +83,12 @@ class QueryService:
             vector_store: Vector store manager 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.embedding_generator = embedding_generator or get_embedding_generator()
         self.bm25_manager = bm25_manager or get_bm25_manager()
+        self.graph_index_manager = graph_index_manager or get_graph_index_manager()
 
     def is_ready(self) -> bool:
         """
@@ -115,6 +123,10 @@ class QueryService:
             results = await self._execute_bm25_query(request)
         elif request.mode == QueryMode.VECTOR:
             results = await self._execute_vector_query(request)
+        elif request.mode == QueryMode.GRAPH:
+            results = await self._execute_graph_query(request)
+        elif request.mode == QueryMode.MULTI:
+            results = await self._execute_multi_query(request)
         else:  # HYBRID
             results = await self._execute_hybrid_query(request)
 
@@ -318,6 +330,197 @@ class QueryService:
 
         return fused_nodes
 
+    async def _execute_graph_query(
+        self,
+        request: QueryRequest,
+        traversal_depth: int = 2,
+    ) -> list[QueryResult]:
+        """Execute graph-only query using entity relationships.
+
+        Uses the knowledge graph to find documents related to
+        entities mentioned in the query.
+
+        Args:
+            request: Query request.
+            traversal_depth: How many hops to traverse in graph.
+
+        Returns:
+            List of QueryResult from graph retrieval.
+
+        Raises:
+            ValueError: If GraphRAG is not enabled.
+        """
+        if not settings.ENABLE_GRAPH_INDEX:
+            raise ValueError(
+                "GraphRAG not enabled. Set ENABLE_GRAPH_INDEX=true in environment."
+            )
+
+        # Query the graph for related entities
+        graph_results = self.graph_index_manager.query(
+            query_text=request.query,
+            top_k=request.top_k,
+            traversal_depth=traversal_depth,
+        )
+
+        if not graph_results:
+            logger.debug("No graph results found, falling back to vector search")
+            return await self._execute_vector_query(request)
+
+        # Convert graph results to QueryResults
+        results: list[QueryResult] = []
+        chunk_ids = [
+            r.get("source_chunk_id") for r in graph_results if r.get("source_chunk_id")
+        ]
+
+        if not chunk_ids:
+            # No source chunks in graph, fall back to vector search
+            return await self._execute_vector_query(request)
+
+        # Look up the actual documents from vector store
+        for graph_result in graph_results:
+            chunk_id = graph_result.get("source_chunk_id")
+            if not chunk_id:
+                continue
+
+            # Get document from vector store by ID
+            try:
+                doc = await self.vector_store.get_by_id(chunk_id)
+                if doc:
+                    result = QueryResult(
+                        text=doc.get("text", ""),
+                        source=doc.get("metadata", {}).get(
+                            "source",
+                            doc.get("metadata", {}).get("file_path", "unknown"),
+                        ),
+                        score=graph_result.get("graph_score", 0.5),
+                        graph_score=graph_result.get("graph_score", 0.5),
+                        chunk_id=chunk_id,
+                        source_type=doc.get("metadata", {}).get("source_type", "doc"),
+                        language=doc.get("metadata", {}).get("language"),
+                        related_entities=[
+                            graph_result.get("subject", ""),
+                            graph_result.get("object", ""),
+                        ],
+                        relationship_path=[graph_result.get("relationship_path", "")],
+                        metadata={
+                            k: v
+                            for k, v in doc.get("metadata", {}).items()
+                            if k
+                            not in ("source", "file_path", "source_type", "language")
+                        },
+                    )
+                    results.append(result)
+            except Exception as e:
+                logger.debug(f"Failed to retrieve chunk {chunk_id}: {e}")
+                continue
+
+        # If no results from graph, fall back to vector search
+        if not results:
+            logger.debug("No documents found from graph, falling back to vector search")
+            return await self._execute_vector_query(request)
+
+        return results[: request.top_k]
+
+    async def _execute_multi_query(self, request: QueryRequest) -> list[QueryResult]:
+        """Execute multi-retrieval query combining vector, BM25, and graph.
+
+        Uses Reciprocal Rank Fusion (RRF) to combine results from
+        all three retrieval methods.
+
+        Args:
+            request: Query request.
+
+        Returns:
+            List of QueryResult with combined scores.
+        """
+        # Get results from each retriever
+        vector_results = await self._execute_vector_query(request)
+        bm25_results = await self._execute_bm25_query(request)
+
+        # Get graph results if enabled
+        graph_results: list[QueryResult] = []
+        if settings.ENABLE_GRAPH_INDEX:
+            try:
+                graph_results = await self._execute_graph_query(request)
+            except ValueError:
+                pass  # Graph not enabled, skip
+
+        # Apply Reciprocal Rank Fusion
+        rrf_k = settings.GRAPH_RRF_K  # Typical value is 60
+        combined_scores: dict[str, dict[str, Any]] = {}
+
+        # Process vector results
+        for rank, result in enumerate(vector_results):
+            chunk_id = result.chunk_id
+            rrf_score = 1.0 / (rrf_k + rank + 1)
+            if chunk_id not in combined_scores:
+                combined_scores[chunk_id] = {
+                    "result": result,
+                    "rrf_score": 0.0,
+                    "vector_rank": None,
+                    "bm25_rank": None,
+                    "graph_rank": None,
+                }
+            combined_scores[chunk_id]["rrf_score"] += rrf_score
+            combined_scores[chunk_id]["vector_rank"] = rank + 1
+
+        # Process BM25 results
+        for rank, result in enumerate(bm25_results):
+            chunk_id = result.chunk_id
+            rrf_score = 1.0 / (rrf_k + rank + 1)
+            if chunk_id not in combined_scores:
+                combined_scores[chunk_id] = {
+                    "result": result,
+                    "rrf_score": 0.0,
+                    "vector_rank": None,
+                    "bm25_rank": None,
+                    "graph_rank": None,
+                }
+            combined_scores[chunk_id]["rrf_score"] += rrf_score
+            combined_scores[chunk_id]["bm25_rank"] = rank + 1
+
+        # Process graph results
+        for rank, result in enumerate(graph_results):
+            chunk_id = result.chunk_id
+            rrf_score = 1.0 / (rrf_k + rank + 1)
+            if chunk_id not in combined_scores:
+                combined_scores[chunk_id] = {
+                    "result": result,
+                    "rrf_score": 0.0,
+                    "vector_rank": None,
+                    "bm25_rank": None,
+                    "graph_rank": None,
+                }
+            combined_scores[chunk_id]["rrf_score"] += rrf_score
+            combined_scores[chunk_id]["graph_rank"] = rank + 1
+            # Preserve graph-specific fields
+            if result.related_entities:
+                combined_scores[chunk_id][
+                    "result"
+                ].related_entities = result.related_entities
+            if result.relationship_path:
+                combined_scores[chunk_id][
+                    "result"
+                ].relationship_path = result.relationship_path
+            if result.graph_score:
+                combined_scores[chunk_id]["result"].graph_score = result.graph_score
+
+        # Sort by RRF score and take top_k
+        sorted_results = sorted(
+            combined_scores.values(),
+            key=lambda x: x["rrf_score"],
+            reverse=True,
+        )
+
+        # Update scores and return
+        final_results: list[QueryResult] = []
+        for data in sorted_results[: request.top_k]:
+            result = data["result"]
+            result.score = data["rrf_score"]
+            final_results.append(result)
+
+        return final_results
+
     async def get_document_count(self) -> int:
         """
         Get the total number of indexed documents.

agent_brain_server/storage/__init__.py

@@ -1,5 +1,21 @@
-"""Storage layer for vector database operations."""
+"""Storage layer for vector database and graph operations."""
 
+from .graph_store import (
+    GraphStoreManager,
+    get_graph_store_manager,
+    initialize_graph_store,
+    reset_graph_store_manager,
+)
 from .vector_store import VectorStoreManager, get_vector_store, initialize_vector_store
 
-__all__ = ["VectorStoreManager", "get_vector_store", "initialize_vector_store"]
+__all__ = [
+    # Vector store
+    "VectorStoreManager",
+    "get_vector_store",
+    "initialize_vector_store",
+    # Graph store (Feature 113)
+    "GraphStoreManager",
+    "get_graph_store_manager",
+    "initialize_graph_store",
+    "reset_graph_store_manager",
+]