src/agents/qa_agent.py

"""Retrieval (document expert) agent.

Answers factual questions strictly from retrieved evidence, with page-level
citations. Default model: Gemini (long context, strong retrieval grounding).
"""

from __future__ import annotations

from langchain_core.messages import HumanMessage, SystemMessage

from src.llm import get_llm
from src.retrieval.hybrid import RetrievedChunk

SYSTEM = """You are a financial document expert. Answer the question using ONLY the
evidence excerpts provided. Rules:
- Quote figures exactly as they appear; never estimate or infer numbers.
- After each factual claim, cite its source in square brackets, e.g. [annual_report_2024 p.38].
- If the evidence does not contain the answer, say so explicitly — do not guess.
- Keep the answer concise and factual; interpretation is another agent's job."""


def format_evidence(retrieved: list[RetrievedChunk]) -> str:
    blocks = []
    for i, r in enumerate(retrieved, 1):
        tag = " (table)" if r.chunk.is_table else ""
        blocks.append(f"[{i}] Source: {r.chunk.citation}{tag}\n{r.chunk.text}")
    return "\n\n---\n\n".join(blocks)


def answer(question: str, retrieved: list[RetrievedChunk],
           history: str = "") -> str:
    llm = get_llm("retrieval")
    context = format_evidence(retrieved) if retrieved else "(no evidence retrieved)"
    prompt = ""
    if history:
        prompt += f"Conversation so far (for pronoun/entity resolution only):\n{history}\n\n"
    prompt += f"Evidence excerpts:\n\n{context}\n\nQuestion: {question}"
    resp = llm.invoke([SystemMessage(content=SYSTEM), HumanMessage(content=prompt)])
    return resp.content