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

agent_brain_server/config/provider_config.py

@@ -0,0 +1,352 @@
+"""Provider configuration models and YAML loader.
+
+This module provides Pydantic models for embedding and summarization
+provider configuration, and functions to load configuration from YAML files.
+"""
+
+import logging
+import os
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Optional
+
+import yaml
+from pydantic import BaseModel, Field, field_validator
+
+from agent_brain_server.providers.base import (
+    EmbeddingProviderType,
+    SummarizationProviderType,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingConfig(BaseModel):
+    """Configuration for embedding provider."""
+
+    provider: EmbeddingProviderType = Field(
+        default=EmbeddingProviderType.OPENAI,
+        description="Embedding provider to use",
+    )
+    model: str = Field(
+        default="text-embedding-3-large",
+        description="Model name for embeddings",
+    )
+    api_key: Optional[str] = Field(
+        default=None,
+        description="API key (alternative to api_key_env for local config files)",
+    )
+    api_key_env: Optional[str] = Field(
+        default="OPENAI_API_KEY",
+        description="Environment variable name containing API key",
+    )
+    base_url: Optional[str] = Field(
+        default=None,
+        description="Custom base URL (for Ollama or compatible APIs)",
+    )
+    params: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Provider-specific parameters",
+    )
+
+    model_config = {"use_enum_values": True}
+
+    @field_validator("provider", mode="before")
+    @classmethod
+    def validate_provider(cls, v: Any) -> EmbeddingProviderType:
+        """Convert string to enum if needed."""
+        if isinstance(v, str):
+            return EmbeddingProviderType(v.lower())
+        if isinstance(v, EmbeddingProviderType):
+            return v
+        return EmbeddingProviderType(v)
+
+    def get_api_key(self) -> Optional[str]:
+        """Resolve API key from config or environment variable.
+
+        Resolution order:
+        1. api_key field in config (direct value)
+        2. Environment variable specified by api_key_env
+
+        Returns:
+            API key value or None if not found/not needed
+        """
+        if self.provider == EmbeddingProviderType.OLLAMA:
+            return None  # Ollama doesn't need API key
+        # Check direct api_key first
+        if self.api_key:
+            return self.api_key
+        # Fall back to environment variable
+        if self.api_key_env:
+            return os.getenv(self.api_key_env)
+        return None
+
+    def get_base_url(self) -> Optional[str]:
+        """Get base URL with defaults for specific providers.
+
+        Returns:
+            Base URL for the provider
+        """
+        if self.base_url:
+            return self.base_url
+        if self.provider == EmbeddingProviderType.OLLAMA:
+            return "http://localhost:11434/v1"
+        return None
+
+
+class SummarizationConfig(BaseModel):
+    """Configuration for summarization provider."""
+
+    provider: SummarizationProviderType = Field(
+        default=SummarizationProviderType.ANTHROPIC,
+        description="Summarization provider to use",
+    )
+    model: str = Field(
+        default="claude-haiku-4-5-20251001",
+        description="Model name for summarization",
+    )
+    api_key: Optional[str] = Field(
+        default=None,
+        description="API key (alternative to api_key_env for local config files)",
+    )
+    api_key_env: Optional[str] = Field(
+        default="ANTHROPIC_API_KEY",
+        description="Environment variable name containing API key",
+    )
+    base_url: Optional[str] = Field(
+        default=None,
+        description="Custom base URL (for Grok or Ollama)",
+    )
+    params: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Provider-specific parameters (max_tokens, temperature)",
+    )
+
+    model_config = {"use_enum_values": True}
+
+    @field_validator("provider", mode="before")
+    @classmethod
+    def validate_provider(cls, v: Any) -> SummarizationProviderType:
+        """Convert string to enum if needed."""
+        if isinstance(v, str):
+            return SummarizationProviderType(v.lower())
+        if isinstance(v, SummarizationProviderType):
+            return v
+        return SummarizationProviderType(v)
+
+    def get_api_key(self) -> Optional[str]:
+        """Resolve API key from config or environment variable.
+
+        Resolution order:
+        1. api_key field in config (direct value)
+        2. Environment variable specified by api_key_env
+
+        Returns:
+            API key value or None if not found/not needed
+        """
+        if self.provider == SummarizationProviderType.OLLAMA:
+            return None  # Ollama doesn't need API key
+        # Check direct api_key first
+        if self.api_key:
+            return self.api_key
+        # Fall back to environment variable
+        if self.api_key_env:
+            return os.getenv(self.api_key_env)
+        return None
+
+    def get_base_url(self) -> Optional[str]:
+        """Get base URL with defaults for specific providers.
+
+        Returns:
+            Base URL for the provider
+        """
+        if self.base_url:
+            return self.base_url
+        if self.provider == SummarizationProviderType.OLLAMA:
+            return "http://localhost:11434/v1"
+        if self.provider == SummarizationProviderType.GROK:
+            return "https://api.x.ai/v1"
+        return None
+
+
+class ProviderSettings(BaseModel):
+    """Top-level provider configuration."""
+
+    embedding: EmbeddingConfig = Field(
+        default_factory=EmbeddingConfig,
+        description="Embedding provider configuration",
+    )
+    summarization: SummarizationConfig = Field(
+        default_factory=SummarizationConfig,
+        description="Summarization provider configuration",
+    )
+
+
+def _find_config_file() -> Optional[Path]:
+    """Find the configuration file in standard locations.
+
+    Search order:
+    1. AGENT_BRAIN_CONFIG environment variable
+    2. State directory config.yaml (if AGENT_BRAIN_STATE_DIR or DOC_SERVE_STATE_DIR set)
+    3. Current directory config.yaml
+    4. Walk up from CWD looking for .claude/agent-brain/config.yaml
+    5. User home ~/.agent-brain/config.yaml
+    6. XDG config ~/.config/agent-brain/config.yaml
+
+    Returns:
+        Path to config file or None if not found
+    """
+    # 1. Environment variable override
+    env_config = os.getenv("AGENT_BRAIN_CONFIG")
+    if env_config:
+        path = Path(env_config)
+        if path.exists():
+            logger.debug(f"Found config via AGENT_BRAIN_CONFIG: {path}")
+            return path
+        logger.warning(f"AGENT_BRAIN_CONFIG points to non-existent file: {env_config}")
+
+    # 2. State directory (check both new and legacy env vars)
+    state_dir = os.getenv("AGENT_BRAIN_STATE_DIR") or os.getenv("DOC_SERVE_STATE_DIR")
+    if state_dir:
+        state_config = Path(state_dir) / "config.yaml"
+        if state_config.exists():
+            logger.debug(f"Found config in state directory: {state_config}")
+            return state_config
+
+    # 3. Current directory
+    cwd_config = Path.cwd() / "config.yaml"
+    if cwd_config.exists():
+        logger.debug(f"Found config in current directory: {cwd_config}")
+        return cwd_config
+
+    # 4. Walk up from CWD looking for .claude/agent-brain/config.yaml
+    current = Path.cwd()
+    root = Path(current.anchor)
+    while current != root:
+        claude_config = current / ".claude" / "agent-brain" / "config.yaml"
+        if claude_config.exists():
+            logger.debug(f"Found config walking up from CWD: {claude_config}")
+            return claude_config
+        current = current.parent
+
+    # 5. User home directory ~/.agent-brain/config.yaml
+    home_config = Path.home() / ".agent-brain" / "config.yaml"
+    if home_config.exists():
+        logger.debug(f"Found config in home directory: {home_config}")
+        return home_config
+
+    # 6. XDG config directory ~/.config/agent-brain/config.yaml
+    xdg_config = Path.home() / ".config" / "agent-brain" / "config.yaml"
+    if xdg_config.exists():
+        logger.debug(f"Found config in XDG config directory: {xdg_config}")
+        return xdg_config
+
+    return None
+
+
+def _load_yaml_config(path: Path) -> dict[str, Any]:
+    """Load YAML configuration from file.
+
+    Args:
+        path: Path to YAML config file
+
+    Returns:
+        Configuration dictionary
+
+    Raises:
+        ConfigurationError: If YAML parsing fails
+    """
+    from agent_brain_server.providers.exceptions import ConfigurationError
+
+    try:
+        with open(path) as f:
+            config = yaml.safe_load(f)
+            return config if config else {}
+    except yaml.YAMLError as e:
+        raise ConfigurationError(
+            f"Failed to parse config file {path}: {e}",
+            "config",
+        ) from e
+    except OSError as e:
+        raise ConfigurationError(
+            f"Failed to read config file {path}: {e}",
+            "config",
+        ) from e
+
+
+@lru_cache
+def load_provider_settings() -> ProviderSettings:
+    """Load provider settings from YAML config or defaults.
+
+    This function:
+    1. Searches for config.yaml in standard locations
+    2. Parses YAML and validates against Pydantic models
+    3. Falls back to defaults (OpenAI embeddings + Anthropic summarization)
+
+    Returns:
+        Validated ProviderSettings instance
+    """
+    config_path = _find_config_file()
+
+    if config_path:
+        logger.info(f"Loading provider config from {config_path}")
+        raw_config = _load_yaml_config(config_path)
+        settings = ProviderSettings(**raw_config)
+    else:
+        logger.info("No config file found, using default providers")
+        settings = ProviderSettings()
+
+    # Log active configuration
+    logger.info(
+        f"Active embedding provider: {settings.embedding.provider} "
+        f"(model: {settings.embedding.model})"
+    )
+    logger.info(
+        f"Active summarization provider: {settings.summarization.provider} "
+        f"(model: {settings.summarization.model})"
+    )
+
+    return settings
+
+
+def clear_settings_cache() -> None:
+    """Clear the cached provider settings (for testing)."""
+    load_provider_settings.cache_clear()
+
+
+def validate_provider_config(settings: ProviderSettings) -> list[str]:
+    """Validate provider configuration and return list of errors.
+
+    Checks:
+    - API keys are available for providers that need them
+    - Models are known for the selected provider
+
+    Args:
+        settings: Provider settings to validate
+
+    Returns:
+        List of validation error messages (empty if valid)
+    """
+    errors: list[str] = []
+
+    # Validate embedding provider
+    if settings.embedding.provider != EmbeddingProviderType.OLLAMA:
+        api_key = settings.embedding.get_api_key()
+        if not api_key:
+            env_var = settings.embedding.api_key_env or "OPENAI_API_KEY"
+            errors.append(
+                f"Missing API key for {settings.embedding.provider} embeddings. "
+                f"Set {env_var} environment variable."
+            )
+
+    # Validate summarization provider
+    if settings.summarization.provider != SummarizationProviderType.OLLAMA:
+        api_key = settings.summarization.get_api_key()
+        if not api_key:
+            env_var = settings.summarization.api_key_env or "ANTHROPIC_API_KEY"
+            errors.append(
+                f"Missing API key for {settings.summarization.provider} summarization. "
+                f"Set {env_var} environment variable."
+            )
+
+    return errors

agent_brain_server/config/settings.py

@@ -26,12 +26,12 @@ class Settings(BaseSettings):
 
     # Anthropic Configuration
     ANTHROPIC_API_KEY: str = ""
-    CLAUDE_MODEL: str = "claude-3-5-haiku-20241022"  # Claude 3.5 Haiku (latest)
+    CLAUDE_MODEL: str = "claude-haiku-4-5-20251001"  # Claude 4.5 Haiku (latest)
 
     # Chroma Configuration
     CHROMA_PERSIST_DIR: str = "./chroma_db"
     BM25_INDEX_PATH: str = "./bm25_index"
-    COLLECTION_NAME: str = "doc_serve_collection"
+    COLLECTION_NAME: str = "agent_brain_collection"
 
     # Chunking Configuration
     DEFAULT_CHUNK_SIZE: int = 512
@@ -48,14 +48,31 @@ class Settings(BaseSettings):
     EMBEDDING_BATCH_SIZE: int = 100
 
     # Multi-instance Configuration
-    DOC_SERVE_STATE_DIR: Optional[str] = None  # Override state directory
-    DOC_SERVE_MODE: str = "project"  # "project" or "shared"
+    AGENT_BRAIN_STATE_DIR: Optional[str] = None  # Override state directory
+    AGENT_BRAIN_MODE: str = "project"  # "project" or "shared"
+
+    # GraphRAG Configuration (Feature 113)
+    ENABLE_GRAPH_INDEX: bool = False  # Master switch for graph indexing
+    GRAPH_STORE_TYPE: str = "simple"  # "simple" (in-memory) or "kuzu" (persistent)
+    GRAPH_INDEX_PATH: str = "./graph_index"  # Path for graph persistence
+    GRAPH_EXTRACTION_MODEL: str = "claude-haiku-4-5"  # Model for entity extraction
+    GRAPH_MAX_TRIPLETS_PER_CHUNK: int = 10  # Max triplets per document chunk
+    GRAPH_USE_CODE_METADATA: bool = True  # Use AST metadata for code entities
+    GRAPH_USE_LLM_EXTRACTION: bool = True  # Use LLM for additional extraction
+    GRAPH_TRAVERSAL_DEPTH: int = 2  # Depth for graph traversal in queries
+    GRAPH_RRF_K: int = 60  # Reciprocal Rank Fusion constant for multi-retrieval
+
+    # Job Queue Configuration (Feature 115)
+    AGENT_BRAIN_MAX_QUEUE: int = 100  # Max pending jobs in queue
+    AGENT_BRAIN_JOB_TIMEOUT: int = 7200  # Job timeout in seconds (2 hours)
+    AGENT_BRAIN_MAX_RETRIES: int = 3  # Max retries for failed jobs
+    AGENT_BRAIN_CHECKPOINT_INTERVAL: int = 50  # Progress checkpoint every N files
 
     model_config = SettingsConfigDict(
         env_file=[
             ".env",  # Current directory
             Path(__file__).parent.parent.parent / ".env",  # Project root
-            Path(__file__).parent.parent / ".env",  # doc-serve-server directory
+            Path(__file__).parent.parent / ".env",  # agent-brain-server directory
         ],
         env_file_encoding="utf-8",
         case_sensitive=True,

agent_brain_server/indexing/__init__.py

@@ -7,6 +7,18 @@ from agent_brain_server.indexing.embedding import (
     EmbeddingGenerator,
     get_embedding_generator,
 )
+from agent_brain_server.indexing.graph_extractors import (
+    CodeMetadataExtractor,
+    LLMEntityExtractor,
+    get_code_extractor,
+    get_llm_extractor,
+    reset_extractors,
+)
+from agent_brain_server.indexing.graph_index import (
+    GraphIndexManager,
+    get_graph_index_manager,
+    reset_graph_index_manager,
+)
 
 __all__ = [
     "DocumentLoader",
@@ -16,4 +28,13 @@ __all__ = [
     "get_embedding_generator",
     "BM25IndexManager",
     "get_bm25_manager",
+    # Graph indexing (Feature 113)
+    "LLMEntityExtractor",
+    "CodeMetadataExtractor",
+    "get_llm_extractor",
+    "get_code_extractor",
+    "reset_extractors",
+    "GraphIndexManager",
+    "get_graph_index_manager",
+    "reset_graph_index_manager",
 ]

agent_brain_server/indexing/bm25_index.py

@@ -89,10 +89,23 @@ class BM25IndexManager:
         if not self._retriever:
             raise RuntimeError("BM25 index not initialized")
 
-        # BM25Retriever similarity_top_k is usually set during initialization.
-        self._retriever.similarity_top_k = top_k
+        # Cap top_k to corpus size to avoid bm25s "k larger than available scores" error
+        corpus_size = len(self._retriever.corpus) if self._retriever.corpus else 0
+        if corpus_size > 0:
+            effective_top_k = min(top_k, corpus_size)
+        else:
+            effective_top_k = top_k
+
+        self._retriever.similarity_top_k = effective_top_k
         return self._retriever
 
+    @property
+    def corpus_size(self) -> int:
+        """Get the number of documents in the BM25 index."""
+        if not self._retriever or not self._retriever.corpus:
+            return 0
+        return len(self._retriever.corpus)
+
     async def search_with_filters(
         self,
         query: str,

agent_brain_server/indexing/document_loader.py

@@ -1,5 +1,6 @@
 """Document loading from various file formats using LlamaIndex."""
 
+import asyncio
 import logging
 import re
 from dataclasses import dataclass, field
@@ -272,9 +273,30 @@ class DocumentLoader:
 
     SUPPORTED_EXTENSIONS: set[str] = DOCUMENT_EXTENSIONS | CODE_EXTENSIONS
 
+    # Default directories to exclude from indexing
+    DEFAULT_EXCLUDE_PATTERNS: list[str] = [
+        "**/node_modules/**",
+        "**/__pycache__/**",
+        "**/.venv/**",
+        "**/venv/**",
+        "**/.git/**",
+        "**/dist/**",
+        "**/build/**",
+        "**/target/**",
+        "**/.next/**",
+        "**/.nuxt/**",
+        "**/coverage/**",
+        "**/.pytest_cache/**",
+        "**/.mypy_cache/**",
+        "**/.tox/**",
+        "**/egg-info/**",
+        "**/*.egg-info/**",
+    ]
+
     def __init__(
         self,
         supported_extensions: Optional[set[str]] = None,
+        exclude_patterns: Optional[list[str]] = None,
     ):
         """
         Initialize the document loader.
@@ -282,8 +304,15 @@ class DocumentLoader:
         Args:
             supported_extensions: Set of file extensions to load.
                                   Defaults to SUPPORTED_EXTENSIONS.
+            exclude_patterns: List of glob patterns to exclude.
+                              Defaults to DEFAULT_EXCLUDE_PATTERNS.
         """
         self.extensions = supported_extensions or self.SUPPORTED_EXTENSIONS
+        self.exclude_patterns = (
+            exclude_patterns
+            if exclude_patterns is not None
+            else self.DEFAULT_EXCLUDE_PATTERNS
+        )
 
     async def load_from_folder(
         self,
@@ -313,16 +342,24 @@ class DocumentLoader:
             raise ValueError(f"Path is not a directory: {folder_path}")
 
         logger.info(f"Loading documents from: {folder_path} (recursive={recursive})")
+        if self.exclude_patterns:
+            logger.info(
+                f"Excluding patterns: {self.exclude_patterns[:3]}... "
+                f"({len(self.exclude_patterns)} total)"
+            )
 
         # Use LlamaIndex's SimpleDirectoryReader
+        # Run in thread pool to avoid blocking the event loop
         try:
             reader = SimpleDirectoryReader(
                 input_dir=str(path),
                 recursive=recursive,
                 required_exts=list(self.extensions),
+                exclude=self.exclude_patterns,
                 filename_as_id=True,
             )
-            llama_documents: list[Document] = reader.load_data()
+            # reader.load_data() is blocking I/O - run in thread pool
+            llama_documents: list[Document] = await asyncio.to_thread(reader.load_data)
         except Exception as e:
             logger.error(f"Failed to load documents: {e}")
             raise
@@ -398,7 +435,8 @@ class DocumentLoader:
             input_files=[str(path)],
             filename_as_id=True,
         )
-        docs = reader.load_data()
+        # Run in thread pool to avoid blocking the event loop
+        docs = await asyncio.to_thread(reader.load_data)
 
         if not docs:
             raise ValueError(f"No content loaded from file: {file_path}")
@@ -456,8 +494,11 @@ class DocumentLoader:
             # Use only document extensions
             effective_extensions = self.DOCUMENT_EXTENSIONS
 
-        # Create a temporary loader with the effective extensions
-        temp_loader = DocumentLoader(supported_extensions=effective_extensions)
+        # Create a temporary loader with the effective extensions and exclude patterns
+        temp_loader = DocumentLoader(
+            supported_extensions=effective_extensions,
+            exclude_patterns=self.exclude_patterns,
+        )
 
         # Load files using the configured extensions
         loaded_docs = await temp_loader.load_from_folder(folder_path, recursive)