diff --git a/src/models/hybrid_rag.py b/src/models/hybrid_rag.py index bc2d7af..32a2596 100644 --- a/src/models/hybrid_rag.py +++ b/src/models/hybrid_rag.py @@ -15,15 +15,18 @@ class HybridRAGConfig(BaseModel): 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): @@ -37,12 +40,13 @@ class RelatedDossier(BaseModel): class HybridRAGResult(BaseModel): """Single result from HybridRAG query.""" - source_type: str = Field(..., description="Source: 'wiki', 'web', 'volatile'") + 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") @@ -57,6 +61,7 @@ class TimingBreakdown(BaseModel): 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") diff --git a/src/services/hybrid_rag_service.py b/src/services/hybrid_rag_service.py index d8ee28e..acfc94d 100644 --- a/src/services/hybrid_rag_service.py +++ b/src/services/hybrid_rag_service.py @@ -109,8 +109,9 @@ class HybridRAGService: timing["graph_ms"] = raw_results.get("timing", {}).get("graph_ms", 0) timing["web_ms"] = raw_results.get("timing", {}).get("web_ms", 0) timing["volatile_ms"] = raw_results.get("timing", {}).get("volatile_ms", 0) + timing["document_ms"] = raw_results.get("timing", {}).get("document_ms", 0) - # Phase 2: Three-Source RRF Fusion + # Phase 2: Four-Source RRF Fusion phase2_start = time.time() # Stage 1: Merge wiki sources (vector + graph) into single ranking @@ -120,12 +121,13 @@ class HybridRAGService: k=config.rrf_k ) - # Stage 2: Final RRF between wiki, volatile, and web + # Stage 2: Final RRF between wiki, volatile, document, and web # Volatile gets priority boost (smaller k = higher contribution per rank) fused_results = self._reciprocal_rank_fusion( wiki_results=wiki_merged, web_results=raw_results.get("web", []), volatile_results=raw_results.get("volatile", []), + document_results=raw_results.get("document", []), k=config.rrf_k ) timing["fusion_ms"] = (time.time() - phase2_start) * 1000 @@ -427,6 +429,63 @@ JSON:""" tasks["volatile"] = volatile_search() + # Paperless document search (separate from wiki vector search) + if config.enable_documents: + async def document_search(): + start = time.time() + try: + # Search in same collection but filter to doc_type=document + from src.core.multi_tenancy import get_qdrant_collection_name + collection_name = get_qdrant_collection_name(user) + + # Check if collection exists + exists = await self.vector.qdrant.collection_exists(collection_name) + if not exists: + return [], (time.time() - start) * 1000 + + # Get query embedding + query_embedding = await self.vector.ollama.embed_text(query) + + # Search with filter for doc_type=document + from qdrant_client.models import Filter, FieldCondition, MatchValue + search_results = self.vector.qdrant.client.search( + collection_name=collection_name, + query_vector=query_embedding, + limit=config.document_limit, + score_threshold=config.document_threshold, + query_filter=Filter( + must=[ + FieldCondition( + key="doc_type", + match=MatchValue(value="document") + ) + ] + ) + ) + + # Format results + formatted = [] + for r in search_results: + payload = r.payload or {} + formatted.append({ + "paperless_id": payload.get("paperless_id"), + "title": payload.get("title", "Untitled Document"), + "content": payload.get("chunk_text", ""), + "score": r.score, + "correspondent": payload.get("correspondent"), + "document_type": payload.get("document_type"), + "tags": payload.get("tags", []), + "original_filename": payload.get("original_filename"), + "source": "document" + }) + + return formatted, (time.time() - start) * 1000 + except Exception as e: + logger.error(f"Document search failed: {e}", exc_info=True) + return [], (time.time() - start) * 1000 + + tasks["document"] = document_search() + # Execute all searches in parallel results_dict = await asyncio.gather(*tasks.values()) @@ -440,7 +499,7 @@ JSON:""" logger.info( f"Parallel retrieval: vector={len(output.get('vector', []))}, " f"graph={len(output.get('graph', []))}, web={len(output.get('web', []))}, " - f"volatile={len(output.get('volatile', []))}" + f"volatile={len(output.get('volatile', []))}, document={len(output.get('document', []))}" ) return output @@ -531,10 +590,11 @@ JSON:""" wiki_results: List[Dict], web_results: List[Dict], volatile_results: Optional[List[Dict]] = None, + document_results: Optional[List[Dict]] = None, k: int = 60 ) -> List[Dict[str, Any]]: """ - Stage 2: Final RRF between wiki, volatile, and web. + Stage 2: Final RRF between wiki, volatile, document, and web. Wiki results are pre-merged from vector+graph. Volatile results get a priority boost (smaller effective k) since they represent @@ -544,6 +604,7 @@ JSON:""" wiki_results: Pre-merged wiki results from _merge_wiki_sources() web_results: Results from web search volatile_results: Results from volatile cache (fresh data) + document_results: Results from Paperless document search k: RRF constant (default 60) Returns: @@ -551,6 +612,7 @@ JSON:""" """ rrf_scores = {} volatile_results = volatile_results or [] + document_results = document_results or [] # Volatile results get priority boost (k/2 = stronger score per rank) volatile_k = k // 2 @@ -567,6 +629,19 @@ JSON:""" "source_type": "volatile" } + # Document results (Paperless) + for rank, result in enumerate(document_results, start=1): + paperless_id = result.get("paperless_id") + if not paperless_id: + continue + result_id = f"doc_{paperless_id}" + rrf_scores[result_id] = { + "result": result, + "rrf_score": 1 / (k + rank), + "sources": ["document"], + "source_type": "document" + } + # Wiki results (single source, already merged) for rank, result in enumerate(wiki_results, start=1): page_id = result.get("page_id") @@ -601,7 +676,8 @@ JSON:""" ) volatile_count = len([r for r in sorted_results if r["source_type"] == "volatile"]) - logger.info(f"Final RRF: {len(sorted_results)} results (wiki + volatile[{volatile_count}] + web)") + document_count = len([r for r in sorted_results if r["source_type"] == "document"]) + logger.info(f"Final RRF: {len(sorted_results)} results (wiki + volatile[{volatile_count}] + document[{document_count}] + web)") return sorted_results @@ -893,23 +969,35 @@ Ranking:""" for result_data in results: result = result_data.get("result", {}) related_dossiers = result_data.get("related_dossiers", []) + source_type = result_data.get("source_type", "unknown") + + # Build metadata based on source type + metadata = { + "entity_matches": result.get("entity_matches"), + "matched_entities": result.get("matched_entities"), + "engine": result.get("engine") + } + + # Add document-specific metadata + if source_type == "document": + metadata["correspondent"] = result.get("correspondent") + metadata["document_type"] = result.get("document_type") + metadata["tags"] = result.get("tags", []) + metadata["original_filename"] = result.get("original_filename") models.append(HybridRAGResult( - source_type=result_data.get("source_type", "unknown"), + source_type=source_type, title=result.get("title", "Untitled"), content=result.get("content", ""), url=result.get("url"), page_id=result.get("page_id"), page_path=result.get("path"), + paperless_id=result.get("paperless_id"), rrf_score=result_data.get("rrf_score", 0), final_rank=result_data.get("final_rank", 0), sources=result_data.get("sources", []), related_dossiers=[RelatedDossier(**d) for d in related_dossiers], - metadata={ - "entity_matches": result.get("entity_matches"), - "matched_entities": result.get("matched_entities"), - "engine": result.get("engine") - } + metadata=metadata )) return models