| """Grounded RAG answering: retrieve top-k chunks, prompt the LLM to answer |
| ONLY from that context with [Company, p.X] citations, abstain when unsupported. |
| |
| CLI: python -m src.answer "your question" |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import re |
| from dataclasses import dataclass, field |
| from functools import lru_cache |
|
|
| from dotenv import load_dotenv |
|
|
| import config |
| from src.retrieve import RetrievedChunk, get_retriever |
|
|
| SYSTEM_PROMPT = f"""You answer questions about company annual reports (Equinor, DNB, and Norsk Hydro, 2025) using ONLY the document excerpts provided in the user message. |
| |
| Rules: |
| 1. Ground every statement in the excerpts. Use no outside knowledge, no estimates, no general knowledge about these companies. |
| 2. Cite the source of every factual claim inline as [Company, p.X], where Company and X (the page number) come from the header line of the excerpt you used. Cite only excerpts you actually relied on. |
| 3. If the excerpts do not contain the information needed to answer, start your reply with exactly this sentence: {config.ABSTAIN_STRING} |
| Then add one or two short sentences that help the reader: say what related information the excerpts DO contain (for example a different year, a target instead of an actual figure, an average instead of a closing value, or a related measure), citing it as in rule 2. If nothing in the excerpts is related, say so plainly. Never state a figure for the thing that was asked if it is not in the excerpts. |
| 4. If the excerpts only partially answer the question, answer the supported part and say what is not stated. |
| 5. Be concise: a direct answer in one to three sentences, with figures quoted exactly as written in the excerpts (keep units and currency). |
| """ |
|
|
| |
| CITATION_RE = re.compile(r"\[([^\[\],]+?),\s*p\.\s*(\d+)\]") |
|
|
|
|
| @dataclass |
| class Citation: |
| """A single (company, page) reference parsed from the answer text.""" |
|
|
| company: str |
| page: int |
|
|
|
|
| @dataclass |
| class AnswerResult: |
| """Everything the UI and the eval harness need about one QA turn.""" |
|
|
| question: str |
| answer: str |
| citations: list[Citation] = field(default_factory=list) |
| chunks: list[RetrievedChunk] = field(default_factory=list) |
| abstained: bool = False |
|
|
|
|
| def parse_citations(text: str) -> list[Citation]: |
| """Extract [Company, p.X] citations, deduplicated, in order of appearance.""" |
| seen: set[tuple[str, int]] = set() |
| citations: list[Citation] = [] |
| for company, page in CITATION_RE.findall(text): |
| key = (company.strip(), int(page)) |
| if key not in seen: |
| seen.add(key) |
| citations.append(Citation(company=key[0], page=key[1])) |
| return citations |
|
|
|
|
| def build_context(chunks: list[RetrievedChunk]) -> str: |
| """Format retrieved chunks as clearly-delimited, citable excerpts.""" |
| blocks: list[str] = [] |
| for i, c in enumerate(chunks, 1): |
| blocks.append(f"--- Excerpt {i} [{c.company}, p.{c.page}] ---\n{c.text}") |
| return "\n\n".join(blocks) |
|
|
|
|
| @lru_cache(maxsize=1) |
| def get_client(): |
| """Anthropic client; reads ANTHROPIC_API_KEY from .env / environment.""" |
| import anthropic |
|
|
| load_dotenv() |
| return anthropic.Anthropic() |
|
|
|
|
| def answer(question: str, top_k: int = config.TOP_K) -> AnswerResult: |
| """Retrieve context and produce a grounded, cited answer (or abstain).""" |
| chunks = get_retriever().search(question, top_k=top_k) |
| context = build_context(chunks) |
|
|
| response = get_client().messages.create( |
| model=config.ANSWER_MODEL, |
| max_tokens=config.MAX_ANSWER_TOKENS, |
| temperature=0.0, |
| system=SYSTEM_PROMPT, |
| messages=[ |
| { |
| "role": "user", |
| "content": f"{context}\n\nQuestion: {question}", |
| } |
| ], |
| ) |
| text = "".join(b.text for b in response.content if b.type == "text").strip() |
| |
| |
| abstained = text.startswith(config.ABSTAIN_STRING) |
|
|
| return AnswerResult( |
| question=question, |
| answer=text, |
| |
| citations=parse_citations(text), |
| chunks=chunks, |
| abstained=abstained, |
| ) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Ask a grounded question.") |
| parser.add_argument("question") |
| parser.add_argument("-k", type=int, default=config.TOP_K) |
| args = parser.parse_args() |
|
|
| result = answer(args.question, top_k=args.k) |
| print(f"Q: {result.question}\n") |
| print(f"A: {result.answer}\n") |
| if result.abstained: |
| print("(abstained)") |
| else: |
| cites = ", ".join(f"[{c.company}, p.{c.page}]" for c in result.citations) |
| print(f"Citations: {cites or '(none found in answer)'}") |
| print("\nRetrieved:") |
| for i, c in enumerate(result.chunks, 1): |
| print(f" #{i} score={c.score:.3f} [{c.company}, p.{c.page}]") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
