agent-brain-rag 1.1.0__py3-none-any.whl

doc_serve_server/api/routers/index.py

@@ -0,0 +1,208 @@
+"""Indexing endpoints for document processing."""
+
+import os
+from pathlib import Path
+
+from fastapi import APIRouter, HTTPException, Request, status
+
+from doc_serve_server.models import IndexRequest, IndexResponse
+
+router = APIRouter()
+
+
+@router.post(
+    "/",
+    response_model=IndexResponse,
+    status_code=status.HTTP_202_ACCEPTED,
+    summary="Index Documents",
+    description="Start indexing documents from a folder.",
+)
+async def index_documents(
+    request_body: IndexRequest, request: Request
+) -> IndexResponse:
+    """Start indexing documents from the specified folder.
+
+    This endpoint initiates a background indexing job and returns immediately.
+    Use the /health/status endpoint to monitor progress.
+
+    Args:
+        request_body: IndexRequest with folder_path and optional configuration.
+        request: FastAPI request for accessing app state.
+
+    Returns:
+        IndexResponse with job_id and status.
+
+    Raises:
+        400: Invalid folder path
+        409: Indexing already in progress
+    """
+    # Validate folder path
+    folder_path = Path(request_body.folder_path).expanduser().resolve()
+
+    if not folder_path.exists():
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Folder not found: {request_body.folder_path}",
+        )
+
+    if not folder_path.is_dir():
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Path is not a directory: {request_body.folder_path}",
+        )
+
+    if not os.access(folder_path, os.R_OK):
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Cannot read folder: {request_body.folder_path}",
+        )
+
+    # Get indexing service from app state
+    indexing_service = request.app.state.indexing_service
+
+    # Check if already indexing
+    if indexing_service.is_indexing:
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail="Indexing already in progress. Please wait for completion.",
+        )
+
+    # Start indexing
+    try:
+        # Update request with resolved path
+        resolved_request = IndexRequest(
+            folder_path=str(folder_path),
+            chunk_size=request_body.chunk_size,
+            chunk_overlap=request_body.chunk_overlap,
+            recursive=request_body.recursive,
+            include_code=request_body.include_code,
+            supported_languages=request_body.supported_languages,
+            code_chunk_strategy=request_body.code_chunk_strategy,
+            include_patterns=request_body.include_patterns,
+            exclude_patterns=request_body.exclude_patterns,
+            generate_summaries=request_body.generate_summaries,
+        )
+        job_id = await indexing_service.start_indexing(resolved_request)
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to start indexing: {str(e)}",
+        ) from e
+
+    return IndexResponse(
+        job_id=job_id,
+        status="started",
+        message=f"Indexing started for {request_body.folder_path}",
+    )
+
+
+@router.post(
+    "/add",
+    response_model=IndexResponse,
+    status_code=status.HTTP_202_ACCEPTED,
+    summary="Add Documents",
+    description="Add documents from another folder to the existing index.",
+)
+async def add_documents(request_body: IndexRequest, request: Request) -> IndexResponse:
+    """Add documents from a new folder to the existing index.
+
+    This is similar to the index endpoint but adds to the existing
+    vector store instead of replacing it.
+
+    Args:
+        request_body: IndexRequest with folder_path and optional configuration.
+        request: FastAPI request for accessing app state.
+
+    Returns:
+        IndexResponse with job_id and status.
+    """
+    # Same validation as index_documents
+    folder_path = Path(request_body.folder_path).expanduser().resolve()
+
+    if not folder_path.exists():
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Folder not found: {request_body.folder_path}",
+        )
+
+    if not folder_path.is_dir():
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Path is not a directory: {request_body.folder_path}",
+        )
+
+    indexing_service = request.app.state.indexing_service
+
+    if indexing_service.is_indexing:
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail="Indexing already in progress. Please wait for completion.",
+        )
+
+    try:
+        resolved_request = IndexRequest(
+            folder_path=str(folder_path),
+            chunk_size=request_body.chunk_size,
+            chunk_overlap=request_body.chunk_overlap,
+            recursive=request_body.recursive,
+            include_code=request_body.include_code,
+            supported_languages=request_body.supported_languages,
+            code_chunk_strategy=request_body.code_chunk_strategy,
+            include_patterns=request_body.include_patterns,
+            exclude_patterns=request_body.exclude_patterns,
+        )
+        job_id = await indexing_service.start_indexing(resolved_request)
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to add documents: {str(e)}",
+        ) from e
+
+    return IndexResponse(
+        job_id=job_id,
+        status="started",
+        message=f"Adding documents from {request_body.folder_path}",
+    )
+
+
+@router.delete(
+    "/",
+    response_model=IndexResponse,
+    summary="Reset Index",
+    description="Delete all indexed documents and reset the vector store.",
+)
+async def reset_index(request: Request) -> IndexResponse:
+    """Reset the index by deleting all stored documents.
+
+    Warning: This permanently removes all indexed content.
+
+    Args:
+        request: FastAPI request for accessing app state.
+
+    Returns:
+        IndexResponse confirming the reset.
+
+    Raises:
+        409: Indexing in progress
+    """
+    indexing_service = request.app.state.indexing_service
+
+    if indexing_service.is_indexing:
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail="Cannot reset while indexing is in progress.",
+        )
+
+    try:
+        await indexing_service.reset()
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to reset index: {str(e)}",
+        ) from e
+
+    return IndexResponse(
+        job_id="reset",
+        status="completed",
+        message="Index has been reset successfully",
+    )

doc_serve_server/api/routers/query.py

@@ -0,0 +1,96 @@
+"""Query endpoints for semantic search."""
+
+import logging
+
+from fastapi import APIRouter, HTTPException, Request, status
+
+from doc_serve_server.models import QueryRequest, QueryResponse
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+
+@router.post(
+    "/",
+    response_model=QueryResponse,
+    summary="Query Documents",
+    description="Perform semantic, keyword, or hybrid search on indexed documents.",
+)
+async def query_documents(
+    request_body: QueryRequest, request: Request
+) -> QueryResponse:
+    """Execute a search query on indexed documents.
+
+    Args:
+        request_body: QueryRequest containing query parameters.
+        request: FastAPI request for accessing app state.
+
+    Returns:
+        QueryResponse with ranked results and timing.
+
+    Raises:
+        400: Invalid query (empty or too long)
+        503: Index not ready (indexing in progress or not initialized)
+    """
+    from doc_serve_server.services import QueryService
+    from doc_serve_server.services.indexing_service import IndexingService
+
+    query_service: QueryService = request.app.state.query_service
+    indexing_service: IndexingService = request.app.state.indexing_service
+
+    # Validate query
+    query = request_body.query.strip()
+    if not query:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Query cannot be empty",
+        )
+
+    # Check if service is ready
+    if not query_service.is_ready():
+        if indexing_service.is_indexing:
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                detail="Index not ready. Indexing is in progress.",
+            )
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                detail="Index not ready. Please index documents first.",
+            )
+
+    # Execute query
+    try:
+        response = await query_service.execute_query(request_body)
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Query failed: {str(e)}",
+        ) from e
+
+    return response
+
+
+@router.get(
+    "/count",
+    summary="Document Count",
+    description="Get the total number of indexed document chunks.",
+)
+async def get_document_count(request: Request) -> dict[str, int | bool]:
+    """Get the total number of indexed document chunks.
+
+    Args:
+        request: FastAPI request for accessing app state.
+
+    Returns:
+        Dictionary with count of indexed chunks.
+    """
+    query_service = request.app.state.query_service
+
+    count = await query_service.get_document_count()
+
+    return {
+        "total_chunks": count,
+        "ready": query_service.is_ready(),
+    }

doc_serve_server/config/__init__.py

@@ -0,0 +1,5 @@
+"""Configuration module."""
+
+from .settings import settings
+
+__all__ = ["settings"]

doc_serve_server/config/settings.py

@@ -0,0 +1,92 @@
+"""Application configuration using Pydantic settings."""
+
+import json
+import logging
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Optional
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+logger = logging.getLogger(__name__)
+
+
+class Settings(BaseSettings):
+    """Application settings loaded from environment variables."""
+
+    # API Configuration
+    API_HOST: str = "127.0.0.1"
+    API_PORT: int = 8000
+    DEBUG: bool = False
+
+    # OpenAI Configuration
+    OPENAI_API_KEY: str = ""
+    EMBEDDING_MODEL: str = "text-embedding-3-large"
+    EMBEDDING_DIMENSIONS: int = 3072
+
+    # Anthropic Configuration
+    ANTHROPIC_API_KEY: str = ""
+    CLAUDE_MODEL: str = "claude-3-5-haiku-20241022"  # Claude 3.5 Haiku (latest)
+
+    # Chroma Configuration
+    CHROMA_PERSIST_DIR: str = "./chroma_db"
+    BM25_INDEX_PATH: str = "./bm25_index"
+    COLLECTION_NAME: str = "doc_serve_collection"
+
+    # Chunking Configuration
+    DEFAULT_CHUNK_SIZE: int = 512
+    DEFAULT_CHUNK_OVERLAP: int = 50
+    MAX_CHUNK_SIZE: int = 2048
+    MIN_CHUNK_SIZE: int = 128
+
+    # Query Configuration
+    DEFAULT_TOP_K: int = 5
+    MAX_TOP_K: int = 50
+    DEFAULT_SIMILARITY_THRESHOLD: float = 0.7
+
+    # Rate Limiting
+    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"
+
+    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
+        ],
+        env_file_encoding="utf-8",
+        case_sensitive=True,
+    )
+
+
+@lru_cache
+def get_settings() -> Settings:
+    """Get cached settings instance."""
+    return Settings()
+
+
+settings = get_settings()
+
+
+def load_project_config(state_dir: Path) -> dict[str, Any]:
+    """Load project configuration from state directory.
+
+    Precedence: CLI flags > env vars > project config > defaults
+
+    Args:
+        state_dir: Path to the state directory containing config.json.
+
+    Returns:
+        Dictionary of configuration values from config.json, or empty dict.
+    """
+    config_path = state_dir / "config.json"
+    if config_path.exists():
+        try:
+            with open(config_path) as f:
+                return json.load(f)  # type: ignore[no-any-return]
+        except (json.JSONDecodeError, OSError) as e:
+            logger.warning(f"Failed to load project config from {config_path}: {e}")
+    return {}

doc_serve_server/indexing/__init__.py

@@ -0,0 +1,19 @@
+"""Indexing pipeline components for document processing."""
+
+from doc_serve_server.indexing.bm25_index import BM25IndexManager, get_bm25_manager
+from doc_serve_server.indexing.chunking import CodeChunker, ContextAwareChunker
+from doc_serve_server.indexing.document_loader import DocumentLoader
+from doc_serve_server.indexing.embedding import (
+    EmbeddingGenerator,
+    get_embedding_generator,
+)
+
+__all__ = [
+    "DocumentLoader",
+    "ContextAwareChunker",
+    "CodeChunker",
+    "EmbeddingGenerator",
+    "get_embedding_generator",
+    "BM25IndexManager",
+    "get_bm25_manager",
+]

doc_serve_server/indexing/bm25_index.py

@@ -0,0 +1,166 @@
+"""BM25 index manager for persistence and retrieval."""
+
+import logging
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Optional
+
+from llama_index.core.schema import BaseNode, NodeWithScore
+from llama_index.retrievers.bm25 import BM25Retriever
+
+from doc_serve_server.config import settings
+
+logger = logging.getLogger(__name__)
+
+
+class BM25IndexManager:
+    """
+    Manages the lifecycle of the BM25 index.
+
+    Handles building the index from nodes, persisting it to disk,
+    and loading it for retrieval.
+    """
+
+    def __init__(self, persist_dir: Optional[str] = None):
+        """
+        Initialize the BM25 index manager.
+
+        Args:
+            persist_dir: Directory for index persistence.
+        """
+        self.persist_dir = persist_dir or settings.BM25_INDEX_PATH
+        self._retriever: Optional[BM25Retriever] = None
+
+    @property
+    def is_initialized(self) -> bool:
+        """Check if the index is initialized."""
+        return self._retriever is not None
+
+    def initialize(self) -> None:
+        """
+        Load the index from disk if it exists.
+        """
+        persist_path = Path(self.persist_dir)
+        if (persist_path / "retriever.json").exists():
+            try:
+                self._retriever = BM25Retriever.from_persist_dir(str(persist_path))
+                logger.info(f"BM25 index loaded from {self.persist_dir}")
+            except Exception as e:
+                logger.error(f"Failed to load BM25 index: {e}")
+                self._retriever = None
+        else:
+            logger.info("No existing BM25 index found")
+
+    def build_index(self, nodes: Sequence[BaseNode]) -> None:
+        """
+        Build a new BM25 index from nodes and persist it.
+
+        Args:
+            nodes: List of LlamaIndex nodes.
+        """
+        logger.info(f"Building BM25 index with {len(nodes)} nodes")
+        self._retriever = BM25Retriever.from_defaults(nodes=nodes)
+        self.persist()
+
+    def persist(self) -> None:
+        """Persist the current index to disk."""
+        if not self._retriever:
+            logger.warning("No BM25 index to persist")
+            return
+
+        persist_path = Path(self.persist_dir)
+        persist_path.mkdir(parents=True, exist_ok=True)
+        self._retriever.persist(str(persist_path))
+        logger.info(f"BM25 index persisted to {self.persist_dir}")
+
+    def get_retriever(self, top_k: int = 5) -> BM25Retriever:
+        """
+        Get the BM25 retriever instance.
+
+        Args:
+            top_k: Number of results to return.
+
+        Returns:
+            The BM25Retriever instance.
+
+        Raises:
+            RuntimeError: If the index is not initialized.
+        """
+        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
+        return self._retriever
+
+    async def search_with_filters(
+        self,
+        query: str,
+        top_k: int = 5,
+        source_types: Optional[list[str]] = None,
+        languages: Optional[list[str]] = None,
+        max_results: Optional[int] = None,
+    ) -> list[NodeWithScore]:
+        """
+        Search the BM25 index with metadata filtering.
+
+        Args:
+            query: Search query string.
+            top_k: Number of results to return.
+            source_types: Filter by source types (doc, code, test).
+            languages: Filter by programming languages.
+
+        Returns:
+            List of NodeWithScore objects, filtered by metadata.
+        """
+        if not self._retriever:
+            raise RuntimeError("BM25 index not initialized")
+
+        # Get results for filtering
+        retriever_top_k = max_results if max_results is not None else (top_k * 3)
+        retriever = self.get_retriever(top_k=retriever_top_k)
+        nodes = await retriever.aretrieve(query)
+
+        # Apply metadata filtering
+        filtered_nodes = []
+        for node in nodes:
+            metadata = node.node.metadata
+
+            # Check source type filter
+            if source_types:
+                source_type = metadata.get("source_type", "doc")
+                if source_type not in source_types:
+                    continue
+
+            # Check language filter
+            if languages:
+                language = metadata.get("language")
+                if not language or language not in languages:
+                    continue
+
+            filtered_nodes.append(node)
+
+        # Return top_k results after filtering
+        return filtered_nodes[:top_k]
+
+    def reset(self) -> None:
+        """Reset the BM25 index by deleting persistent files."""
+        self._retriever = None
+        persist_path = Path(self.persist_dir)
+        if persist_path.exists():
+            for file in persist_path.glob("*"):
+                file.unlink()
+            persist_path.rmdir()
+        logger.info("BM25 index reset")
+
+
+# Global singleton instance
+_bm25_manager: Optional[BM25IndexManager] = None
+
+
+def get_bm25_manager() -> BM25IndexManager:
+    """Get the global BM25 manager instance."""
+    global _bm25_manager
+    if _bm25_manager is None:
+        _bm25_manager = BM25IndexManager()
+    return _bm25_manager