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

agent_brain_server/models/graph.py

@@ -0,0 +1,253 @@
+"""Models for GraphRAG feature (Feature 113).
+
+Defines Pydantic models for graph entities, relationships, and status.
+All models are configured with frozen=True for immutability.
+"""
+
+from datetime import datetime
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class GraphTriple(BaseModel):
+    """Represents a subject-predicate-object triple in the knowledge graph.
+
+    Triples are the fundamental unit of knowledge representation in GraphRAG.
+    They capture relationships between entities extracted from documents.
+
+    Attributes:
+        subject: The subject entity (e.g., "FastAPI").
+        subject_type: Optional type classification (e.g., "Framework").
+        predicate: The relationship type (e.g., "uses").
+        object: The object entity (e.g., "Pydantic").
+        object_type: Optional type classification (e.g., "Library").
+        source_chunk_id: Optional ID of the source document chunk.
+    """
+
+    model_config = ConfigDict(
+        frozen=True,
+        json_schema_extra={
+            "examples": [
+                {
+                    "subject": "FastAPI",
+                    "subject_type": "Framework",
+                    "predicate": "uses",
+                    "object": "Pydantic",
+                    "object_type": "Library",
+                    "source_chunk_id": "chunk_abc123",
+                },
+                {
+                    "subject": "UserController",
+                    "subject_type": "Class",
+                    "predicate": "calls",
+                    "object": "authenticate_user",
+                    "object_type": "Function",
+                    "source_chunk_id": "chunk_def456",
+                },
+            ]
+        },
+    )
+
+    subject: str = Field(
+        ...,
+        min_length=1,
+        description="Subject entity in the triple",
+    )
+    subject_type: Optional[str] = Field(
+        default=None,
+        description="Type classification for subject entity",
+    )
+    predicate: str = Field(
+        ...,
+        min_length=1,
+        description="Relationship type connecting subject to object",
+    )
+    object: str = Field(
+        ...,
+        min_length=1,
+        description="Object entity in the triple",
+    )
+    object_type: Optional[str] = Field(
+        default=None,
+        description="Type classification for object entity",
+    )
+    source_chunk_id: Optional[str] = Field(
+        default=None,
+        description="ID of the source document chunk",
+    )
+
+
+class GraphEntity(BaseModel):
+    """Represents an entity node in the knowledge graph.
+
+    Entities are the nodes in the graph, representing concepts,
+    code elements, or other named items extracted from documents.
+
+    Attributes:
+        name: Unique name/identifier of the entity.
+        entity_type: Classification type (e.g., "Class", "Function", "Concept").
+        description: Optional description of the entity.
+        source_chunk_ids: List of source chunk IDs where entity appears.
+        properties: Additional metadata properties.
+    """
+
+    model_config = ConfigDict(
+        frozen=True,
+        json_schema_extra={
+            "examples": [
+                {
+                    "name": "VectorStoreManager",
+                    "entity_type": "Class",
+                    "description": "Manages Chroma vector store operations",
+                    "source_chunk_ids": ["chunk_001", "chunk_002"],
+                    "properties": {"module": "storage.vector_store"},
+                },
+            ]
+        },
+    )
+
+    name: str = Field(
+        ...,
+        min_length=1,
+        description="Unique name/identifier of the entity",
+    )
+    entity_type: Optional[str] = Field(
+        default=None,
+        description="Classification type for the entity",
+    )
+    description: Optional[str] = Field(
+        default=None,
+        description="Description of the entity",
+    )
+    source_chunk_ids: list[str] = Field(
+        default_factory=list,
+        description="List of source chunk IDs where entity appears",
+    )
+    properties: dict[str, str] = Field(
+        default_factory=dict,
+        description="Additional metadata properties",
+    )
+
+
+class GraphIndexStatus(BaseModel):
+    """Status of the graph index.
+
+    Provides information about the graph index state,
+    including whether it's enabled, initialized, and statistics.
+
+    Attributes:
+        enabled: Whether graph indexing is enabled.
+        initialized: Whether the graph store is initialized.
+        entity_count: Number of entities in the graph.
+        relationship_count: Number of relationships in the graph.
+        last_updated: Timestamp of last graph update.
+        store_type: Type of graph store backend.
+    """
+
+    model_config = ConfigDict(
+        frozen=True,
+        json_schema_extra={
+            "examples": [
+                {
+                    "enabled": True,
+                    "initialized": True,
+                    "entity_count": 150,
+                    "relationship_count": 320,
+                    "last_updated": "2024-12-15T10:30:00Z",
+                    "store_type": "simple",
+                },
+                {
+                    "enabled": False,
+                    "initialized": False,
+                    "entity_count": 0,
+                    "relationship_count": 0,
+                    "last_updated": None,
+                    "store_type": "simple",
+                },
+            ]
+        },
+    )
+
+    enabled: bool = Field(
+        default=False,
+        description="Whether graph indexing is enabled",
+    )
+    initialized: bool = Field(
+        default=False,
+        description="Whether the graph store is initialized",
+    )
+    entity_count: int = Field(
+        default=0,
+        ge=0,
+        description="Number of entities in the graph",
+    )
+    relationship_count: int = Field(
+        default=0,
+        ge=0,
+        description="Number of relationships in the graph",
+    )
+    last_updated: Optional[datetime] = Field(
+        default=None,
+        description="Timestamp of last graph update",
+    )
+    store_type: str = Field(
+        default="simple",
+        description="Type of graph store backend (simple or kuzu)",
+    )
+
+
+class GraphQueryContext(BaseModel):
+    """Context information from graph-based retrieval.
+
+    Contains additional context extracted from the knowledge graph
+    during query processing.
+
+    Attributes:
+        related_entities: List of related entity names.
+        relationship_paths: List of relationship paths as strings.
+        subgraph_triplets: Relevant triplets from the graph.
+        graph_score: Score from graph-based retrieval.
+    """
+
+    model_config = ConfigDict(
+        frozen=True,
+        json_schema_extra={
+            "examples": [
+                {
+                    "related_entities": ["FastAPI", "Pydantic", "Uvicorn"],
+                    "relationship_paths": [
+                        "FastAPI -> uses -> Pydantic",
+                        "FastAPI -> runs_on -> Uvicorn",
+                    ],
+                    "subgraph_triplets": [
+                        {
+                            "subject": "FastAPI",
+                            "predicate": "uses",
+                            "object": "Pydantic",
+                        },
+                    ],
+                    "graph_score": 0.85,
+                },
+            ]
+        },
+    )
+
+    related_entities: list[str] = Field(
+        default_factory=list,
+        description="List of related entity names",
+    )
+    relationship_paths: list[str] = Field(
+        default_factory=list,
+        description="Relationship paths as formatted strings",
+    )
+    subgraph_triplets: list[GraphTriple] = Field(
+        default_factory=list,
+        description="Relevant triplets from the graph",
+    )
+    graph_score: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="Score from graph-based retrieval",
+    )

agent_brain_server/models/health.py

@@ -1,7 +1,7 @@
 """Health status models."""
 
 from datetime import datetime, timezone
-from typing import Literal, Optional
+from typing import Any, Literal, Optional
 
 from pydantic import BaseModel, Field
 
@@ -22,7 +22,7 @@ class HealthStatus(BaseModel):
         description="Timestamp of the health check",
     )
     version: str = Field(
-        default="1.2.0",
+        default="2.0.0",
         description="Server version",
     )
     mode: Optional[str] = Field(
@@ -49,7 +49,7 @@ class HealthStatus(BaseModel):
                     "status": "healthy",
                     "message": "Server is running and ready for queries",
                     "timestamp": "2024-12-15T10:30:00Z",
-                    "version": "1.2.0",
+                    "version": "2.0.0",
                 }
             ]
         }
@@ -105,6 +105,11 @@ class IndexingStatus(BaseModel):
         default_factory=list,
         description="List of folders that have been indexed",
     )
+    # Graph index status (Feature 113)
+    graph_index: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="Graph index status with entity_count, relationship_count, etc.",
+    )
 
     model_config = {
         "json_schema_extra": {
@@ -120,6 +125,13 @@ class IndexingStatus(BaseModel):
                     "last_indexed_at": "2024-12-15T10:30:00Z",
                     "indexed_folders": ["/path/to/docs"],
                     "supported_languages": ["python", "typescript", "java"],
+                    "graph_index": {
+                        "enabled": True,
+                        "initialized": True,
+                        "entity_count": 120,
+                        "relationship_count": 250,
+                        "store_type": "simple",
+                    },
                 }
             ]
         }

agent_brain_server/models/query.py

@@ -14,6 +14,8 @@ class QueryMode(str, Enum):
     VECTOR = "vector"
     BM25 = "bm25"
     HYBRID = "hybrid"
+    GRAPH = "graph"  # Graph-only retrieval (Feature 113)
+    MULTI = "multi"  # Multi-retrieval: vector + BM25 + graph with RRF (Feature 113)
 
 
 class QueryRequest(BaseModel):
@@ -39,7 +41,7 @@ class QueryRequest(BaseModel):
     )
     mode: QueryMode = Field(
         default=QueryMode.HYBRID,
-        description="Retrieval mode (vector, bm25, hybrid)",
+        description="Retrieval mode (vector, bm25, hybrid, graph, multi)",
     )
     alpha: float = Field(
         default=0.5,
@@ -131,6 +133,17 @@ class QueryResult(BaseModel):
         default=None, description="Programming language for code files"
     )
 
+    # GraphRAG fields (Feature 113)
+    graph_score: float | None = Field(
+        default=None, description="Score from graph-based retrieval"
+    )
+    related_entities: list[str] | None = Field(
+        default=None, description="Related entities from knowledge graph"
+    )
+    relationship_path: list[str] | None = Field(
+        default=None, description="Relationship paths in the graph"
+    )
+
     # Additional metadata
     metadata: dict[str, Any] = Field(
         default_factory=dict, description="Additional metadata"

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",
+]