Coverage for src / quber / playground / benchmark / run_shapes.py: 30%
173 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Measure four ways of returning an answer over one set of questions.
3Every arm answers from the SAME retrieved context: retrieval runs once per
4question and its chunks are handed to all four, so a difference in the results
5is a difference in the return shape and not in what was retrieved.
7The arms:
9 baseline the playground's current agent — prose plus cited ids
10 fixed one envelope, payload chosen from the closed union in `shapes`
11 planned a planner declares fields per question; they are compiled and applied
12 declared the caller states the shape; nothing is inferred
14Scoring is deterministic. Gold values were copied out of printed table cells,
15so a returned figure is compared to the gold string directly and, separately,
16as a number — no judge is asked whether two figures agree. The baseline has no
17value field to compare, so it is scored on whether the gold figure appears
18anywhere in its prose: that measures whether the information was found, which
19is the point the comparison has to be fair about.
21Usage:
22 python -m quber.playground.benchmark.run_shapes --gold <path>
23 python -m quber.playground.benchmark.run_shapes --gold <path> --limit 4
24"""
26from __future__ import annotations
28import argparse
29import asyncio
30import json
31import re
32import time
33from decimal import Decimal
34from pathlib import Path
35from typing import Any, Dict, List, Optional
37from loguru import logger
38from pydantic import BaseModel, Field
40from quber.playground import db
41from quber.playground.agent import answer as baseline_answer
42from quber.playground.answers import declared, fixed, planned
43from quber.playground.answers.figures import as_number, norm_printed
44from quber.playground.retrieval import RetrievedChunk, retrieve
46OUTPUT = Path("output/benchmark/answer_shapes.json")
47CONCURRENCY = 4
48K = 10
50ARMS = ("baseline", "fixed", "planned", "declared")
53class ProseAnswer(BaseModel):
54 """The shape a caller declares when it wants an explanation."""
56 text: str = Field(description="The explanation, in prose.")
57 cited_ids: List[str] = Field(default_factory=list)
60# ---------------------------------------------------------------- normalizing
62# `norm_printed` and `as_number` live in `answers.figures`, which is the only
63# place a printed figure is read as a number. Scoring a run and exporting a
64# batch both compare figures, and each keeping its own reading is how the same
65# cell comes out with two different signs.
68def values_agree(got: Optional[str], gold: str) -> Dict[str, bool]:
69 """Exact-as-printed and numeric agreement between a returned and gold figure."""
70 if got is None:
71 return {"exact": False, "numeric": False}
72 exact = norm_printed(got) == norm_printed(gold)
73 a, b = as_number(got), as_number(gold)
74 return {"exact": exact, "numeric": bool(a is not None and b is not None and a == b)}
77def looks_like_prose(s: Optional[str]) -> bool:
78 """True when a value field holds a sentence rather than a figure.
80 A printed figure is short and at most a few tokens ('$1.2 million',
81 '4.6x'). Four or more tokens, or a long string, means the value came back
82 wrapped in words — the failure this whole exercise is about.
83 """
84 if not s:
85 return False
86 t = norm_printed(s)
87 return len(t) > 40 or len(t.split()) >= 4
90def gold_in_prose(prose: str, gold: str) -> bool:
91 """Whether the gold figure appears in a prose answer, as printed or as a number.
93 Whitespace is ignored on the literal comparison: a cell prints "$ (10,550)"
94 where prose writes "$(10,550)", and that is the same figure.
96 Scale words are deliberately not part of a token. Both sides drop the scale
97 when they parse — "$241.4M" and "$241.4 million" each read as 241.4 — so
98 absorbing a trailing "thousand" into the token achieves nothing and breaks
99 any figure whose closing parenthesis then sits mid-token.
100 """
101 if not prose:
102 return False
103 if re.sub(r"\s+", "", gold) in re.sub(r"\s+", "", prose):
104 return True
105 g = as_number(gold)
106 if g is None:
107 return False
108 # A decimal part must be a real decimal: `\.?\d*` would also match the full
109 # stop ending a sentence, leaving a token that cannot parse. The optional
110 # parenthesis is allowed on either side of the currency sign because prose
111 # writes "$(10,550)" and cells print "(10,550)".
112 for tok in re.findall(r"\(?-?\$?\(?[\d][\d,]*(?:\.\d+)?\)?%?[xX]?", prose):
113 if as_number(tok) == g:
114 return True
115 return False
118# ------------------------------------------------------------------ grounding
121def cell_text(doc_key: str, cell_id: str) -> Optional[str]:
122 """The text printed in one tagged table cell of one document, if it exists."""
123 with db.connect() as conn:
124 row = conn.execute(
125 """SELECT c.content FROM ade_playground.chunks c
126 JOIN ade_playground.documents d ON d.id = c.document_id
127 WHERE d.doc_key = %s AND c.content LIKE %s
128 LIMIT 1""",
129 (doc_key, f'%id="{cell_id}"%'),
130 ).fetchone()
131 if not row:
132 return None
133 m = re.search(rf'<td id="{re.escape(cell_id)}"[^>]*>(.*?)</td>', row[0], re.S)
134 return norm_printed(re.sub(r"<[^>]+>", "", m.group(1))) if m else None
137def retrieval_hit(chunks: List[RetrievedChunk], cell_id: str) -> bool:
138 return any(f'id="{cell_id}"' in c.content for c in chunks)
141# ----------------------------------------------------------------------- arms
144def scalar_field(obj: Any) -> Optional[str]:
145 """The single figure a compiled or declared model returned, if it has one.
147 A planner-compiled model names its field whatever suited the question, so
148 the figure is found by taking the model's first string field that is not
149 bookkeeping.
150 """
151 if obj is None:
152 return None
153 data = obj.model_dump() if isinstance(obj, BaseModel) else dict(obj)
154 for key, val in data.items():
155 if key in ("cited_ids", "not_found"):
156 continue
157 if isinstance(val, str) and val.strip():
158 return val
159 if isinstance(val, (int, float, Decimal)):
160 return str(val)
161 return None
164async def run_arm(arm: str, q: Dict[str, Any], chunks: List[RetrievedChunk]) -> Dict[str, Any]:
165 """One arm's attempt at one question, with what it returned and how long it took."""
166 started = time.perf_counter()
167 out: Dict[str, Any] = {"arm": arm}
168 is_value = q["kind"] == "value"
169 try:
170 if arm == "baseline":
171 res = await baseline_answer(q["question"], chunks)
172 out.update(
173 shape="prose",
174 value=None,
175 prose=res.answer,
176 cited_ids=res.cited_ids,
177 )
179 elif arm == "fixed":
180 res = await fixed.answer(q["question"], chunks)
181 payload = res.payload
182 out.update(shape=payload.kind, cited_ids=res.cited_ids)
183 if payload.kind == "scalar":
184 # The whole scalar, not just the printed form. Anything reading
185 # this file downstream needs `number` to have a numeric column
186 # without parsing `value` back into one, and a second parser of
187 # printed figures is how a negative silently becomes positive.
188 out.update(
189 value=payload.value,
190 number=payload.number,
191 unit=payload.unit,
192 period=payload.period,
193 label=payload.label,
194 source_id=payload.source_id,
195 )
196 elif payload.kind == "prose":
197 out.update(value=None, prose=payload.text)
198 else:
199 out.update(value=None, payload=payload.model_dump(mode="json"))
201 elif arm == "planned":
202 plan = await planned.plan_for(q["question"])
203 out["plan"] = plan.model_dump()
204 if plan.wants_value and plan.fields:
205 model = planned.compile_model(plan)
206 res = await declared.answer(q["question"], chunks, model)
207 data = res.model_dump()
208 out.update(
209 shape="value",
210 value=scalar_field(res),
211 cited_ids=data.get("cited_ids", []),
212 fields=data,
213 )
214 else:
215 res = await declared.answer(q["question"], chunks, ProseAnswer)
216 out.update(shape="prose", value=None, prose=res.text, cited_ids=res.cited_ids)
218 elif arm == "declared":
219 # The caller knows which it wants, so the shape is supplied, not inferred.
220 if is_value:
221 res = await declared.answer(q["question"], chunks, declared.DeclaredScalar)
222 out.update(
223 shape="value",
224 value=res.value,
225 unit=res.unit,
226 period=res.period,
227 cited_ids=res.cited_ids,
228 not_found=res.not_found,
229 )
230 else:
231 res = await declared.answer(q["question"], chunks, ProseAnswer)
232 out.update(shape="prose", value=None, prose=res.text, cited_ids=res.cited_ids)
234 except Exception as exc:
235 out.update(error=f"{type(exc).__name__}: {exc}", shape=None, value=None, cited_ids=[])
237 out["seconds"] = round(time.perf_counter() - started, 2)
238 return score(out, q)
241def score(out: Dict[str, Any], q: Dict[str, Any]) -> Dict[str, Any]:
242 """Add the measurements this comparison turns on."""
243 is_value = q["kind"] == "value"
244 shape = out.get("shape")
245 out["shape_correct"] = shape in ("scalar", "value") if is_value else shape in ("prose", "unanswerable")
246 out["cites_any"] = bool(out.get("cited_ids"))
248 if is_value:
249 gold = q["gold_value"]
250 if out.get("value") is not None:
251 out.update(values_agree(out["value"], gold))
252 out["prose_leak"] = looks_like_prose(out["value"])
253 out["found_gold"] = out["exact"] or out["numeric"]
254 else:
255 # No value field: the arm either answered in prose or failed.
256 prose = out.get("prose") or ""
257 out.update(exact=False, numeric=False, prose_leak=bool(prose))
258 out["found_gold"] = gold_in_prose(prose, gold)
259 return out
262# ------------------------------------------------------------------ the sweep
265async def run_question(q: Dict[str, Any], sem: asyncio.Semaphore) -> Dict[str, Any]:
266 async with sem:
267 chunks = await retrieve(q["slug"], q["question"], k=K)
268 row: Dict[str, Any] = {
269 **q,
270 "retrieved_pages": sorted({c.page + 1 for c in chunks}),
271 "retrieval_hit": retrieval_hit(chunks, q["gold_cell_id"]) if q["kind"] == "value" else None,
272 }
273 results = []
274 for arm in ARMS:
275 async with sem:
276 results.append(await run_arm(arm, q, chunks))
277 row["arms"] = results
278 return row
281def summarize(rows: List[Dict[str, Any]]) -> Dict[str, Any]:
282 """Per-arm totals over the value questions and the prose controls."""
283 value_rows = [r for r in rows if r["kind"] == "value"]
284 hit_rows = [r for r in value_rows if r["retrieval_hit"]]
285 prose_rows = [r for r in rows if r["kind"] == "prose"]
286 out: Dict[str, Any] = {
287 "value_questions": len(value_rows),
288 "value_questions_with_retrieval_hit": len(hit_rows),
289 "prose_controls": len(prose_rows),
290 "arms": {},
291 }
293 def arm_of(row: Dict[str, Any], arm: str) -> Dict[str, Any]:
294 return next(a for a in row["arms"] if a["arm"] == arm)
296 for arm in ARMS:
297 vs = [arm_of(r, arm) for r in hit_rows]
298 ps = [arm_of(r, arm) for r in prose_rows]
299 n = len(vs) or 1
300 out["arms"][arm] = {
301 "value_exact_as_printed": sum(a.get("exact", False) for a in vs),
302 "value_numeric_match": sum(a.get("numeric", False) for a in vs),
303 "found_gold_anywhere": sum(a.get("found_gold", False) for a in vs),
304 "returned_a_value_field": sum(a.get("value") is not None for a in vs),
305 "prose_leak": sum(a.get("prose_leak", False) for a in vs),
306 "shape_correct_on_values": sum(a["shape_correct"] for a in vs),
307 "shape_correct_on_prose": sum(a["shape_correct"] for a in ps),
308 "cites_any": sum(a["cites_any"] for a in vs),
309 "errors": sum("error" in a for a in vs + ps),
310 "mean_seconds": round(sum(a["seconds"] for a in vs) / n, 2),
311 }
312 return out
315async def run(gold_path: Path, limit: Optional[int]) -> None:
316 gold = json.loads(gold_path.read_text())
317 questions = gold[:limit] if limit else gold
318 logger.info("{n} questions x {a} arms", n=len(questions), a=len(ARMS))
319 sem = asyncio.Semaphore(CONCURRENCY)
320 rows = await asyncio.gather(*(run_question(q, sem) for q in questions))
321 summary = summarize(list(rows))
323 OUTPUT.parent.mkdir(parents=True, exist_ok=True)
324 OUTPUT.write_text(json.dumps({"summary": summary, "rows": rows}, indent=2, default=str))
325 logger.success("written to {p}", p=OUTPUT)
326 print(json.dumps(summary, indent=2))
329def rescore(path: Path) -> None:
330 """Recompute the measurements over a finished run's stored answers.
332 Every arm's raw output is kept in the results file, so a correction to the
333 scoring is applied by re-reading it — no model is called again, and the
334 answers being scored are byte-identical to the ones originally returned.
335 """
336 data = json.loads(path.read_text())
337 for row in data["rows"]:
338 row["arms"] = [score(arm, row) for arm in row["arms"]]
339 data["summary"] = summarize(data["rows"])
340 path.write_text(json.dumps(data, indent=2, default=str))
341 logger.success("rescored {p}", p=path)
342 print(json.dumps(data["summary"], indent=2))
345def main() -> None:
346 p = argparse.ArgumentParser(description=__doc__)
347 p.add_argument("--gold", type=Path, help="gold question set JSON")
348 p.add_argument("--limit", type=int, default=None)
349 p.add_argument("--rescore", type=Path, help="recompute scores over a finished run")
350 args = p.parse_args()
351 if args.rescore:
352 rescore(args.rescore)
353 return
354 if not args.gold:
355 p.error("--gold is required unless --rescore is given")
356 asyncio.run(run(args.gold, args.limit))
359if __name__ == "__main__":
360 main()