jpmschweitzerandClaude Fable 5
5c2284922d
Each retrieval leg (vector, graph, web, volatile, documents) now returns
(results, timing, error) instead of swallowing exceptions to an empty list.
The response reports per-leg status ('ok'/'failed'/'disabled') in
source_status and sets degraded=true when any enabled leg failed. Failed
legs still contribute no results (behavior unchanged) and are logged at
WARNING. Both fields are additive with defaults, so deploy order relative
to consumers does not matter.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1061 lines
40 KiB
Python
1061 lines
40 KiB
Python
"""
|
|
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__)
|
|
|
|
# Timeout for auxiliary LLM calls (keyword extraction, re-ranking).
|
|
# A hung Ollama call must not gate retrieval for the full client timeout.
|
|
LLM_CALL_TIMEOUT_SECONDS = 12.0
|
|
|
|
|
|
class HybridRAGService:
|
|
"""
|
|
Service for HybridRAG multi-source search with fusion and re-ranking.
|
|
"""
|
|
|
|
# Maps internal retrieval leg names to source_status keys in the response
|
|
SOURCE_STATUS_KEYS = {
|
|
"vector": "vector",
|
|
"graph": "graph",
|
|
"web": "web",
|
|
"volatile": "volatile",
|
|
"document": "documents",
|
|
}
|
|
|
|
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_llm_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)
|
|
timing["document_ms"] = raw_results.get("timing", {}).get("document_ms", 0)
|
|
|
|
# Phase 2: Four-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, 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
|
|
|
|
# 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
|
|
|
|
# Degradation signaling: a failed leg contributes no results, but the
|
|
# response says so instead of silently pretending the leg was empty
|
|
source_status = raw_results.get("source_status", {})
|
|
degraded = any(status == "failed" for status in source_status.values())
|
|
if degraded:
|
|
failed_legs = [leg for leg, status in source_status.items() if status == "failed"]
|
|
logger.warning(f"HybridRAG search degraded: failed legs: {failed_legs}")
|
|
|
|
# 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,
|
|
source_status=source_status,
|
|
degraded=degraded
|
|
)
|
|
|
|
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 asyncio.wait_for(
|
|
self.ollama.generate_text(
|
|
prompt=prompt,
|
|
model=self.reranker_model,
|
|
temperature=0.0 # Deterministic for consistent extraction
|
|
),
|
|
timeout=LLM_CALL_TIMEOUT_SECONDS
|
|
)
|
|
|
|
# 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 asyncio.TimeoutError:
|
|
logger.warning(
|
|
f"Keyword extraction timed out after {LLM_CALL_TIMEOUT_SECONDS}s, using fallback"
|
|
)
|
|
return {
|
|
"core_keywords": query.split(),
|
|
"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, Any]:
|
|
"""
|
|
Phase 1: Retrieve results from all sources in parallel.
|
|
|
|
Each retrieval leg returns (results, timing_ms, error) so that a
|
|
failed leg still contributes no results but is reported in
|
|
"source_status" instead of being silently swallowed.
|
|
|
|
Args:
|
|
query: Search query
|
|
user: User identifier
|
|
config: Search configuration
|
|
keywords_data: Extracted keywords/synonyms
|
|
|
|
Returns:
|
|
Dictionary with results from each source, timing, and per-leg
|
|
"source_status" ('ok', 'failed', or 'disabled')
|
|
"""
|
|
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, None
|
|
except Exception as e:
|
|
logger.warning(f"Vector search failed: {e}", exc_info=True)
|
|
return [], (time.time() - start) * 1000, e
|
|
|
|
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, None
|
|
except Exception as e:
|
|
logger.warning(f"Graph search failed: {e}", exc_info=True)
|
|
return [], (time.time() - start) * 1000, e
|
|
|
|
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, None
|
|
except Exception as e:
|
|
logger.warning(f"Web search failed: {e}", exc_info=True)
|
|
return [], (time.time() - start) * 1000, e
|
|
|
|
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, None
|
|
except Exception as e:
|
|
logger.warning(f"Volatile search failed: {e}", exc_info=True)
|
|
return [], (time.time() - start) * 1000, e
|
|
|
|
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, None
|
|
|
|
# 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, None
|
|
except Exception as e:
|
|
logger.warning(f"Document search failed: {e}", exc_info=True)
|
|
return [], (time.time() - start) * 1000, e
|
|
|
|
tasks["document"] = document_search()
|
|
|
|
# Execute all searches in parallel
|
|
# No return_exceptions needed: each leg captures its own exception
|
|
# and reports it via the (results, timing, error) tuple.
|
|
results_dict = await asyncio.gather(*tasks.values())
|
|
|
|
# Combine results with timing and per-leg status
|
|
output = {"timing": {}, "source_status": {}}
|
|
for i, source in enumerate(tasks.keys()):
|
|
results, source_timing, error = results_dict[i]
|
|
output[source] = results
|
|
output["timing"][f"{source}_ms"] = source_timing
|
|
status_key = self.SOURCE_STATUS_KEYS[source]
|
|
output["source_status"][status_key] = "failed" if error is not None else "ok"
|
|
|
|
# Legs that were not attempted are reported as disabled
|
|
for status_key in self.SOURCE_STATUS_KEYS.values():
|
|
output["source_status"].setdefault(status_key, "disabled")
|
|
|
|
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', []))}, document={len(output.get('document', []))}"
|
|
)
|
|
|
|
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,
|
|
document_results: Optional[List[Dict]] = None,
|
|
k: int = 60
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
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
|
|
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)
|
|
document_results: Results from Paperless document search
|
|
k: RRF constant (default 60)
|
|
|
|
Returns:
|
|
Final merged and sorted results
|
|
"""
|
|
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
|
|
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"
|
|
}
|
|
|
|
# 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")
|
|
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"])
|
|
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
|
|
|
|
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 asyncio.wait_for(
|
|
self.ollama.generate_text(
|
|
prompt=prompt,
|
|
model=self.reranker_model,
|
|
temperature=0.0 # Deterministic for consistent rankings
|
|
),
|
|
timeout=LLM_CALL_TIMEOUT_SECONDS
|
|
)
|
|
|
|
# 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 asyncio.TimeoutError:
|
|
logger.warning(
|
|
f"LLM re-ranking timed out after {LLM_CALL_TIMEOUT_SECONDS}s, using RRF order"
|
|
)
|
|
return results # Fallback to RRF order
|
|
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", [])
|
|
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=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=metadata
|
|
))
|
|
|
|
return models
|