Phase C of memory system: Documents are now a retrieval source alongside wiki, volatile, and web search. Changes: - Add enable_documents, document_limit, document_threshold to HybridRAGConfig - Add paperless_id field to HybridRAGResult - Add document_ms timing to TimingBreakdown - Add document search to parallel retrieval (filters doc_type=document) - Update RRF fusion to include documents as fourth source - Add document metadata (correspondent, document_type, tags) to results HybridRAG now searches 4 sources in parallel: - Wiki (vector + graph merged) - Volatile cache (priority boost) - Paperless documents (new) - Web search 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
97 lines
5.3 KiB
Python
97 lines
5.3 KiB
Python
"""
|
|
HybridRAG models for multi-source search with RRF fusion.
|
|
|
|
Combines vector search (Qdrant), knowledge graph (Neo4j), and web search (SearXNG)
|
|
with Reciprocal Rank Fusion and LLM re-ranking.
|
|
"""
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing import List, Optional, Dict, Any
|
|
|
|
|
|
class HybridRAGConfig(BaseModel):
|
|
"""Configuration for HybridRAG query."""
|
|
vector_limit: int = Field(default=10, ge=1, le=50, description="Max vector results")
|
|
graph_limit: int = Field(default=10, ge=1, le=50, description="Max graph results")
|
|
web_limit: int = Field(default=5, ge=1, le=20, description="Max web results")
|
|
volatile_limit: int = Field(default=1, ge=1, le=5, description="Max volatile results (typically 1)")
|
|
document_limit: int = Field(default=5, ge=1, le=20, description="Max Paperless document results")
|
|
enable_vector: bool = Field(default=True, description="Enable vector search")
|
|
enable_graph: bool = Field(default=True, description="Enable graph search")
|
|
enable_web: bool = Field(default=True, description="Enable web search")
|
|
enable_volatile: bool = Field(default=True, description="Enable volatile cache search")
|
|
enable_documents: bool = Field(default=True, description="Enable Paperless document search")
|
|
enable_reranking: bool = Field(default=True, description="Enable LLM re-ranking")
|
|
enable_enrichment: bool = Field(default=True, description="Enable graph enrichment")
|
|
final_result_count: int = Field(default=10, ge=1, le=50, description="Final results to return")
|
|
rrf_k: int = Field(default=60, ge=1, le=100, description="RRF constant")
|
|
volatile_threshold: float = Field(default=0.8, ge=0.5, le=1.0, description="Volatile similarity threshold")
|
|
document_threshold: float = Field(default=0.6, ge=0.3, le=1.0, description="Document similarity threshold")
|
|
|
|
|
|
class RelatedDossier(BaseModel):
|
|
"""Related document metadata from graph enrichment."""
|
|
page_id: int
|
|
title: str
|
|
path: str
|
|
tag: str
|
|
shared_entities: int
|
|
|
|
|
|
class HybridRAGResult(BaseModel):
|
|
"""Single result from HybridRAG query."""
|
|
source_type: str = Field(..., description="Source: 'wiki', 'web', 'volatile', 'document'")
|
|
title: str
|
|
content: str
|
|
url: Optional[str] = Field(None, description="URL for web results")
|
|
page_id: Optional[int] = Field(None, description="Page ID for wiki results")
|
|
page_path: Optional[str] = Field(None, description="Wiki page path")
|
|
paperless_id: Optional[int] = Field(None, description="Paperless document ID")
|
|
rrf_score: float = Field(..., description="Reciprocal Rank Fusion score")
|
|
final_rank: int = Field(..., description="Final rank after re-ranking")
|
|
sources: List[str] = Field(..., description="Which sources included this result")
|
|
related_dossiers: List[RelatedDossier] = Field(default=[], description="Related documents via shared entities")
|
|
metadata: Dict[str, Any] = Field(default={}, description="Additional metadata")
|
|
|
|
|
|
class TimingBreakdown(BaseModel):
|
|
"""Performance timing breakdown for each phase."""
|
|
query_enhancement_ms: float = Field(..., description="Phase 0: Keyword/synonym extraction")
|
|
vector_ms: float = Field(..., description="Phase 1: Vector search")
|
|
graph_ms: float = Field(..., description="Phase 1: Graph search")
|
|
web_ms: float = Field(..., description="Phase 1: Web search")
|
|
volatile_ms: float = Field(default=0, description="Phase 1: Volatile cache search")
|
|
document_ms: float = Field(default=0, description="Phase 1: Paperless document search")
|
|
fusion_ms: float = Field(..., description="Phase 2: RRF fusion")
|
|
enrichment_ms: float = Field(..., description="Phase 3: Graph enrichment")
|
|
reranking_ms: float = Field(..., description="Phase 4: LLM re-ranking")
|
|
persistence_ms: float = Field(..., description="Phase 6: Search persistence")
|
|
total_ms: float = Field(..., description="Total end-to-end time")
|
|
|
|
|
|
class KeywordExtraction(BaseModel):
|
|
"""Extracted keywords and synonyms from query enhancement."""
|
|
core_keywords: List[str] = Field(default=[], description="Primary keywords")
|
|
entities: List[str] = Field(default=[], description="Named entities")
|
|
synonyms: Dict[str, List[str]] = Field(default={}, description="Synonyms map")
|
|
expansions: Dict[str, List[str]] = Field(default={}, description="Abbreviation expansions")
|
|
|
|
|
|
class HybridRAGResponse(BaseModel):
|
|
"""Response from HybridRAG query."""
|
|
query: str = Field(..., description="Original search query")
|
|
keywords: KeywordExtraction = Field(..., description="Extracted keywords/synonyms")
|
|
results: List[HybridRAGResult] = Field(..., description="Ranked search results")
|
|
context: str = Field(..., description="Formatted context for LLM consumption")
|
|
source_counts: Dict[str, int] = Field(..., description="Result counts by source")
|
|
total_results: int = Field(..., description="Total number of results")
|
|
timing: TimingBreakdown = Field(..., description="Performance breakdown")
|
|
config_used: HybridRAGConfig = Field(..., description="Configuration used")
|
|
search_id: Optional[str] = Field(None, description="Search ID for Librarian tracking")
|
|
|
|
|
|
class HybridRAGRequest(BaseModel):
|
|
"""Request for HybridRAG query."""
|
|
query: str = Field(..., min_length=1, max_length=500, description="Search query")
|
|
config: Optional[HybridRAGConfig] = Field(None, description="Custom configuration")
|