Agentic RAG & research agents

GRADE_PROMPT = """Score retrieval for the question.
Return JSON: {"relevant": bool, "sufficient": bool, "follow_up_query": str|null}

Question: {question}
Chunks: {chunks}"""

def grade_retrieval(question: str, chunks: list[dict], llm) -> dict:
    raw = llm.complete(GRADE_PROMPT.format(question=question, chunks=chunks[:5]))
    grade = json.loads(raw)
    if grade.get("relevant") and grade.get("sufficient"):
        return {"action": "synthesize", "chunks": chunks}
    if grade.get("follow_up_query"):
        return {"action": "retrieve_again", "query": grade["follow_up_query"]}
    return {"action": "rewrite", "query": question}

MAX_RETRIEVE_STEPS = 4

DECOMPOSE_PROMPT = """Break this question into 1-4 search sub-queries needed to answer it.
Return JSON: {"sub_queries": ["...", "..."]}"""

def agentic_retrieve(question: str, retriever, llm) -> list[dict]:
    plan = json.loads(llm.complete(DECOMPOSE_PROMPT + question))
    evidence: list[dict] = []
    seen_ids: set[str] = set()

    for sub_q in plan.get("sub_queries", [question]):
        if len(evidence) >= 20:
            break
        hits = retriever.search(sub_q, k=5)
        for h in hits:
            if h["id"] not in seen_ids:
                seen_ids.add(h["id"])
                evidence.append({**h, "from_query": sub_q})

    # Optional: critic step — need more?
    critique = llm.complete(f"Evidence count {len(evidence)}. Need more retrieval for: {question}? yes/no")
    if "yes" in critique.lower() and len(plan.get("sub_queries", [])) < MAX_RETRIEVE_STEPS:
        follow_up = llm.complete(f"One follow-up search query for: {question}")
        evidence.extend(retriever.search(follow_up.strip(), k=3))
    return evidence

@dataclass
class Claim:
    text: str
    source_id: str
    quote: str

def research_agent(question: str, retriever, llm) -> dict:
    brief = llm.complete_json(CLARIFY_PROMPT, question)
    evidence = agentic_retrieve(brief["restated_question"], retriever, llm)

    claims: list[Claim] = []
    for chunk in evidence:
        extracted = llm.complete_json(EXTRACT_PROMPT, chunk["text"])
        for c in extracted.get("claims", []):
            claims.append(Claim(c["text"], chunk["id"], c["quote"]))

    synthesis = llm.complete(SYNTHESIZE_PROMPT, {
        "question": question,
        "claims": [c.__dict__ for c in claims],
    })
    return {
        "answer": synthesis,
        "citations": {c.source_id: c.quote for c in claims},
        "gaps": llm.complete(GAP_PROMPT, {"question": question, "claims": len(claims)}),
    }

@dataclass
class TopicNode:
    topic: str
    depth: int
    parent: str | None

def deep_research(root_question: str, retriever, llm, max_nodes: int = 24) -> dict:
    frontier: list[TopicNode] = [TopicNode(root_question, 0, None)]
    visited_topics: set[str] = set()
    source_pool: dict[str, dict] = {}  # id -> chunk

    while frontier and len(visited_topics) < max_nodes:
        node = frontier.pop(0)
        key = node.topic.lower().strip()
        if key in visited_topics:
            continue
        visited_topics.add(key)

        hits = retriever.search(node.topic, k=6)
        for h in dedupe_hits(hits, source_pool):
            source_pool[h["id"]] = h

        if node.depth < 3:
            subtopics = llm.complete_json(EXPAND_PROMPT, node.topic).get("subtopics", [])
            for st in subtopics[:4]:
                frontier.append(TopicNode(st, node.depth + 1, key))

    outline = llm.complete(OUTLINE_PROMPT, list(source_pool.values()))
    draft = llm.complete(REPORT_PROMPT, {"outline": outline, "sources": list(source_pool.values())})
    grounded = grounding_filter(draft, source_pool, llm)
    return {"report": grounded, "sources": list(source_pool.values()), "nodes_explored": len(visited_topics)}

from e2b_code_interpreter import Sandbox

MAX_FIX_ATTEMPTS = 5

def run_code_agent(question: str, data_csv: str, llm) -> dict:
    code = llm.complete(CODEGEN_PROMPT, {"question": question, "schema": data_csv[:500]})
    errors: list[str] = []

    with Sandbox(timeout=120) as sbx:
        sbx.files.write("/data/input.csv", data_csv)
        for attempt in range(MAX_FIX_ATTEMPTS):
            execution = sbx.run_code(code)
            if execution.error is None:
                return {
                    "answer": llm.complete(SUMMARIZE_PROMPT, execution.logs.stdout),
                    "code": code,
                    "artifacts": execution.results,
                }
            errors.append(str(execution.error))
            code = llm.complete(FIX_PROMPT, {"code": code, "errors": errors})
    return {"error": "max_attempts", "last_code": code}

@dataclass
class GoldenTask:
    id: str
    question: str
    expected_tools: set[str]
    required_source_ids: set[str]

def eval_agent(agent_fn, tasks: list[GoldenTask], judge_llm) -> dict:
    results = []
    for task in tasks:
        trace = agent_fn.run_with_trace(task.question)
        tools_used = {s["name"] for s in trace.steps if s["type"] == "tool"}
        sources = set(trace.retrieval_ids)

        tool_score = len(task.expected_tools & tools_used) / max(len(task.expected_tools), 1)
        recall = len(task.required_source_ids & sources) / max(len(task.required_source_ids), 1)
        completion = judge_llm.score(task.question, trace.answer, task.rubric)

        results.append({
            "task_id": task.id,
            "completion": completion,
            "tool_accuracy": tool_score,
            "retrieval_recall": recall,
            "steps": len(trace.steps),
        })
    return {"tasks": results, "avg_completion": mean(r["completion"] for r in results)}

def route_and_run(question: str, attachments: list, cfg) -> dict:
    route = "fast"
    if cfg.deep_research_phrase in question.lower():
        route = "deep"
    elif attachments and any(a.endswith(".csv") for a in attachments):
        route = "code"
    elif classifier.needs_agentic_rag(question):
        route = "agentic"

    log.info("route_selected", route=route, question_hash=hash(question))
    handlers = {
        "fast": lambda: single_shot_rag(question),
        "agentic": lambda: research_agent(question, retriever, llm),
        "deep": lambda: enqueue_deep_research(question),
        "code": lambda: run_code_agent(question, attachments[0], llm),
    }
    return {"route": route, "result": handlers[route]()}

def agent_loop(question: str, tools, llm, max_steps: int = 8):
    messages = [{"role": "user", "content": question}]
    queries_seen: set[str] = set()
    for step in range(max_steps):
        resp = llm.chat(messages, tools=tools)
        if not resp.tool_calls:
            return resp.content
        for call in resp.tool_calls:
            if call.name == "search_kb":
                q = json.loads(call.arguments)["query"]
                if q in queries_seen:
                    messages.append({"role": "user", "content": "Duplicate query. Synthesize now with existing evidence."})
                    continue
                queries_seen.add(q)
            messages.append(run_tool(call))
    return llm.chat(messages + [{"role": "user", "content": "Step limit reached. Answer with citations and list gaps."}]).content