""" HybridRAG service combining vector, graph, and web search. 6-Phase Pipeline: 0. Query Enhancement - Extract keywords/synonyms with LLM 1. Parallel Retrieval - Vector + Graph + Web search 2. RRF Fusion - Merge results with Reciprocal Rank Fusion 3. Enrichment - Add related dossiers via graph 4. LLM Re-ranking - Re-rank with configured Ollama model 5. Context Formatting - Format for LLM consumption 6. Persistence - Store for Librarian processing """ import asyncio import time import json import uuid from typing import List, Dict, Any, Optional import logging from src.services.vector_service import VectorService from src.services.graph_service import GraphService from src.services.volatile_service import VolatileCacheService from src.clients.searxng_client import SearXNGClient from src.clients.ollama_client import OllamaClient from src.clients.content_extractor import ContentExtractor from src.config import Settings from src.models.hybrid_rag import ( HybridRAGConfig, HybridRAGRequest, HybridRAGResponse, HybridRAGResult, TimingBreakdown, KeywordExtraction, RelatedDossier ) from src.core.multi_tenancy import get_neo4j_user_base_label, get_neo4j_user_label logger = logging.getLogger(__name__) class HybridRAGService: """ Service for HybridRAG multi-source search with fusion and re-ranking. """ def __init__( self, vector_service: VectorService, graph_service: GraphService, searxng_client: SearXNGClient, ollama_client: OllamaClient, content_extractor: ContentExtractor, settings: Settings, volatile_service: Optional[VolatileCacheService] = None ): """ Initialize HybridRAG service. Args: vector_service: Service for Qdrant vector search graph_service: Service for Neo4j graph search searxng_client: Client for web search ollama_client: Client for LLM (keyword extraction, re-ranking) content_extractor: Client for extracting full content from URLs settings: Application settings volatile_service: Service for volatile cache search (optional) """ self.vector = vector_service self.graph = graph_service self.searxng = searxng_client self.ollama = ollama_client self.content_extractor = content_extractor self.settings = settings self.volatile = volatile_service self.reranker_model = settings.ollama_model async def search( self, query: str, user: str, config: Optional[HybridRAGConfig] = None ) -> HybridRAGResponse: """ Execute HybridRAG search across all sources. Args: query: Search query user: User identifier config: Optional configuration override Returns: Complete search response with ranked results and timing """ start_time = time.time() timing = {} # Use default config if not provided if not config: config = HybridRAGConfig() logger.info(f"HybridRAG search: '{query}' for user '{user}'") # Phase 0: Query Enhancement phase0_start = time.time() keywords_data = await self._extract_keywords_and_synonyms(query) timing["query_enhancement_ms"] = (time.time() - phase0_start) * 1000 # Phase 1: Parallel Retrieval phase1_start = time.time() raw_results = await self._retrieve_parallel(query, user, config, keywords_data) timing["vector_ms"] = raw_results.get("timing", {}).get("vector_ms", 0) 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) # Phase 2: Three-Source RRF Fusion phase2_start = time.time() # Stage 1: Merge wiki sources (vector + graph) into single ranking wiki_merged = self._merge_wiki_sources( vector_results=raw_results.get("vector", []), graph_results=raw_results.get("graph", []), k=config.rrf_k ) # Stage 2: Final RRF between wiki, volatile, 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", []), k=config.rrf_k ) timing["fusion_ms"] = (time.time() - phase2_start) * 1000 # Phase 3: Enrichment phase3_start = time.time() if config.enable_enrichment: enriched_results = await self._enrich_with_related_dossiers(fused_results, user) else: enriched_results = fused_results timing["enrichment_ms"] = (time.time() - phase3_start) * 1000 # Phase 4: LLM Re-ranking phase4_start = time.time() if config.enable_reranking and len(enriched_results) > 1: reranked_results = await self._rerank_with_llm(enriched_results[:20], query) else: reranked_results = enriched_results timing["reranking_ms"] = (time.time() - phase4_start) * 1000 # Limit to final result count final_results = reranked_results[:config.final_result_count] # Update final ranks for i, result in enumerate(final_results, start=1): result["final_rank"] = i # Convert to HybridRAGResult models result_models = self._convert_to_result_models(final_results) # Phase 5: Context Formatting context = self._format_context_for_llm(result_models) # Calculate source counts source_counts = {} for result in result_models: for source in result.sources: source_counts[source] = source_counts.get(source, 0) + 1 timing["total_ms"] = (time.time() - start_time) * 1000 # Phase 6: Persistence (async, non-blocking) phase6_start = time.time() search_id = await self._persist_search_for_librarian( query=query, user=user, keywords_data=keywords_data, raw_results=raw_results, final_results=final_results, timing=timing ) timing["persistence_ms"] = (time.time() - phase6_start) * 1000 # Build response return HybridRAGResponse( query=query, keywords=KeywordExtraction(**keywords_data), results=result_models, context=context, source_counts=source_counts, total_results=len(result_models), timing=TimingBreakdown(**timing), config_used=config, search_id=search_id ) async def _extract_keywords_and_synonyms(self, query: str) -> Dict[str, Any]: """ Phase 0: Extract keywords, entities, and synonyms using LLM. Args: query: Search query Returns: Dictionary with keywords, entities, synonyms, expansions """ prompt = f"""Extract search terms from this query. Query: "{query}" RULES: - Extract ONLY keywords explicitly present or directly implied in the query - Do NOT invent terms, concepts, or synonyms not clearly related - Do NOT add general knowledge or associations - Provide synonyms ONLY for technical terms with well-known alternatives - Return valid JSON only, no commentary Return format: {{ "core_keywords": ["words", "from", "query"], "synonyms": {{"term": ["direct", "alternatives"]}} }} JSON:""" try: response = await self.ollama.generate_text( prompt=prompt, model=self.reranker_model, temperature=0.0 # Deterministic for consistent extraction ) # Parse JSON response (handle potential extra text) response_clean = response.strip() # Try to extract JSON if wrapped in text if '{' in response_clean: json_start = response_clean.find('{') json_end = response_clean.rfind('}') + 1 response_clean = response_clean[json_start:json_end] keywords_data = json.loads(response_clean) # Ensure all required fields exist result = { "core_keywords": keywords_data.get("core_keywords", []), "entities": keywords_data.get("entities", []), "synonyms": keywords_data.get("synonyms", {}), "expansions": keywords_data.get("expansions", {}) } logger.info(f"Extracted keywords: {result['core_keywords'][:5]}, synonyms: {len(result['synonyms'])} terms") return result except json.JSONDecodeError as e: logger.warning(f"Failed to parse LLM keyword extraction: {e}, using fallback") # Fallback to simple extraction words = query.split() return { "core_keywords": words, "entities": [], "synonyms": {}, "expansions": {} } except Exception as e: logger.error(f"Keyword extraction failed: {e}", exc_info=True) return { "core_keywords": query.split(), "entities": [], "synonyms": {}, "expansions": {} } async def _retrieve_parallel( self, query: str, user: str, config: HybridRAGConfig, keywords_data: Dict[str, Any] ) -> Dict[str, List]: """ Phase 1: Retrieve results from all sources in parallel. Args: query: Search query user: User identifier config: Search configuration keywords_data: Extracted keywords/synonyms Returns: Dictionary with results from each source and timing """ tasks = {} timing = {} # Vector search if config.enable_vector: async def vector_search(): start = time.time() try: response = await self.vector.search( query=query, user=user, limit=config.vector_limit, score_threshold=self.settings.vector_similarity_threshold ) results = [ { "page_id": r.page_id, "title": r.page_title, "content": r.content, "path": r.page_path, "score": r.score, "source": "vector" } for r in response.results ] return results, (time.time() - start) * 1000 except Exception as e: logger.error(f"Vector search failed: {e}", exc_info=True) return [], (time.time() - start) * 1000 tasks["vector"] = vector_search() # Graph search if config.enable_graph: async def graph_search(): start = time.time() try: # Skip synonyms for graph search - only use core keywords # Synonyms like "author" can match unrelated entities like "author2000" graph_keywords = { "core_keywords": keywords_data.get("core_keywords", []), "entities": keywords_data.get("entities", []), "synonyms": {}, # No synonyms for exact entity matching "expansions": {} } results = await self.graph.search_documents( query=query, user=user, limit=config.graph_limit, keywords_data=graph_keywords ) formatted = [ { "page_id": r["page_id"], "title": r["title"], "content": "", # Graph doesn't return content "path": r["path"], "entity_matches": r.get("entity_matches", 0), "matched_entities": r.get("matched_entities", []), "source": "graph" } for r in results ] return formatted, (time.time() - start) * 1000 except Exception as e: logger.error(f"Graph search failed: {e}", exc_info=True) return [], (time.time() - start) * 1000 tasks["graph"] = graph_search() # Web search with content extraction if config.enable_web: async def web_search(): start = time.time() try: results = await self.searxng.search_general( query=query, limit=config.web_limit ) # Extract full content from URLs using Trafilatura urls = [r.get("url") for r in results if r.get("url")] extraction_results = await self.content_extractor.extract_batch(urls) # Map extracted content back to results by URL url_to_content = { ext.url: ext.content for ext in extraction_results if ext.success and ext.content } formatted = [ { "url": r.get("url"), "title": r.get("title", ""), "content": url_to_content.get(r.get("url"), r.get("content", "")), "snippet": r.get("content", ""), # Keep original snippet "engine": r.get("engine", ""), "source": "web" } for r in results ] return formatted, (time.time() - start) * 1000 except Exception as e: logger.error(f"Web search failed: {e}", exc_info=True) return [], (time.time() - start) * 1000 tasks["web"] = web_search() # Volatile cache search if config.enable_volatile and self.volatile: async def volatile_search(): start = time.time() try: results = await self.volatile.search( user=user, query=query, limit=config.volatile_limit, score_threshold=config.volatile_threshold ) formatted = [ { "key": r.key, "namespace": r.namespace, "title": f"{r.namespace}: {r.key}", "content": r.data.get("text", "") if isinstance(r.data, dict) else str(r.data), "raw_data": r.data, "source_api": r.source, "ttl_remaining": r.ttl_remaining, "source": "volatile" } for r in results ] return formatted, (time.time() - start) * 1000 except Exception as e: logger.error(f"Volatile search failed: {e}", exc_info=True) return [], (time.time() - start) * 1000 tasks["volatile"] = volatile_search() # Execute all searches in parallel results_dict = await asyncio.gather(*tasks.values()) # Combine results with timing output = {"timing": {}} for i, source in enumerate(tasks.keys()): results, source_timing = results_dict[i] output[source] = results output["timing"][f"{source}_ms"] = source_timing 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', []))}" ) return output def _merge_wiki_sources( self, vector_results: List[Dict], graph_results: List[Dict], k: int = 60 ) -> List[Dict[str, Any]]: """ Stage 1: Merge vector and graph into single wiki ranking using RRF. Both sources search the same wiki pool, so we combine them before final RRF with web to avoid double-counting wiki pages. Args: vector_results: Results from vector search graph_results: Results from graph search k: RRF constant (default 60) Returns: Merged wiki results sorted by wiki RRF score """ wiki_scores = {} # Process vector results for rank, result in enumerate(vector_results, start=1): page_id = result.get("page_id") if not page_id: continue result_id = f"page_{page_id}" if result_id not in wiki_scores: wiki_scores[result_id] = { "result": dict(result), # Copy to avoid mutation "wiki_rrf_score": 0.0, "found_by": [] } wiki_scores[result_id]["wiki_rrf_score"] += 1 / (k + rank) wiki_scores[result_id]["found_by"].append("vector") # Process graph results for rank, result in enumerate(graph_results, start=1): page_id = result.get("page_id") if not page_id: continue result_id = f"page_{page_id}" if result_id not in wiki_scores: wiki_scores[result_id] = { "result": dict(result), "wiki_rrf_score": 0.0, "found_by": [] } wiki_scores[result_id]["wiki_rrf_score"] += 1 / (k + rank) wiki_scores[result_id]["found_by"].append("graph") # Add graph metadata to existing result wiki_scores[result_id]["result"]["entity_matches"] = result.get("entity_matches") wiki_scores[result_id]["result"]["matched_entities"] = result.get("matched_entities") # Sort by wiki RRF score sorted_wiki = sorted( wiki_scores.values(), key=lambda x: x["wiki_rrf_score"], reverse=True ) # Return merged results with wiki ranking merged = [] for wiki_rank, item in enumerate(sorted_wiki, start=1): merged.append({ **item["result"], "wiki_rank": wiki_rank, "wiki_rrf_score": item["wiki_rrf_score"], "found_by": item["found_by"], "source": "wiki" }) logger.info(f"Wiki merge: {len(merged)} unique pages from vector+graph") return merged def _reciprocal_rank_fusion( self, wiki_results: List[Dict], web_results: List[Dict], volatile_results: Optional[List[Dict]] = None, k: int = 60 ) -> List[Dict[str, Any]]: """ Stage 2: Final RRF between wiki, volatile, and web. Wiki results are pre-merged from vector+graph. Volatile results get a priority boost (smaller effective k) since they represent current, time-sensitive information. Args: wiki_results: Pre-merged wiki results from _merge_wiki_sources() web_results: Results from web search volatile_results: Results from volatile cache (fresh data) k: RRF constant (default 60) Returns: Final merged and sorted results """ rrf_scores = {} volatile_results = volatile_results or [] # Volatile results get priority boost (k/2 = stronger score per rank) volatile_k = k // 2 for rank, result in enumerate(volatile_results, start=1): key = result.get("key") namespace = result.get("namespace", "unknown") if not key: continue result_id = f"volatile_{namespace}_{key}" rrf_scores[result_id] = { "result": result, "rrf_score": 1 / (volatile_k + rank), # Priority boost "sources": ["volatile"], "source_type": "volatile" } # Wiki results (single source, already merged) for rank, result in enumerate(wiki_results, start=1): page_id = result.get("page_id") if not page_id: continue result_id = f"page_{page_id}" rrf_scores[result_id] = { "result": result, "rrf_score": 1 / (k + rank), "sources": result.get("found_by", ["wiki"]), "source_type": "wiki" } # Web results (single source) for rank, result in enumerate(web_results, start=1): url = result.get("url") if not url: continue result_id = f"url_{hash(url)}" rrf_scores[result_id] = { "result": result, "rrf_score": 1 / (k + rank), "sources": ["web"], "source_type": "web" } # Sort by RRF score descending sorted_results = sorted( rrf_scores.values(), key=lambda x: x["rrf_score"], reverse=True ) 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)") return sorted_results async def _enrich_with_related_dossiers( self, results: List[Dict[str, Any]], user: str ) -> List[Dict[str, Any]]: """ Phase 3: Enrich results with related documents via shared entities. Args: results: Fused results user: User identifier Returns: Results with related_dossiers added """ for result in results: result_data = result.get("result", {}) page_id = result_data.get("page_id") if page_id: try: related_docs = await self.graph.get_related_documents( page_id=page_id, user=user, limit=5 ) # Convert to RelatedDossier format related_dossiers = [] for doc in related_docs: for tag in doc.get("tags", [])[:3]: # Max 3 tags per doc related_dossiers.append({ "page_id": doc["page_id"], "title": doc["title"], "path": doc["path"], "tag": tag, "shared_entities": doc["shared_entities"] }) result["related_dossiers"] = related_dossiers[:5] # Limit to 5 total except Exception as e: logger.warning(f"Failed to get related docs for page {page_id}: {e}") result["related_dossiers"] = [] else: result["related_dossiers"] = [] return results async def _rerank_with_llm( self, results: List[Dict[str, Any]], query: str ) -> List[Dict[str, Any]]: """ Phase 4: Re-rank results using LLM for better relevance. Args: results: Results to re-rank (top 20) query: Original search query Returns: Re-ranked results """ if len(results) <= 1: return results try: # Build prompt with numbered results docs_text = "\n".join([ f"{i+1}. {r['result'].get('title', 'Untitled')} - {r['result'].get('content', '')[:200]}..." for i, r in enumerate(results) ]) prompt = f"""Rank these documents by relevance to the query. Query: {query} Documents: {docs_text} RULES: - Rank ONLY by how well content answers the query - Do NOT consider document length, formatting, or style - Do NOT add explanation or commentary - Return ONLY comma-separated numbers, most relevant first Example output: 3,1,5,2,4 Ranking:""" response = await self.ollama.generate_text( prompt=prompt, model=self.reranker_model, temperature=0.0 # Deterministic for consistent rankings ) # Parse response: "3,1,5,2,4" → [2, 0, 4, 1, 3] (0-indexed) indices_str = response.strip().split('\n')[0] # Take first line indices = [int(x.strip()) - 1 for x in indices_str.split(",") if x.strip().isdigit()] # Reorder results according to LLM ranking reranked = [] for idx in indices: if 0 <= idx < len(results): reranked.append(results[idx]) # Add any results that weren't in the LLM response for i, result in enumerate(results): if i not in indices and result not in reranked: reranked.append(result) logger.info(f"LLM re-ranking: reordered {len(reranked)} results") return reranked except Exception as e: logger.warning(f"LLM re-ranking failed: {e}, using RRF order") return results # Fallback to RRF order def _format_context_for_llm(self, results: List[HybridRAGResult]) -> str: """ Phase 5: Format results into context for LLM consumption. Args: results: Ranked results Returns: Formatted context string """ context_parts = [] for i, result in enumerate(results[:10], start=1): # Source indicator source_tag = f"[{result.source_type.upper()}]" # Related dossiers if available related = "" if result.related_dossiers: tags = ", ".join([d.tag for d in result.related_dossiers[:3]]) related = f"\n Related research: {tags}" # Build context entry content_preview = result.content[:300] if result.content else "(no content)" context_parts.append( f"{i}. {source_tag} {result.title}\n" f" {content_preview}...{related}" ) return "\n\n".join(context_parts) async def _persist_search_for_librarian( self, query: str, user: str, keywords_data: Dict[str, Any], raw_results: Dict[str, List], final_results: List[Dict[str, Any]], timing: Dict[str, float] ) -> Optional[str]: """ Phase 6: Store search query and results for Librarian processing. Creates SearchQuery node in Neo4j with relationships to found documents and web results for offline knowledge consolidation. Args: query: Search query user: User identifier keywords_data: Extracted keywords/synonyms raw_results: Results from each source final_results: Final ranked results timing: Performance timing Returns: Search ID for tracking """ try: user_base_label = get_neo4j_user_base_label(user) search_id = str(uuid.uuid4()) # Create SearchQuery node create_query = f""" CREATE (sq:{user_base_label}_SearchQuery:SearchQuery {{ id: $search_id, query: $query, user: $user, timestamp: datetime(), processed: false, total_results: $total_results, vector_count: $vector_count, graph_count: $graph_count, web_count: $web_count, keywords: $keywords, synonyms: $synonyms, timing_ms: $timing_ms }}) RETURN sq.id as id """ result = await self.graph.neo4j.execute_query(create_query, { "search_id": search_id, "query": query, "user": user, "total_results": len(final_results), "vector_count": len(raw_results.get("vector", [])), "graph_count": len(raw_results.get("graph", [])), "web_count": len(raw_results.get("web", [])), "keywords": keywords_data.get("core_keywords", []), "synonyms": json.dumps(keywords_data.get("synonyms", {})), "timing_ms": timing.get("total_ms", 0) }) # Link to found wiki documents (top 20) for rank, result_data in enumerate(final_results[:20], start=1): result = result_data.get("result", {}) page_id = result.get("page_id") if page_id: link_doc_query = f""" MATCH (sq:{user_base_label}_SearchQuery:SearchQuery {{id: $search_id}}) MATCH (d:Document {{page_id: $page_id}}) MERGE (sq)-[f:FOUND]->(d) SET f.source = $source, f.rank = $rank, f.rrf_score = $rrf_score, f.final_rank = $final_rank """ await self.graph.neo4j.execute_query(link_doc_query, { "search_id": search_id, "page_id": page_id, "source": result_data.get("source_type", "unknown"), "rank": rank, "rrf_score": result_data.get("rrf_score", 0), "final_rank": result_data.get("final_rank", rank) }) # Store web results as WebResult nodes (top 10) web_results = [r for r in final_results[:10] if r.get("result", {}).get("url")] for rank, result_data in enumerate(web_results, start=1): result = result_data.get("result", {}) create_web_query = f""" MATCH (sq:{user_base_label}_SearchQuery:SearchQuery {{id: $search_id}}) CREATE (wr:{user_base_label}_WebResult:WebResult {{ url: $url, title: $title, content: $content, search_id: $search_id, timestamp: datetime() }}) CREATE (sq)-[:FOUND {{ source: "web", rank: $rank, rrf_score: $rrf_score }}]->(wr) """ await self.graph.neo4j.execute_query(create_web_query, { "search_id": search_id, "url": result.get("url"), "title": result.get("title", ""), "content": result.get("content", "")[:1000], # Truncate "rank": rank, "rrf_score": result_data.get("rrf_score", 0) }) logger.info(f"Persisted search {search_id} for Librarian processing") return search_id except Exception as e: logger.error(f"Failed to persist search for Librarian: {e}", exc_info=True) return None def _convert_to_result_models(self, results: List[Dict[str, Any]]) -> List[HybridRAGResult]: """ Convert internal result format to HybridRAGResult models. Args: results: Internal result dictionaries Returns: List of HybridRAGResult models """ models = [] for result_data in results: result = result_data.get("result", {}) related_dossiers = result_data.get("related_dossiers", []) models.append(HybridRAGResult( source_type=result_data.get("source_type", "unknown"), title=result.get("title", "Untitled"), content=result.get("content", ""), url=result.get("url"), page_id=result.get("page_id"), page_path=result.get("path"), 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") } )) return models