feat(library-desk): implement HybridRAG query system
HybridRAG Service: - Combine vector (Qdrant), graph (Neo4j), and web (SearXNG) search - Reciprocal Rank Fusion (RRF) for result merging - LLM re-ranking with mistral-nemo - Graph enrichment with related dossiers - Query enhancement with keyword/synonym extraction - Search result persistence for offline processing Router: - POST /query/hybrid endpoint - Configurable search limits per source - Enable/disable individual sources - Timing breakdown for performance monitoring Models: - HybridRAGRequest, HybridRAGResponse - HybridRAGResult with source tracking - KeywordExtraction for query analysis - TimingBreakdown for performance metrics Tests: - End-to-end HybridRAG query tests - RRF fusion algorithm validation - Multi-source result merging
This commit is contained in:
@@ -0,0 +1,87 @@
|
||||
"""
|
||||
HybridRAG models for multi-source search with RRF fusion.
|
||||
|
||||
Combines vector search (Qdrant), knowledge graph (Neo4j), and web search (SearXNG)
|
||||
with Reciprocal Rank Fusion and LLM re-ranking.
|
||||
"""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
|
||||
class HybridRAGConfig(BaseModel):
|
||||
"""Configuration for HybridRAG query."""
|
||||
vector_limit: int = Field(default=10, ge=1, le=50, description="Max vector results")
|
||||
graph_limit: int = Field(default=10, ge=1, le=50, description="Max graph results")
|
||||
web_limit: int = Field(default=5, ge=1, le=20, description="Max web results")
|
||||
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_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")
|
||||
|
||||
|
||||
class RelatedDossier(BaseModel):
|
||||
"""Related document metadata from graph enrichment."""
|
||||
page_id: int
|
||||
title: str
|
||||
path: str
|
||||
tag: str
|
||||
shared_entities: int
|
||||
|
||||
|
||||
class HybridRAGResult(BaseModel):
|
||||
"""Single result from HybridRAG query."""
|
||||
source_type: str = Field(..., description="Source: 'vector', 'graph', 'web'")
|
||||
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")
|
||||
rrf_score: float = Field(..., description="Reciprocal Rank Fusion score")
|
||||
final_rank: int = Field(..., description="Final rank after re-ranking")
|
||||
sources: List[str] = Field(..., description="Which sources included this result")
|
||||
related_dossiers: List[RelatedDossier] = Field(default=[], description="Related documents via shared entities")
|
||||
metadata: Dict[str, Any] = Field(default={}, description="Additional metadata")
|
||||
|
||||
|
||||
class TimingBreakdown(BaseModel):
|
||||
"""Performance timing breakdown for each phase."""
|
||||
query_enhancement_ms: float = Field(..., description="Phase 0: Keyword/synonym extraction")
|
||||
vector_ms: float = Field(..., description="Phase 1: Vector search")
|
||||
graph_ms: float = Field(..., description="Phase 1: Graph search")
|
||||
web_ms: float = Field(..., description="Phase 1: Web search")
|
||||
fusion_ms: float = Field(..., description="Phase 2: RRF fusion")
|
||||
enrichment_ms: float = Field(..., description="Phase 3: Graph enrichment")
|
||||
reranking_ms: float = Field(..., description="Phase 4: LLM re-ranking")
|
||||
persistence_ms: float = Field(..., description="Phase 6: Search persistence")
|
||||
total_ms: float = Field(..., description="Total end-to-end time")
|
||||
|
||||
|
||||
class KeywordExtraction(BaseModel):
|
||||
"""Extracted keywords and synonyms from query enhancement."""
|
||||
core_keywords: List[str] = Field(default=[], description="Primary keywords")
|
||||
entities: List[str] = Field(default=[], description="Named entities")
|
||||
synonyms: Dict[str, List[str]] = Field(default={}, description="Synonyms map")
|
||||
expansions: Dict[str, List[str]] = Field(default={}, description="Abbreviation expansions")
|
||||
|
||||
|
||||
class HybridRAGResponse(BaseModel):
|
||||
"""Response from HybridRAG query."""
|
||||
query: str = Field(..., description="Original search query")
|
||||
keywords: KeywordExtraction = Field(..., description="Extracted keywords/synonyms")
|
||||
results: List[HybridRAGResult] = Field(..., description="Ranked search results")
|
||||
context: str = Field(..., description="Formatted context for LLM consumption")
|
||||
source_counts: Dict[str, int] = Field(..., description="Result counts by source")
|
||||
total_results: int = Field(..., description="Total number of results")
|
||||
timing: TimingBreakdown = Field(..., description="Performance breakdown")
|
||||
config_used: HybridRAGConfig = Field(..., description="Configuration used")
|
||||
search_id: Optional[str] = Field(None, description="Search ID for Librarian tracking")
|
||||
|
||||
|
||||
class HybridRAGRequest(BaseModel):
|
||||
"""Request for HybridRAG query."""
|
||||
query: str = Field(..., min_length=1, max_length=500, description="Search query")
|
||||
config: Optional[HybridRAGConfig] = Field(None, description="Custom configuration")
|
||||
@@ -0,0 +1,117 @@
|
||||
"""
|
||||
HybridRAG router for multi-source search API.
|
||||
|
||||
Provides endpoint for combining vector, graph, and web search
|
||||
with RRF fusion and LLM re-ranking.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Depends, Query
|
||||
import logging
|
||||
|
||||
from src.models.hybrid_rag import HybridRAGRequest, HybridRAGResponse
|
||||
from src.services.hybrid_rag_service import HybridRAGService
|
||||
from src.services.vector_service import VectorService
|
||||
from src.services.graph_service import GraphService
|
||||
from src.clients.searxng_client import SearXNGClient
|
||||
from src.clients.ollama_client import OllamaClient
|
||||
from src.core.dependencies import (
|
||||
Neo4jDep, WikiJSDep, QdrantDep, OllamaDep,
|
||||
SearXNGDep, verify_api_key, get_settings
|
||||
)
|
||||
from src.config import Settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/query", tags=["HybridRAG"])
|
||||
|
||||
|
||||
# Dependency to get HybridRAG service
|
||||
def get_hybrid_rag_service(
|
||||
neo4j_client: Neo4jDep,
|
||||
wiki_client: WikiJSDep,
|
||||
qdrant_client: QdrantDep,
|
||||
ollama_client: OllamaDep,
|
||||
searxng_client: SearXNGDep,
|
||||
settings: Settings = Depends(get_settings)
|
||||
) -> HybridRAGService:
|
||||
"""Get HybridRAG service instance with all dependencies."""
|
||||
from src.services.vector_service import VectorService
|
||||
from src.services.graph_service import GraphService
|
||||
|
||||
# Create component services
|
||||
vector_service = VectorService(qdrant_client, wiki_client, ollama_client)
|
||||
graph_service = GraphService(neo4j_client, wiki_client)
|
||||
|
||||
# Create HybridRAG service
|
||||
return HybridRAGService(
|
||||
vector_service=vector_service,
|
||||
graph_service=graph_service,
|
||||
searxng_client=searxng_client,
|
||||
ollama_client=ollama_client,
|
||||
settings=settings
|
||||
)
|
||||
|
||||
|
||||
@router.post("/hybrid", response_model=HybridRAGResponse)
|
||||
async def hybrid_search(
|
||||
request: HybridRAGRequest,
|
||||
user: str = Query(default="jpmschweitzer", description="User identifier for multi-tenancy"),
|
||||
hybrid_rag_service: HybridRAGService = Depends(get_hybrid_rag_service),
|
||||
api_key: str = Depends(verify_api_key)
|
||||
):
|
||||
"""
|
||||
Execute HybridRAG query combining vector, graph, and web search.
|
||||
|
||||
**6-Phase Pipeline:**
|
||||
1. **Query Enhancement**: Extract keywords/synonyms with LLM
|
||||
2. **Parallel Retrieval**: Search vector (Qdrant), graph (Neo4j), web (SearXNG)
|
||||
3. **RRF Fusion**: Merge results with Reciprocal Rank Fusion
|
||||
4. **Enrichment**: Add related documents via shared entities
|
||||
5. **LLM Re-ranking**: Re-rank with mistral-nemo for relevance
|
||||
6. **Context Formatting**: Format for LLM consumption
|
||||
7. **Persistence**: Store for Librarian knowledge consolidation
|
||||
|
||||
**Example Request:**
|
||||
```json
|
||||
{
|
||||
"query": "How does Docker orchestration work with Kubernetes?",
|
||||
"user": "jpmschweitzer",
|
||||
"config": {
|
||||
"vector_limit": 10,
|
||||
"graph_limit": 10,
|
||||
"web_limit": 5,
|
||||
"enable_reranking": true,
|
||||
"final_result_count": 10
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Returns:**
|
||||
- Ranked results from all sources
|
||||
- Extracted keywords/synonyms
|
||||
- Related dossiers (via graph)
|
||||
- Formatted context for LLM
|
||||
- Performance timing breakdown
|
||||
- Search ID for Librarian tracking
|
||||
"""
|
||||
try:
|
||||
logger.info(f"HybridRAG request: '{request.query}' for user '{user}'")
|
||||
|
||||
response = await hybrid_rag_service.search(
|
||||
query=request.query,
|
||||
user=user,
|
||||
config=request.config
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"HybridRAG completed: {response.total_results} results in {response.timing.total_ms:.0f}ms"
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
except ValueError as e:
|
||||
logger.error(f"Invalid request: {e}")
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"HybridRAG search failed: {e}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail="Search failed")
|
||||
@@ -0,0 +1,739 @@
|
||||
"""
|
||||
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 mistral-nemo
|
||||
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.clients.searxng_client import SearXNGClient
|
||||
from src.clients.ollama_client import OllamaClient
|
||||
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,
|
||||
settings: Settings
|
||||
):
|
||||
"""
|
||||
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)
|
||||
settings: Application settings
|
||||
"""
|
||||
self.vector = vector_service
|
||||
self.graph = graph_service
|
||||
self.searxng = searxng_client
|
||||
self.ollama = ollama_client
|
||||
self.settings = settings
|
||||
self.reranker_model = settings.reranker_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)
|
||||
|
||||
# Phase 2: RRF Fusion
|
||||
phase2_start = time.time()
|
||||
fused_results = self._reciprocal_rank_fusion(
|
||||
results_by_source={
|
||||
"vector": raw_results.get("vector", []),
|
||||
"graph": raw_results.get("graph", []),
|
||||
"web": raw_results.get("web", [])
|
||||
},
|
||||
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. For each important word, provide synonyms and expansions.
|
||||
|
||||
Query: "{query}"
|
||||
|
||||
Return ONLY valid JSON:
|
||||
{{
|
||||
"core_keywords": ["key", "words", "from", "query"],
|
||||
"synonyms": {{
|
||||
"word": ["alternative", "terms"]
|
||||
}}
|
||||
}}
|
||||
|
||||
Example for "Docker container hosting":
|
||||
{{
|
||||
"core_keywords": ["docker", "container", "hosting"],
|
||||
"synonyms": {{
|
||||
"docker": ["containerization", "container runtime"],
|
||||
"hosting": ["server", "infrastructure"]
|
||||
}}
|
||||
}}
|
||||
|
||||
JSON:"""
|
||||
|
||||
try:
|
||||
response = await self.ollama.generate_text(
|
||||
prompt=prompt,
|
||||
model=self.reranker_model
|
||||
)
|
||||
|
||||
# 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
|
||||
)
|
||||
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:
|
||||
results = await self.graph.search_documents(
|
||||
query=query,
|
||||
user=user,
|
||||
limit=config.graph_limit,
|
||||
keywords_data=keywords_data
|
||||
)
|
||||
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
|
||||
if config.enable_web:
|
||||
async def web_search():
|
||||
start = time.time()
|
||||
try:
|
||||
results = await self.searxng.search_general(
|
||||
query=query,
|
||||
limit=config.web_limit
|
||||
)
|
||||
formatted = [
|
||||
{
|
||||
"url": r.get("url"),
|
||||
"title": r.get("title", ""),
|
||||
"content": r.get("content", ""),
|
||||
"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()
|
||||
|
||||
# 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', []))}"
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
def _reciprocal_rank_fusion(
|
||||
self,
|
||||
results_by_source: Dict[str, List],
|
||||
k: int = 60
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Phase 2: Merge results using Reciprocal Rank Fusion.
|
||||
|
||||
RRF formula: score = sum(1 / (k + rank)) for each source
|
||||
|
||||
Args:
|
||||
results_by_source: Results from each source
|
||||
k: RRF constant (default 60)
|
||||
|
||||
Returns:
|
||||
Merged and sorted results
|
||||
"""
|
||||
rrf_scores = {}
|
||||
|
||||
for source, results in results_by_source.items():
|
||||
for rank, result in enumerate(results, start=1):
|
||||
# Use page_id for wiki results, url hash for web results
|
||||
if result.get("page_id"):
|
||||
result_id = f"page_{result['page_id']}"
|
||||
elif result.get("url"):
|
||||
result_id = f"url_{hash(result['url'])}"
|
||||
else:
|
||||
continue # Skip results without ID
|
||||
|
||||
if result_id not in rrf_scores:
|
||||
rrf_scores[result_id] = {
|
||||
"result": result,
|
||||
"rrf_score": 0.0,
|
||||
"sources": [],
|
||||
"source_type": source
|
||||
}
|
||||
|
||||
# RRF formula: sum of 1/(k + rank) across sources
|
||||
rrf_scores[result_id]["rrf_score"] += 1 / (k + rank)
|
||||
rrf_scores[result_id]["sources"].append(source)
|
||||
|
||||
# If result appears in multiple sources, update source_type
|
||||
if len(rrf_scores[result_id]["sources"]) > 1:
|
||||
rrf_scores[result_id]["source_type"] = "+".join(
|
||||
sorted(set(rrf_scores[result_id]["sources"]))
|
||||
)
|
||||
|
||||
# Sort by RRF score descending
|
||||
sorted_results = sorted(
|
||||
rrf_scores.values(),
|
||||
key=lambda x: x["rrf_score"],
|
||||
reverse=True
|
||||
)
|
||||
|
||||
logger.info(f"RRF fusion: {len(sorted_results)} unique results from {len(results_by_source)} sources")
|
||||
|
||||
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"""Given this search query and documents, rank them by relevance.
|
||||
|
||||
Query: {query}
|
||||
|
||||
Documents:
|
||||
{docs_text}
|
||||
|
||||
Return only the numbers in order of relevance (most relevant first).
|
||||
Example: 3,1,5,2,4
|
||||
|
||||
Ranking:"""
|
||||
|
||||
response = await self.ollama.generate_text(
|
||||
prompt=prompt,
|
||||
model=self.reranker_model
|
||||
)
|
||||
|
||||
# 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
|
||||
Reference in New Issue
Block a user