#!/usr/bin/env python
"""
Multi-hop retrieval agent for CognitiveRAG.
Coordinates query decomposition, iterative retrieval, evidence validation,
and final answer synthesis with provenance-aware citations.
"""
from __future__ import annotations
import json
import os
import re
import time
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple
import boto3
from langsmith import traceable
try:
from retriever_multi import get_unified_context_for_llm
except ImportError:
import sys
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "04_vectoredb"))
from retriever_multi import get_unified_context_for_llm
@dataclass
class MultiHopConfig:
max_hops: int = 3
per_hop_top_k: int = 8
total_doc_budget: int = 24
sufficiency_threshold: float = 0.75
time_limit_sec: int = 60
strategy: str = "hybrid"
include_related: bool = False
memory_top_k: int = 3
@dataclass
class EvidenceItem:
chunk_id: str
doc_id: str
text: str
source: str
modality: str
section: str
page: Optional[int]
score: float
location: str
def provenance_key(self) -> str:
if self.chunk_id:
return self.chunk_id
base = f"{self.doc_id}|{self.page}|{self.section}|{self.location}"
return f"{base}|{hash(self.text)}"
@dataclass
class HopRecord:
hop_index: int
sub_question: str
retrieval_strategy: str
retrieved_count: int
new_evidence_count: int
total_evidence_count: int
validation: Dict[str, Any]
retrieval_metadata: List[Dict[str, Any]]
started_at: str
completed_at: str
@dataclass
class MultiHopState:
query: str
started_at: str
hops: List[HopRecord] = field(default_factory=list)
evidence_pool: List[EvidenceItem] = field(default_factory=list)
termination_reason: str = ""
status: str = "running"
completed_at: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
return {
"query": self.query,
"started_at": self.started_at,
"completed_at": self.completed_at,
"status": self.status,
"termination_reason": self.termination_reason,
"evidence_pool": [item.__dict__ for item in self.evidence_pool],
"hops": [
{
"hop_index": hop.hop_index,
"sub_question": hop.sub_question,
"retrieval_strategy": hop.retrieval_strategy,
"retrieved_count": hop.retrieved_count,
"new_evidence_count": hop.new_evidence_count,
"total_evidence_count": hop.total_evidence_count,
"validation": hop.validation,
"retrieval_metadata": hop.retrieval_metadata,
"started_at": hop.started_at,
"completed_at": hop.completed_at,
}
for hop in self.hops
],
}
@dataclass
class MultiHopResult:
answer: str
citations: List[Dict[str, Any]]
confidence: float
status: str
termination_reason: str
trace_path: str
reasoning_trace: str = ""
evidence: List[Dict[str, Any]] = field(default_factory=list)
class MultiHopAgent:
def __init__(
self,
embed_fn,
config: MultiHopConfig,
bedrock_region: str = "ca-central-1",
output_dir: Optional[str] = None,
) -> None:
self.embed_fn = embed_fn
self.config = config
self.bedrock = boto3.client(
service_name="bedrock-runtime",
region_name=bedrock_region,
)
self.output_dir = output_dir or os.path.join(
os.path.dirname(__file__), "output", "multihop_traces"
)
os.makedirs(self.output_dir, exist_ok=True)
@traceable(name="multihop_decompose")
def decompose(self, query: str, evidence_summary: str) -> Dict[str, Any]:
system_prompt = (
"You are a query decomposition agent. "
"Respond with JSON only."
)
user_prompt = f"""
Query:
{query}
Evidence Summary:
{evidence_summary}
If the query can be answered with current evidence, return:
{{"action":"answer","sub_question":""}}
If more evidence is needed, return:
{{"action":"ask","sub_question":"<focused sub-question>"}}
"""
response = self._invoke_claude(system_prompt, user_prompt)
payload = self._parse_json_response(response)
if not payload:
return {"action": "ask", "sub_question": query}
action = payload.get("action", "ask")
sub_question = payload.get("sub_question", "").strip()
if action == "answer":
return {"action": "answer", "sub_question": ""}
return {"action": "ask", "sub_question": sub_question or query}
@traceable(name="multihop_retrieve")
def retrieve(self, sub_question: str) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
query_embedding = None
if self.config.strategy in ("dense", "hybrid") or self.config.memory_top_k > 0:
query_embedding = self.embed_fn(sub_question)
context = get_unified_context_for_llm(
query_embedding=query_embedding,
query_text=sub_question,
strategy=self.config.strategy,
doc_top_k=self.config.per_hop_top_k,
memory_top_k=self.config.memory_top_k,
include_related=self.config.include_related,
)
doc_chunks = context.get("document_chunks", [])
memory_chunks = context.get("memory_chunks", [])
return doc_chunks + memory_chunks, context
def _dedupe_evidence(
self, existing: Dict[str, EvidenceItem], candidates: List[EvidenceItem]
) -> Tuple[List[EvidenceItem], int]:
new_items = []
for item in candidates:
key = item.provenance_key()
if key not in existing:
existing[key] = item
new_items.append(item)
return new_items, len(new_items)
def _build_evidence_item(self, chunk: Dict[str, Any]) -> EvidenceItem:
doc_id = str(chunk.get("doc_id") or chunk.get("session_id") or "")
chunk_id = str(chunk.get("chunk_id") or "")
modality = str(chunk.get("modality") or "text")
section = str(chunk.get("section") or "")
page = chunk.get("page")
source = str(chunk.get("source") or "document")
score = float(
chunk.get("cross_encoder_score")
or chunk.get("score")
or chunk.get("rrf_score")
or 0.0
)
location = ""
if page is not None:
location = f"page:{page}"
elif section:
location = f"section:{section}"
else:
location = "unknown"
return EvidenceItem(
chunk_id=chunk_id,
doc_id=doc_id,
text=str(chunk.get("text") or ""),
source=source,
modality=modality,
section=section,
page=page,
score=score,
location=location,
)
def _summarize_evidence(self, evidence: List[EvidenceItem], max_chars: int = 1200) -> str:
if not evidence:
return "No evidence collected yet."
lines = []
for item in evidence[:6]:
snippet = item.text.replace("\n", " ").strip()[:180]
lines.append(
f"- [{item.doc_id or 'unknown'}|{item.location}] {snippet}"
)
summary = "\n".join(lines)
if len(summary) > max_chars:
return summary[:max_chars] + "..."
return summary
def _heuristic_validation(self, query: str, evidence: List[EvidenceItem]) -> Dict[str, float]:
if not evidence:
return {"keyword_coverage": 0.0, "doc_diversity": 0.0, "heuristic_score": 0.0}
query_tokens = [
token
for token in re.findall(r"[A-Za-z0-9]{4,}", query.lower())
]
if not query_tokens:
return {"keyword_coverage": 0.0, "doc_diversity": 0.0, "heuristic_score": 0.0}
evidence_text = " ".join(item.text.lower() for item in evidence)
matches = sum(1 for token in set(query_tokens) if token in evidence_text)
keyword_coverage = matches / max(len(set(query_tokens)), 1)
unique_docs = len(set(item.doc_id for item in evidence if item.doc_id))
doc_diversity = min(unique_docs / 3.0, 1.0)
heuristic_score = round((0.6 * keyword_coverage) + (0.4 * doc_diversity), 4)
return {
"keyword_coverage": round(keyword_coverage, 4),
"doc_diversity": round(doc_diversity, 4),
"heuristic_score": heuristic_score,
}
@traceable(name="multihop_validate")
def validate(self, query: str, evidence: List[EvidenceItem]) -> Dict[str, Any]:
heuristic = self._heuristic_validation(query, evidence)
summary = self._summarize_evidence(evidence)
system_prompt = "You are a strict evidence sufficiency critic. Respond with JSON only."
user_prompt = f"""
Question:
{query}
Evidence Summary:
{summary}
Return JSON with:
{{"sufficiency_score":0-1,"relevance_score":0-1,"missing_info":"...","should_stop":true/false}}
"""
response = self._invoke_claude(system_prompt, user_prompt)
payload = self._parse_json_response(response) or {}
sufficiency = float(payload.get("sufficiency_score", 0.0))
relevance = float(payload.get("relevance_score", 0.0))
should_stop = bool(payload.get("should_stop", False))
combined = (0.55 * sufficiency) + (0.25 * relevance) + (0.20 * heuristic["heuristic_score"])
return {
"heuristic": heuristic,
"sufficiency_score": round(sufficiency, 4),
"relevance_score": round(relevance, 4),
"combined_score": round(combined, 4),
"missing_info": payload.get("missing_info", ""),
"should_stop": should_stop,
}
@traceable(name="multihop_synthesize")
def synthesize(self, query: str, evidence: List[EvidenceItem]) -> Dict[str, Any]:
evidence_payload = [
{
"chunk_id": item.chunk_id,
"doc_id": item.doc_id,
"source": item.source,
"location": item.location,
"text": item.text,
}
for item in evidence
]
system_prompt = (
"You are a grounded QA system. "
"Use only the provided evidence. "
"Respond with JSON only."
)
user_prompt = f"""
Question:
{query}
Evidence:
{json.dumps(evidence_payload[:12], ensure_ascii=True, indent=2)}
Return JSON:
{{"answer":"...","citations":[{{"chunk_id":"","doc_id":"","location":""}}],"confidence":0-1,"reasoning_trace":"short"}}
If evidence is insufficient, say so directly in the answer and set confidence <= 0.3.
"""
response = self._invoke_claude(system_prompt, user_prompt)
payload = self._parse_json_response(response) or {}
return {
"answer": payload.get("answer", "The provided evidence is insufficient."),
"citations": payload.get("citations", []),
"confidence": float(payload.get("confidence", 0.0) or 0.0),
"reasoning_trace": payload.get("reasoning_trace", ""),
}
def _invoke_claude(self, system_prompt: str, user_prompt: str) -> str:
body = json.dumps(
{
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": 500,
"temperature": 0,
"system": system_prompt,
"messages": [{"role": "user", "content": user_prompt}],
}
)
response = self.bedrock.invoke_model(
modelId="anthropic.claude-3-haiku-20240307-v1:0",
body=body,
contentType="application/json",
accept="application/json",
)
response_body = json.loads(response["body"].read())
return response_body["content"][0]["text"].strip()
def _parse_json_response(self, text: str) -> Optional[Dict[str, Any]]:
try:
return json.loads(text)
except json.JSONDecodeError:
match = re.search(r"\{.*\}", text, re.DOTALL)
if not match:
return None
try:
return json.loads(match.group(0))
except json.JSONDecodeError:
return None
def _write_trace(self, state: MultiHopState, final: Dict[str, Any]) -> str:
trace_id = uuid.uuid4().hex[:8]
timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
path = os.path.join(self.output_dir, f"multihop_trace_{timestamp}_{trace_id}.json")
payload = {
"state": state.to_dict(),
"final": final,
}
with open(path, "w", encoding="utf-8") as handle:
json.dump(payload, handle, indent=2, ensure_ascii=True)
return path
@traceable(name="multihop_run")
def run(self, query: str) -> MultiHopResult:
start_time = time.monotonic()
state = MultiHopState(
query=query,
started_at=datetime.utcnow().isoformat(timespec="seconds") + "Z",
)
evidence_index: Dict[str, EvidenceItem] = {}
retrieved_total = 0
for hop in range(1, self.config.max_hops + 1):
hop_start = datetime.utcnow().isoformat(timespec="seconds") + "Z"
elapsed = time.monotonic() - start_time
if elapsed > self.config.time_limit_sec:
state.status = "terminated"
state.termination_reason = "time_limit_exceeded"
break
evidence_summary = self._summarize_evidence(state.evidence_pool)
decomposition = self.decompose(query, evidence_summary)
if decomposition["action"] == "answer":
state.status = "terminated"
state.termination_reason = "decomposer_complete"
break
sub_question = decomposition["sub_question"]
chunks, context = self.retrieve(sub_question)
retrieved_total += len(chunks)
candidates = [self._build_evidence_item(chunk) for chunk in chunks]
new_items, new_count = self._dedupe_evidence(evidence_index, candidates)
state.evidence_pool.extend(new_items)
validation = self.validate(query, state.evidence_pool)
should_stop = (
validation["combined_score"] >= self.config.sufficiency_threshold
or validation.get("should_stop", False)
)
hop_record = HopRecord(
hop_index=hop,
sub_question=sub_question,
retrieval_strategy=context.get("strategy", self.config.strategy),
retrieved_count=len(chunks),
new_evidence_count=new_count,
total_evidence_count=len(state.evidence_pool),
validation=validation,
retrieval_metadata=[
{
"chunk_id": c.get("chunk_id", ""),
"doc_id": c.get("doc_id", ""),
"section": c.get("section", ""),
"page": c.get("page"),
"source": c.get("source", ""),
"score": c.get("score", c.get("rrf_score", 0)),
}
for c in chunks
],
started_at=hop_start,
completed_at=datetime.utcnow().isoformat(timespec="seconds") + "Z",
)
state.hops.append(hop_record)
if retrieved_total >= self.config.total_doc_budget:
state.status = "terminated"
state.termination_reason = "retrieval_budget_exceeded"
break
if should_stop:
state.status = "terminated"
state.termination_reason = "evidence_sufficient"
break
if state.status != "terminated":
state.status = "terminated"
if not state.termination_reason:
state.termination_reason = "max_hops_reached"
synthesis = self.synthesize(query, state.evidence_pool)
state.completed_at = datetime.utcnow().isoformat(timespec="seconds") + "Z"
trace_path = self._write_trace(state, synthesis)
evidence_payload = [
{
"chunk_id": item.chunk_id,
"doc_id": item.doc_id,
"text": item.text,
"source": item.source,
"location": item.location,
}
for item in state.evidence_pool
]
return MultiHopResult(
answer=synthesis["answer"],
citations=synthesis.get("citations", []),
confidence=synthesis.get("confidence", 0.0),
reasoning_trace=synthesis.get("reasoning_trace", ""),
status=state.status,
termination_reason=state.termination_reason,
trace_path=trace_path,
evidence=evidence_payload,
)