Financial RAG Accuracy Comes From Tools, Not More Retrieval: A Calculator, A Token Budget And Evidence Scoring, In Code

"""Financial QA with a calculator tool, against a documents-only baseline.

Python 3.12+, anthropic==1.13.0. The SDK reads ANTHROPIC_API_KEY from the environment.
The model picks evidence spans and writes the formula. The server binds each variable
to a span's parsed value at full precision, so no figure is ever re-typed by the model
on its way into the arithmetic (the failure FinSkillBench saw with large numeric inputs).
"""
from __future__ import annotations

import ast
import json
import operator
import os
from dataclasses import dataclass, field
from decimal import Decimal, getcontext
from typing import Literal

import anthropic

getcontext().prec = 34
MODEL = os.environ.get("FINQA_MODEL", "claude-opus-5-5")
client = anthropic.Anthropic()

Mode = Literal["tools", "documents_only"]


@dataclass(frozen=True)
class Span:
    span_id: str                 # e.g. "ACME-AR-FY2025:p61:t4:r7"
    doc_id: str                  # the filing, RNS or transcript it came from
    text: str                    # the row or sentence exactly as the model sees it
    value: Decimal | None = None # parsed once at ingestion, never by the model
    unit: str | None = None      # "GBP m", "%", "pence"
    period: str | None = None    # "FY2025", "H1 2026"


@dataclass
class Answer:
    question: str
    mode: Mode
    text: str = ""
    value: Decimal | None = None
    unit: str | None = None
    cited_span_ids: list[str] = field(default_factory=list)
    calculation: dict | None = None
    computed_by: Literal["calculator", "quoted", "model", "none"] = "none"
    input_tokens: int = 0
    output_tokens: int = 0
    served_by: str = MODEL


# The calculator: an AST whitelist over Decimal. No names except bound spans, no attributes.
_BIN = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul,
        ast.Div: operator.truediv, ast.Pow: operator.pow}
_FUNCS = {"abs": abs, "min": min, "max": max}


def safe_eval(expr: str, env: dict[str, Decimal]) -> Decimal:
    if len(expr) > 300:
        raise ValueError("expression longer than 300 characters")

    def ev(node: ast.AST) -> Decimal:
        match node:
            case ast.Expression(body=body):
                return ev(body)
            case ast.Constant(value=v) if isinstance(v, (int, float)) and not isinstance(v, bool):
                return Decimal(str(v))
            case ast.Name(id=name):
                if name not in env:
                    raise ValueError(f"unbound variable {name!r}: bind it to a span_id")
                return env[name]
            case ast.UnaryOp(op=ast.USub(), operand=operand):
                return -ev(operand)
            case ast.BinOp(left=left, op=op, right=right) if type(op) in _BIN:
                lv, rv = ev(left), ev(right)
                if isinstance(op, ast.Pow) and (rv != rv.to_integral_value() or abs(rv) > 12):
                    raise ValueError("exponent must be a whole number between -12 and 12")
                return _BIN[type(op)](lv, rv)
            case ast.Call(func=ast.Name(id=fn), args=args, keywords=[]) if fn in _FUNCS:
                return _FUNCS[fn](*(ev(a) for a in args))
        raise ValueError(f"operation not allowed: {type(node).__name__}")

    return ev(ast.parse(expr, mode="eval"))


NULLABLE = {"anyOf": [{"type": "string"}, {"type": "null"}]}
CALCULATE = {
    "name": "calculate",
    "description": ("Evaluate arithmetic such as (rev_25 - rev_24) / rev_24 * 100. Bind every "
                    "variable to a span_id; the server substitutes that span's exact value."),
    "strict": True,
    "input_schema": {
        "type": "object",
        "properties": {
            "expression": {"type": "string"},
            "bindings": {"type": "array", "items": {
                "type": "object",
                "properties": {"name": {"type": "string"}, "span_id": {"type": "string"}},
                "required": ["name", "span_id"], "additionalProperties": False}},
            "result_unit": {"type": "string"},
        },
        "required": ["expression", "bindings", "result_unit"],
        "additionalProperties": False,
    },
}
SUBMIT = {
    "name": "submit_answer",
    "description": ("Submit the final answer once. cited_span_ids lists every span the answer "
                    "relies on. value is the headline figure as a plain decimal string, or null."),
    "strict": True,
    "input_schema": {
        "type": "object",
        "properties": {
            "answer_text": {"type": "string"},
            "value": NULLABLE,
            "unit": NULLABLE,
            "cited_span_ids": {"type": "array", "items": {"type": "string"}},
            "calculation_id": NULLABLE,
        },
        "required": ["answer_text", "value", "unit", "cited_span_ids", "calculation_id"],
        "additionalProperties": False,
    },
}
SYSTEM = ("You answer questions about company filings and call transcripts using only the "
          "evidence spans supplied, each prefixed with its [span_id]. Cite every span you rely "
          "on. {rule} Finish by calling submit_answer. If the evidence does not contain the "
          "answer, submit value null and say what is missing. Describe figures only; never "
          "give investment advice or opinions on securities.")
RULES = {
    "tools": "Never do arithmetic yourself: every derived figure must come from calculate.",
    "documents_only": "You have no calculator; show your working in answer_text.",
}


def build_request(question: str, spans: list[Span], mode: Mode) -> dict:
    """Shared by answer_question and by the token budgeter, so counts match real calls."""
    evidence = "\n".join(f"[{s.span_id}] ({s.doc_id}, {s.period or 'n/a'}) {s.text}" for s in spans)
    return {
        "system": SYSTEM.format(rule=RULES[mode]),
        "tools": [CALCULATE, SUBMIT] if mode == "tools" else [SUBMIT],
        "messages": [{"role": "user", "content": f"Evidence:\n{evidence}\n\nQuestion: {question}"}],
    }


def _run_calculation(inp: dict, by_id: dict[str, Span], calcs: dict[str, dict]) -> str:
    env, units = {}, set()
    for b in inp["bindings"]:
        span = by_id.get(b["span_id"])
        if span is None or span.value is None:
            raise ValueError(f"{b['span_id']} is not a numeric span in the evidence")
        env[b["name"]] = span.value
        units.add(span.unit)
    tree = ast.parse(inp["expression"], mode="eval")
    adds = any(isinstance(n, ast.BinOp) and isinstance(n.op, (ast.Add, ast.Sub)) for n in ast.walk(tree))
    if adds and len(units) > 1:
        raise ValueError(f"adding or subtracting mixed units {sorted(map(str, units))}: convert first")
    result = safe_eval(inp["expression"], env)
    calc_id = f"calc-{len(calcs) + 1}"
    calcs[calc_id] = {"id": calc_id, "expression": inp["expression"], "result": str(result),
                      "unit": inp["result_unit"],
                      "bindings": {b["name"]: b["span_id"] for b in inp["bindings"]}}
    return json.dumps({"calculation_id": calc_id, "result": str(result), "unit": inp["result_unit"]})


def _accept(inp: dict, by_id: dict[str, Span], calcs: dict[str, dict], ans: Answer) -> None:
    cited = inp["cited_span_ids"]
    if unknown := [s for s in cited if s not in by_id]:
        raise ValueError(f"cited spans not in the evidence: {unknown}")
    value = Decimal(inp["value"]) if inp["value"] is not None else None
    if value is not None and not value.is_finite():
        raise ValueError("value must be a finite decimal")
    if value is not None and not cited:
        raise ValueError("a figure with no cited evidence is not an answer")
    calc = calcs.get(inp["calculation_id"] or "")
    if inp["calculation_id"] and calc is None:
        raise ValueError(f"unknown calculation_id {inp['calculation_id']}")
    if value is None:
        computed_by = "none"
    elif calc is not None:
        if not set(calc["bindings"].values()) <= set(cited):
            raise ValueError("cite every span the calculation used")
        step = Decimal(1).scaleb(value.as_tuple().exponent)       # the precision the model reported
        if Decimal(calc["result"]).quantize(step) != value:
            raise ValueError(f"value {value} does not match {calc['id']} = {calc['result']}")
        computed_by = "calculator"
    elif any(by_id[s].value == value for s in cited):
        computed_by = "quoted"
    elif ans.mode == "tools":
        raise ValueError("value is neither quoted from a cited span nor produced by calculate")
    else:
        computed_by = "model"                                    # baseline: unverified arithmetic
    ans.text, ans.value, ans.unit = inp["answer_text"], value, inp["unit"]
    ans.cited_span_ids, ans.calculation, ans.computed_by = cited, calc, computed_by


def answer_question(question: str, spans: list[Span], mode: Mode = "tools", max_turns: int = 6) -> Answer:
    by_id = {s.span_id: s for s in spans}
    calcs: dict[str, dict] = {}
    ans = Answer(question=question, mode=mode)
    req = build_request(question, spans, mode)
    messages = req["messages"]
    for _ in range(max_turns):
        resp = client.beta.messages.create(
            model=MODEL, max_tokens=16000, system=req["system"], tools=req["tools"],
            messages=messages, output_config={"effort": "medium"},
            betas=["server-side-fallback-2026-07-01"], fallbacks="default",
        )
        ans.input_tokens += resp.usage.input_tokens
        ans.output_tokens += resp.usage.output_tokens
        ans.served_by = resp.model
        if resp.stop_reason == "refusal":
            ans.text = "declined by the model; route to a person"
            return ans
        messages.append({"role": "assistant", "content": resp.content})
        tool_uses = [b for b in resp.content if b.type == "tool_use"]
        if not tool_uses:
            messages.append({"role": "user", "content": "Finish by calling submit_answer."})
            continue
        results = []
        for tu in tool_uses:
            try:
                if tu.name == "calculate":
                    results.append({"type": "tool_result", "tool_use_id": tu.id,
                                    "content": _run_calculation(tu.input, by_id, calcs)})
                else:
                    _accept(tu.input, by_id, calcs, ans)
                    return ans
            except (ValueError, ArithmeticError, SyntaxError) as exc:
                results.append({"type": "tool_result", "tool_use_id": tu.id,
                                "content": f"rejected: {exc}", "is_error": True})
        messages.append({"role": "user", "content": results})
    ans.text = f"no accepted answer within {max_turns} turns"
    return ans

"""Top-1 first, then escalate to top-5 by predicted gain per token under a batch budget.

Python 3.12+, anthropic==1.13.0, scikit-learn==1.9.1, numpy 2.x.
Same shape as arXiv 2610.07760: everyone gets the cheap route; the budget, set as a
fraction of what escalating the whole batch would cost, buys deeper context only where
predicted gain per additional token is highest. Every question's tokens are logged.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split

from fin_qa.service import MODEL, Answer, Span, answer_question, build_request, client

COMPARATIVE = re.compile(r"\b(change|growth|increase|decrease|ratio|margin|versus|compared|between|per share)\b", re.I)
PERIODS = re.compile(r"\b(?:FY|H[12] |Q[1-4] )?20\d\d\b")


class PageIndex:
    """TF-IDF page retrieval, the cheap first route used in the paper."""

    def __init__(self, pages: list[Span]):
        self.pages = pages
        self.vectoriser = TfidfVectorizer(ngram_range=(1, 2), sublinear_tf=True)
        self.matrix = self.vectoriser.fit_transform([p.text for p in pages])

    def search(self, question: str, k: int) -> tuple[list[Span], np.ndarray]:
        scores = (self.matrix @ self.vectoriser.transform([question]).T).toarray().ravel()
        order = np.argsort(-scores)[:k]
        return [self.pages[i] for i in order], scores[order]


def features(question: str, top1: Answer, scores: np.ndarray) -> list[float]:
    """Signals available after the top-1 answer. Expect them to be weak: the paper's AUROC was 0.60."""
    return [
        float(scores[0]),
        float(scores[0] - scores[1]) if len(scores) > 1 else float(scores[0]),
        1.0 if top1.value is None else 0.0,                  # top-1 abstained
        float(len(set(PERIODS.findall(question)))),          # multi-period questions need more pages
        1.0 if COMPARATIVE.search(question) else 0.0,
        1.0 if top1.computed_by == "calculator" else 0.0,
    ]


def fit_gain_model(history_jsonl: Path) -> LogisticRegression:
    """Rows from past runs: features plus 'helped' (did top-5 beat top-1 under adjudication?)."""
    rows = [json.loads(line) for line in history_jsonl.read_text().splitlines() if line.strip()]
    X = np.array([r["features"] for r in rows])
    y = np.array([int(r["helped"]) for r in rows])
    X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.3, stratify=y, random_state=7)
    model = LogisticRegression(max_iter=1000, class_weight="balanced").fit(X_tr, y_tr)
    auroc = roc_auc_score(y_te, model.predict_proba(X_te)[:, 1])
    print(f"gain model held-out AUROC {auroc:.2f} (0.50 is a coin toss)")
    return model.fit(X, y)


def prompt_tokens(question: str, spans: list[Span]) -> int:
    """Token counting is a separate, free, rate-limited endpoint; the count is an estimate."""
    req = build_request(question, spans, "tools")
    return client.messages.count_tokens(model=MODEL, **req).input_tokens


@dataclass
class Candidate:
    qid: str
    question: str
    top1: Answer
    p_gain: float
    predicted_cost: int          # predicted tokens for the extra top-5 call
    features: list[float]

    @property
    def gain_per_token(self) -> float:
        return self.p_gain / max(self.predicted_cost, 1)


def run_batch(questions: list[dict], index: PageIndex, gain_model: LogisticRegression,
              budget_fraction: float = 0.10, log_path: Path = Path("logs/qa_tokens.jsonl")) -> dict:
    cands: list[Candidate] = []
    for q in questions:
        pages, scores = index.search(q["question"], 5)
        top1 = answer_question(q["question"], pages[:1])
        # Cost is the predictable half: scale the observed top-1 spend by the prompt-size ratio.
        ratio = prompt_tokens(q["question"], pages) / max(prompt_tokens(q["question"], pages[:1]), 1)
        x = features(q["question"], top1, scores)
        cands.append(Candidate(
            qid=q["id"], question=q["question"], top1=top1, features=x,
            p_gain=float(gain_model.predict_proba([x])[0, 1]),
            predicted_cost=int((top1.input_tokens + top1.output_tokens) * ratio)))

    # Budget = fraction of the predicted cost of escalating everyone; greedy by gain per token.
    budget = budget_fraction * sum(c.predicted_cost for c in cands)
    chosen, committed = set(), 0
    for c in sorted(cands, key=lambda c: c.gain_per_token, reverse=True):
        if committed + c.predicted_cost <= budget:
            chosen.add(c.qid)
            committed += c.predicted_cost

    log_path.parent.mkdir(parents=True, exist_ok=True)
    total = 0
    with log_path.open("a") as log:
        for c in cands:
            final, extra = c.top1, 0
            if c.qid in chosen:
                pages, _ = index.search(c.question, 5)
                final = answer_question(c.question, pages)
                extra = final.input_tokens + final.output_tokens
            top1_tokens = c.top1.input_tokens + c.top1.output_tokens
            total += top1_tokens + extra
            log.write(json.dumps({
                "qid": c.qid, "route": "top5" if c.qid in chosen else "top1",
                "top1_tokens": top1_tokens, "escalation_tokens": extra,
                "predicted_escalation_tokens": c.predicted_cost, "p_gain": round(c.p_gain, 4),
                "gain_per_1k_tokens": round(c.gain_per_token * 1000, 5), "features": c.features,
                "value": None if final.value is None else str(final.value), "unit": final.unit,
                "cited": final.cited_span_ids, "computed_by": final.computed_by,
                "served_by": final.served_by,
            }) + "\n")
    n = len(cands)
    return {"questions": n, "escalated": len(chosen), "budget_tokens": int(budget),
            "tokens_per_question": round(total / max(n, 1)),
            "predicted_top5_for_all_per_question": round(sum(c.predicted_cost for c in cands) / max(n, 1))}

"""Evidence-level scoring, plus an always-on refresher that re-runs only dependent claims.

Python 3.12+, standard library only (uses fin_qa.service for Span and Answer).
Scoring follows FinFIRST's split: information acquisition, source verification and
computation, scored as atomic items next to the final answer, so a right number on
missing evidence is counted as what it is. The refresher never publishes on its own.
"""
from __future__ import annotations

import hashlib
import json
import os
import time
import uuid
from collections.abc import Callable
from dataclasses import dataclass, field
from decimal import Decimal
from pathlib import Path

from fin_qa.service import Answer, Span

# ---------- 1. Evidence-level evaluation ----------

@dataclass(frozen=True)
class Reference:
    task_id: str
    question: str
    required_span_ids: frozenset[str]      # what an analyst would have to find
    authoritative_docs: frozenset[str]     # e.g. the annual report, not a news summary of it
    periods: frozenset[str]                # every period the answer legitimately needs
    expected: Decimal
    unit: str
    tolerance: Decimal                     # absolute, in the reference unit


def score(ref: Reference, ans: Answer, retrieved: list[Span]) -> dict[str, bool]:
    by_id = {s.span_id: s for s in retrieved}
    cited = [by_id[i] for i in ans.cited_span_ids if i in by_id]
    calc_ok = ans.computed_by == "quoted" or (
        ans.computed_by == "calculator" and ans.calculation is not None
        and set(ans.calculation["bindings"].values()) <= set(ans.cited_span_ids))
    return {
        "acquisition.required_spans_retrieved": ref.required_span_ids <= by_id.keys(),
        "verification.required_spans_cited": ref.required_span_ids <= set(ans.cited_span_ids),
        "verification.sources_authoritative": bool(cited) and all(s.doc_id in ref.authoritative_docs for s in cited),
        "verification.periods_aligned": bool(cited) and all(s.period in ref.periods for s in cited),
        "computation.reproducible_from_evidence": calc_ok,
        "computation.unit_matches": ans.unit == ref.unit,
        "final.value_within_tolerance": ans.value is not None and abs(ans.value - ref.expected) <= ref.tolerance,
    }


def summarise(results: list[dict[str, bool]]) -> dict[str, float]:
    n = len(results) or 1
    items = [v for r in results for v in r.values()]
    out = {"atomic_score": sum(items) / max(len(items), 1),
           "strict_pass": sum(all(r.values()) for r in results) / n}
    for dim in ("acquisition", "verification", "computation", "final"):
        vals = [v for r in results for k, v in r.items() if k.startswith(dim)]
        out[dim] = sum(vals) / max(len(vals), 1)
    correct = [r for r in results if r["final.value_within_tolerance"]]
    out["correct_but_unsupported"] = sum(not all(r.values()) for r in correct) / max(len(correct), 1)
    return out


def compare_modes(refs: list[Reference], retrieve: Callable[[str], list[Span]],
                  answer: Callable[..., Answer]) -> dict[str, dict]:
    report = {}
    for mode in ("tools", "documents_only"):
        rows, tokens = [], 0
        for ref in refs:
            spans = retrieve(ref.question)
            ans = answer(ref.question, spans, mode=mode)
            rows.append(score(ref, ans, spans))
            tokens += ans.input_tokens + ans.output_tokens
        report[mode] = summarise(rows) | {"tokens_per_question": tokens / max(len(refs), 1)}
    return report


# ---------- 2. Always-on refresher ----------

@dataclass
class Claim:
    claim_id: str
    question: str
    entity: str
    period_rule: str                    # "latest" follows new filings; "FY2025" is fixed
    answer: Answer
    depends_on: set[str]                # doc_ids behind the current answer
    published_in: list[str] = field(default_factory=list)   # lesson or note ids


@dataclass(frozen=True)
class NewDocument:
    doc_id: str
    entity: str
    doc_type: str                       # "annual report", "RNS", "call transcript"
    period: str
    supersedes: str | None = None       # set for restatements and amended filings


class AuditLog:
    """Append-only JSONL; each record carries the hash of the previous one."""

    def __init__(self, path: Path):
        self.path = path
        path.parent.mkdir(parents=True, exist_ok=True)
        lines = path.read_text().splitlines() if path.exists() else []
        self.prev = json.loads(lines[-1])["hash"] if lines else "genesis"

    def write(self, event: str, **data) -> None:
        rec = {"ts": time.time(), "event": event, "prev": self.prev, **data}
        rec["hash"] = hashlib.sha256(json.dumps(rec, sort_keys=True, default=str).encode()).hexdigest()
        with self.path.open("a") as f:
            f.write(json.dumps(rec, default=str) + "\n")
        self.prev = rec["hash"]


class Refresher:
    def __init__(self, claims: dict[str, Claim], retrieve: Callable[[str, str], list[Span]],
                 answer: Callable[..., Answer], notify: Callable[[dict], None], audit: AuditLog,
                 max_reruns: int = 50, token_budget: int = 400_000,
                 material_rel: Decimal = Decimal("0.005")):
        self.claims, self.retrieve, self.answer = claims, retrieve, answer
        self.notify, self.audit = notify, audit
        self.max_reruns, self.token_budget, self.material_rel = max_reruns, token_budget, material_rel
        self.reviews: dict[str, dict] = {}

    def dependents(self, doc: NewDocument) -> list[Claim]:
        return [c for c in self.claims.values()
                if (doc.supersedes is not None and doc.supersedes in c.depends_on)
                or (c.entity == doc.entity and c.period_rule == "latest")]

    def diff(self, old: Answer, new: Answer) -> dict:
        if old.value is None or new.value is None or old.unit != new.unit:
            material = (old.value, old.unit) != (new.value, new.unit)
            rel = None
        else:
            rel = abs(new.value - old.value) / max(abs(old.value), Decimal("1e-9"))
            material = rel > self.material_rel
        return {"old": str(old.value), "new": str(new.value), "unit": new.unit,
                "relative_change": None if rel is None else str(rel), "material": material,
                "evidence_added": sorted(set(new.cited_span_ids) - set(old.cited_span_ids)),
                "evidence_dropped": sorted(set(old.cited_span_ids) - set(new.cited_span_ids))}

    def on_new_document(self, doc: NewDocument) -> list[str]:
        if os.environ.get("REFRESHER_KILL_SWITCH") == "1":
            self.audit.write("halted_by_kill_switch", doc=doc.doc_id)
            return []
        affected = self.dependents(doc)
        if len(affected) > self.max_reruns:
            self.audit.write("rerun_cap_hit", doc=doc.doc_id, affected=len(affected), cap=self.max_reruns)
            affected = affected[: self.max_reruns]
        spent, opened = 0, []
        for claim in affected:
            if spent >= self.token_budget:
                self.audit.write("token_budget_exhausted", doc=doc.doc_id, spent=spent)
                break
            spans = self.retrieve(claim.question, claim.entity)
            new = self.answer(claim.question, spans, mode="tools")
            spent += new.input_tokens + new.output_tokens
            d = self.diff(claim.answer, new)
            self.audit.write("claim_rerun", doc=doc.doc_id, claim=claim.claim_id, diff=d)
            docs = {s.doc_id for s in spans if s.span_id in new.cited_span_ids}
            if not d["material"]:
                claim.depends_on |= docs        # same figure, newer support: track it, publish nothing
                continue
            rid = str(uuid.uuid4())
            self.reviews[rid] = {"claim": claim.claim_id, "trigger": doc.doc_id, "diff": d,
                                 "new": new, "docs": docs, "status": "pending_review"}
            self.notify({"review_id": rid, "claim": claim.claim_id, "trigger": doc.doc_id,
                         "published_in": claim.published_in, **d})
            opened.append(rid)
        return opened

    def approve(self, review_id: str, reviewer: str) -> None:
        if not reviewer.strip():
            raise PermissionError("approval needs a named reviewer")
        review = self.reviews[review_id]
        if review["status"] != "pending_review":
            raise PermissionError(f"review {review_id} is {review['status']}")
        review |= {"status": "approved", "reviewer": reviewer}
        self.audit.write("review_approved", review=review_id, reviewer=reviewer)

    def apply(self, review_id: str, publish: Callable[[Claim], None]) -> None:
        """The only path to customers: a named person approved this exact diff."""
        review = self.reviews[review_id]
        if review["status"] != "approved":
            raise PermissionError("not approved; nothing is republished automatically")
        claim = self.claims[review["claim"]]
        claim.answer, claim.depends_on = review["new"], review["docs"]
        publish(claim)
        review["status"] = "applied"
        self.audit.write("claim_republished", review=review_id, claim=claim.claim_id)