agent-brain-rag 1.2.0__py3-none-any.whl → 3.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,26 @@ 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.",
+    )
+    # Queue status (Feature 115)
+    queue_pending: int = Field(
+        default=0,
+        ge=0,
+        description="Number of pending jobs in the queue",
+    )
+    queue_running: int = Field(
+        default=0,
+        ge=0,
+        description="Number of running jobs (0 or 1)",
+    )
+    current_job_running_time_ms: Optional[int] = Field(
+        None,
+        description="Running time of current job in milliseconds",
+    )
 
     model_config = {
         "json_schema_extra": {
@@ -120,6 +140,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/job.py

@@ -0,0 +1,289 @@
+"""Job queue models for indexing job management."""
+
+import hashlib
+from datetime import datetime, timezone
+from enum import Enum
+from pathlib import Path
+from typing import Optional
+
+from pydantic import BaseModel, Field, computed_field
+
+
+class JobStatus(str, Enum):
+    """Status of an indexing job."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    DONE = "done"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class JobProgress(BaseModel):
+    """Progress tracking for an indexing job."""
+
+    files_processed: int = Field(default=0, ge=0, description="Files processed so far")
+    files_total: int = Field(default=0, ge=0, description="Total files to process")
+    chunks_created: int = Field(default=0, ge=0, description="Chunks created so far")
+    current_file: str = Field(default="", description="Currently processing file")
+    updated_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+        description="Last progress update timestamp",
+    )
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def percent_complete(self) -> float:
+        """Calculate completion percentage."""
+        if self.files_total == 0:
+            return 0.0
+        return round((self.files_processed / self.files_total) * 100, 1)
+
+
+class JobRecord(BaseModel):
+    """Persistent job record for the queue."""
+
+    id: str = Field(..., description="Unique job identifier (job_<uuid12>)")
+    dedupe_key: str = Field(..., description="SHA256 hash for deduplication")
+
+    # Request parameters (normalized)
+    folder_path: str = Field(..., description="Resolved, normalized folder path")
+    include_code: bool = Field(default=False, description="Whether to index code files")
+    operation: str = Field(
+        default="index", description="Operation type: 'index' or 'add'"
+    )
+
+    # Optional request parameters
+    chunk_size: int = Field(default=512, description="Chunk size in tokens")
+    chunk_overlap: int = Field(default=50, description="Chunk overlap in tokens")
+    recursive: bool = Field(default=True, description="Recursive folder scan")
+    generate_summaries: bool = Field(
+        default=False, description="Generate LLM summaries"
+    )
+    supported_languages: Optional[list[str]] = Field(
+        default=None, description="Languages to index"
+    )
+    include_patterns: Optional[list[str]] = Field(
+        default=None, description="File patterns to include"
+    )
+    exclude_patterns: Optional[list[str]] = Field(
+        default=None, description="File patterns to exclude"
+    )
+
+    # Job state
+    status: JobStatus = Field(
+        default=JobStatus.PENDING, description="Current job status"
+    )
+    cancel_requested: bool = Field(
+        default=False, description="Flag for graceful cancellation"
+    )
+
+    # Timestamps
+    enqueued_at: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+        description="When the job was enqueued",
+    )
+    started_at: Optional[datetime] = Field(
+        default=None, description="When the job started running"
+    )
+    finished_at: Optional[datetime] = Field(
+        default=None, description="When the job finished (done, failed, or cancelled)"
+    )
+
+    # Results and metadata
+    error: Optional[str] = Field(default=None, description="Error message if failed")
+    retry_count: int = Field(default=0, ge=0, description="Number of retry attempts")
+    progress: Optional[JobProgress] = Field(
+        default=None, description="Progress tracking"
+    )
+    total_chunks: int = Field(default=0, ge=0, description="Total chunks indexed")
+    total_documents: int = Field(default=0, ge=0, description="Total documents indexed")
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def execution_time_ms(self) -> Optional[int]:
+        """Calculate execution time in milliseconds."""
+        if self.started_at is None:
+            return None
+        end_time = self.finished_at or datetime.now(timezone.utc)
+        delta = end_time - self.started_at
+        return int(delta.total_seconds() * 1000)
+
+    @staticmethod
+    def compute_dedupe_key(
+        folder_path: str,
+        include_code: bool,
+        operation: str,
+        include_patterns: Optional[list[str]] = None,
+        exclude_patterns: Optional[list[str]] = None,
+    ) -> str:
+        """Compute deduplication key from job parameters.
+
+        Args:
+            folder_path: Normalized, resolved folder path.
+            include_code: Whether to include code files.
+            operation: Operation type (index or add).
+            include_patterns: Optional include patterns.
+            exclude_patterns: Optional exclude patterns.
+
+        Returns:
+            SHA256 hash of normalized parameters.
+        """
+        # Normalize path (resolve and lowercase on case-insensitive systems)
+        resolved = str(Path(folder_path).resolve())
+
+        # Build dedupe string
+        parts = [
+            resolved,
+            str(include_code),
+            operation,
+            ",".join(sorted(include_patterns or [])),
+            ",".join(sorted(exclude_patterns or [])),
+        ]
+        dedupe_string = "|".join(parts)
+
+        return hashlib.sha256(dedupe_string.encode()).hexdigest()
+
+
+class JobEnqueueResponse(BaseModel):
+    """Response when enqueueing a job."""
+
+    job_id: str = Field(..., description="Unique job identifier")
+    status: str = Field(default="pending", description="Job status")
+    queue_position: int = Field(
+        default=0, ge=0, description="Position in the queue (0 = first)"
+    )
+    queue_length: int = Field(default=0, ge=0, description="Total jobs in queue")
+    message: str = Field(..., description="Human-readable status message")
+    dedupe_hit: bool = Field(
+        default=False, description="True if this was a duplicate request"
+    )
+
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "job_id": "job_abc123def456",
+                    "status": "pending",
+                    "queue_position": 2,
+                    "queue_length": 5,
+                    "message": "Job queued for /path/to/docs",
+                    "dedupe_hit": False,
+                }
+            ]
+        }
+    }
+
+
+class JobListResponse(BaseModel):
+    """Response for listing jobs."""
+
+    jobs: list["JobSummary"] = Field(default_factory=list, description="List of jobs")
+    total: int = Field(default=0, ge=0, description="Total number of jobs")
+    pending: int = Field(default=0, ge=0, description="Number of pending jobs")
+    running: int = Field(default=0, ge=0, description="Number of running jobs")
+    completed: int = Field(default=0, ge=0, description="Number of completed jobs")
+    failed: int = Field(default=0, ge=0, description="Number of failed jobs")
+
+
+class JobSummary(BaseModel):
+    """Summary view of a job for list responses."""
+
+    id: str = Field(..., description="Job identifier")
+    status: JobStatus = Field(..., description="Current status")
+    folder_path: str = Field(..., description="Folder being indexed")
+    operation: str = Field(..., description="Operation type")
+    include_code: bool = Field(..., description="Whether indexing code")
+    enqueued_at: datetime = Field(..., description="When queued")
+    started_at: Optional[datetime] = Field(default=None, description="When started")
+    finished_at: Optional[datetime] = Field(default=None, description="When finished")
+    progress_percent: float = Field(default=0.0, description="Completion percentage")
+    error: Optional[str] = Field(default=None, description="Error message if failed")
+
+    @classmethod
+    def from_record(cls, record: JobRecord) -> "JobSummary":
+        """Create a summary from a full job record."""
+        return cls(
+            id=record.id,
+            status=record.status,
+            folder_path=record.folder_path,
+            operation=record.operation,
+            include_code=record.include_code,
+            enqueued_at=record.enqueued_at,
+            started_at=record.started_at,
+            finished_at=record.finished_at,
+            progress_percent=(
+                record.progress.percent_complete if record.progress else 0.0
+            ),
+            error=record.error,
+        )
+
+
+class JobDetailResponse(BaseModel):
+    """Detailed response for a single job."""
+
+    id: str = Field(..., description="Job identifier")
+    status: JobStatus = Field(..., description="Current status")
+    folder_path: str = Field(..., description="Folder being indexed")
+    operation: str = Field(..., description="Operation type")
+    include_code: bool = Field(..., description="Whether indexing code")
+
+    # Timestamps
+    enqueued_at: datetime = Field(..., description="When queued")
+    started_at: Optional[datetime] = Field(default=None, description="When started")
+    finished_at: Optional[datetime] = Field(default=None, description="When finished")
+    execution_time_ms: Optional[int] = Field(
+        default=None, description="Execution time in ms"
+    )
+
+    # Progress
+    progress: Optional[JobProgress] = Field(
+        default=None, description="Progress details"
+    )
+
+    # Results
+    total_documents: int = Field(default=0, description="Documents indexed")
+    total_chunks: int = Field(default=0, description="Chunks created")
+    error: Optional[str] = Field(default=None, description="Error message if failed")
+    retry_count: int = Field(default=0, description="Retry attempts")
+    cancel_requested: bool = Field(
+        default=False, description="Whether cancellation requested"
+    )
+
+    @classmethod
+    def from_record(cls, record: JobRecord) -> "JobDetailResponse":
+        """Create a detail response from a full job record."""
+        return cls(
+            id=record.id,
+            status=record.status,
+            folder_path=record.folder_path,
+            operation=record.operation,
+            include_code=record.include_code,
+            enqueued_at=record.enqueued_at,
+            started_at=record.started_at,
+            finished_at=record.finished_at,
+            execution_time_ms=record.execution_time_ms,
+            progress=record.progress,
+            total_documents=record.total_documents,
+            total_chunks=record.total_chunks,
+            error=record.error,
+            retry_count=record.retry_count,
+            cancel_requested=record.cancel_requested,
+        )
+
+
+class QueueStats(BaseModel):
+    """Statistics about the job queue."""
+
+    pending: int = Field(default=0, ge=0, description="Pending jobs count")
+    running: int = Field(default=0, ge=0, description="Running jobs count")
+    completed: int = Field(default=0, ge=0, description="Completed jobs count")
+    failed: int = Field(default=0, ge=0, description="Failed jobs count")
+    cancelled: int = Field(default=0, ge=0, description="Cancelled jobs count")
+    total: int = Field(default=0, ge=0, description="Total jobs count")
+    current_job_id: Optional[str] = Field(
+        default=None, description="Currently running job ID"
+    )
+    current_job_running_time_ms: Optional[int] = Field(
+        default=None, description="Current job running time in ms"
+    )

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):
@@ -32,14 +34,14 @@ class QueryRequest(BaseModel):
         description="Number of results to return",
     )
     similarity_threshold: float = Field(
-        default=0.7,
+        default=0.3,
         ge=0.0,
         le=1.0,
         description="Minimum similarity score (0-1)",
     )
     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,
@@ -90,7 +92,7 @@ class QueryRequest(BaseModel):
                 {
                     "query": "How do I configure authentication?",
                     "top_k": 5,
-                    "similarity_threshold": 0.7,
+                    "similarity_threshold": 0.3,
                     "mode": "hybrid",
                     "alpha": 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/project_root.py

@@ -1,4 +1,4 @@
-"""Project root resolution for per-project doc-serve instances."""
+"""Project root resolution for per-project Agent Brain instances."""
 
 import logging
 import subprocess