Files
library-desk/src/models/hybrid_rag.py
T
jpmschweitzerandClaude Opus 4.5 c01033505b feat: add Paperless document recall to HybridRAG
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>
2025-12-29 15:11:11 +01:00

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")