jpmschweitzerandClaude Opus 4.5
37f8e1819e
Build and Push / build (release) Successful in 28s
- Migrate volatile backend from Redis to Qdrant for semantic search
- Add natural language conversion for structured data embedding
- Simplify API: /volatile/search, /volatile/store, /{namespace}/{key}
- Integrate volatile into HybridRAG with priority boost in RRF fusion
- Add POST /maintenance/cleanup/volatile for expiry purging
- Update tests for new Qdrant-based architecture (37/37 pass)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
916 lines
33 KiB
Python
916 lines
33 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__)
|
|
|
|
|
|
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
|