Coverage for src / quber / playground / ranking.py: 40%
173 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Jev ranking of retrieved chunks for the playground's answer path.
3Ownership. This module decides which of a question's retrieved chunks the
4answer model reads when the playground's ranker setting is `jev`. `selection`
5holds the Haiku equivalent, and `retrieval.retrieve` chooses between them and
6composes the final context. Both satisfy `ChunkRanker`: the question and the
7candidates in, indices into the candidates out, most relevant first.
9How it ranks. Jev, TypeSafe's scoring model, does not read the candidates as
10one prompt. Every page and every chunk is judged alone, in its own request,
11so the ranker can afford the whole fused list rather than a window of it.
12Two passes. Each distinct page in the candidate list is asked whether it has
13material relevant to the question, and a page answered no is dropped with
14every candidate on it. Each candidate on a page that passed is then graded
15on a four-level rubric, and `Ranking.order` keeps the ones at or above the
16configured cut, highest first. Nothing is added back from fused order when
17fewer clear the cut: anything added that way is a chunk Jev graded as
18unrelated or same-topic, or one on a page Jev rejected.
20What leaves the host. The question, each page's text and each chunk's text
21are sent to TypeSafe's API, one text per request. TypeSafe states it does not
22train on requests; zero data retention is an enterprise option not yet taken.
24Failures. The SDK retries 429, 5xx, connection and timeout errors itself.
25After each pass, the calls that still failed are run once more. A page that
26fails twice counts as relevant, so its chunks are still graded. A chunk that
27fails twice goes after the graded ranking in fused order. Every failed page
28or chunk is named in the log and the trace, because it may be the one that
29holds the answer. Only when every call in a pass fails does `rank` raise, and
30`retrieve` falls back to fused order.
31"""
33from __future__ import annotations
35import asyncio
36import time
37from collections import Counter
38from functools import lru_cache
39from typing import (
40 TYPE_CHECKING,
41 Any,
42 Awaitable,
43 Callable,
44 Dict,
45 List,
46 Optional,
47 Protocol,
48 TypeVar,
49)
51import httpx2
52from loguru import logger
53from pydantic import BaseModel, Field
54from typesafe_sdk import (
55 AsyncTypeSafeClient,
56 Noul,
57 NoulAnswer,
58 RetryPolicy,
59 Score,
60 ScoreAnswer,
61 TypeSafeAPIError,
62 TypeSafeError,
63)
65from quber.agents.langsmith_tracer import usage_metadata_from
66from quber.playground import db
67from quber.playground.tracing import run_metadata, tracer
68from quber.settings import get_settings
70if TYPE_CHECKING:
71 from quber.playground.retrieval import RetrievedChunk
74class ChunkRanker(Protocol):
75 """The question and a candidate list in; indices into that list out, most relevant first."""
77 async def __call__(self, question: str, cands: List["RetrievedChunk"]) -> List[int]: ...
80# The questions Jev is asked. PydanticAI sends a field's description as the
81# question and nothing else, so the yes and no criteria are written into the
82# question text. The wording is positive throughout: Jev's documentation says
83# negated questions score worse.
84PAGE_QUESTION = (
85 "Does this page have material relevant to the query? "
86 "Yes when the page discusses the metric, line item, or statement the query asks about. "
87 "No when the page is about other subjects and nothing on it bears on the query."
88)
90# The chunk question is short because the levels carry the meaning. The levels
91# describe situations, not degrees, and carry no numbers. Level 2 is what keeps
92# a footnote that defines the metric, the components that reconcile to it, or
93# the prior period column in the context the answer model reads.
94CHUNK_QUESTION = "How does this chunk relate to the query?"
96LEVEL_UNRELATED = "The chunk is about a different subject; nothing in it bears on the query."
97LEVEL_SAME_TOPIC = (
98 "The chunk is on the same topic as the query but gives neither the asked-for item "
99 "nor anything that qualifies it."
100)
101LEVEL_SUPPORTS = (
102 "The chunk supports the asked-for item: a definition, a footnote, a component, "
103 "a reconciliation, or a comparison period for it."
104)
105LEVEL_STATES = "The chunk states the asked-for item directly, for the entity and period the query names."
108# The two questions as the SDK sends them. A Noul is a yes/no question: Jev
109# returns the probability of yes. A Score is a rubric: the levels are the
110# criteria in order, 0 first, and Jev returns a probability per level and the
111# probability-weighted level as the score. The SDK is called directly rather
112# than through a PydanticAI agent because this path makes a few hundred calls
113# per question, and the agent's per-call work measured three to four times
114# the CPU of the call it wraps, enough to saturate one event loop under a
115# batch of questions.
116PAGE = Noul(instructions=PAGE_QUESTION)
117CHUNK = Score(
118 instructions=CHUNK_QUESTION, criteria=[LEVEL_UNRELATED, LEVEL_SAME_TOPIC, LEVEL_SUPPORTS, LEVEL_STATES]
119)
121# A page is relevant when Jev's probability of yes reaches this. It is where
122# the answer turns from no to yes, not a tuned bar, so it is not a setting.
123PAGE_THRESHOLD = 0.5
126class PageScore(BaseModel):
127 """One page's verdict. `page` is the stored 0-based index."""
129 page: int
130 relevant: bool
131 # How far Jev's probability sat from the yes/no threshold, 0 to 1.
132 confidence: Optional[float] = None
133 error: Optional[str] = None
136class ChunkScore(BaseModel):
137 """One candidate's grade. `index` is its position in the candidate list."""
139 index: int
140 chunk_id: str
141 page: int
142 # The nearest level, 0 to 3.
143 grade: Optional[int] = None
144 # Jev's position on the rubric, the probability-weighted level. This is
145 # what the cut and the ordering use: two chunks can share a grade and
146 # differ in score.
147 score: Optional[float] = None
148 error: Optional[str] = None
151class Ranking(BaseModel):
152 """Everything both passes returned for one question."""
154 pages: List[PageScore]
155 # Every candidate on a page that passed, in candidate order.
156 scores: List[ChunkScore]
157 grade_cut: float
158 model: Optional[str] = None
159 seconds: float = 0.0
160 input_tokens: int = 0
161 # Count per error class or HTTP status, after the SDK's retries and the
162 # ranker's second attempt.
163 failures: Dict[str, int] = Field(default_factory=dict)
165 def order(self) -> List[int]:
166 """Candidate indices at or above the cut, highest score first, then the
167 chunks whose calls failed, in candidate order."""
168 graded = [c for c in self.scores if c.score is not None and c.score >= self.grade_cut]
169 graded.sort(key=lambda c: c.score or 0.0, reverse=True)
170 failed = [c for c in self.scores if c.error is not None]
171 return [c.index for c in graded] + [c.index for c in failed]
173 def cleared(self) -> int:
174 return sum(1 for c in self.scores if c.score is not None and c.score >= self.grade_cut)
176 def kept_pages(self) -> List[int]:
177 return [p.page for p in self.pages if p.relevant]
180@lru_cache(maxsize=1)
181def client() -> AsyncTypeSafeClient:
182 """The one SDK client both passes share, so every call reuses one
183 connection pool. The pool is unbounded: a batch of questions, each with
184 `concurrency` calls in flight, would otherwise queue on httpx's default
185 of 100 connections and reopen the ones past its keep-alive cap of 20 on
186 every call. The per-question semaphore and the batch's question cap are
187 what bound the load."""
188 ts = get_settings().typesafe
189 if not ts.api_key:
190 raise RuntimeError("TYPESAFE_API_KEY is not set; the jev ranker needs it.")
191 http = httpx2.AsyncClient(
192 timeout=ts.timeout_seconds, limits=httpx2.Limits(max_connections=None, max_keepalive_connections=None)
193 )
194 return AsyncTypeSafeClient(
195 api_key=ts.api_key, retry=RetryPolicy(max_retries=ts.max_retries), http_client=http
196 )
199async def page_texts(cands: List["RetrievedChunk"]) -> Dict[int, str]:
200 """The text of each distinct candidate page: its chunks that are not line
201 records, joined in stored order.
203 Candidates carry no document key. Their chunk ids are the parse's own
204 UUIDs, so the document that holds the most of them is the document, and
205 the pages are read from it.
206 """
207 pages = sorted({c.page for c in cands})
208 ids = [c.chunk_id for c in cands]
209 sql = """
210 WITH doc AS (
211 SELECT document_id FROM ade_playground.chunks
212 WHERE chunk_id = ANY(%(ids)s)
213 GROUP BY document_id ORDER BY count(*) DESC LIMIT 1
214 )
215 SELECT k.page, k.content
216 FROM ade_playground.chunks k JOIN doc ON doc.document_id = k.document_id
217 WHERE k.page = ANY(%(pages)s) AND k.chunk_type <> 'line_item'
218 ORDER BY k.page, k.id
219 """
220 texts: Dict[int, List[str]] = {p: [] for p in pages}
221 async with db.connect_async() as conn:
222 for page, content in await (await conn.execute(sql, {"ids": ids, "pages": pages})).fetchall():
223 texts[page].append(content)
224 return {p: "\n\n".join(parts) for p, parts in texts.items()}
227T = TypeVar("T")
230class _Tally:
231 """What every call adds up to: tokens, the resolved model id, failures.
232 Shaped so `usage_metadata_from` can read it like a run usage."""
234 def __init__(self) -> None:
235 self.input_tokens = 0
236 self.output_tokens = 0
237 self.model: Optional[str] = None
238 self.failures: Counter[str] = Counter()
240 def record(self, response: Any) -> None:
241 self.input_tokens += response.usage.input_tokens or 0
242 self.output_tokens += response.usage.output_tokens or 0
243 self.model = response.model or self.model
245 def failed(self, exc: Exception) -> str:
246 key = str(exc.status) if isinstance(exc, TypeSafeAPIError) else type(exc).__name__
247 self.failures[key] += 1
248 return f"{key}: {exc}"
251async def _twice(
252 items: List[T], call: Callable[[T], Awaitable[Any]], failed: Callable[[Any], bool], concurrency: int
253) -> List[Any]:
254 """Run `call` over `items` at most `concurrency` at a time, then once more
255 over the items whose result `failed`. Results in item order."""
256 sem = asyncio.Semaphore(concurrency)
258 async def guarded(item: T) -> Any:
259 async with sem:
260 return await call(item)
262 results = list(await asyncio.gather(*(guarded(i) for i in items)))
263 retry = [n for n, r in enumerate(results) if failed(r)]
264 if retry:
265 again = await asyncio.gather(*(guarded(items[n]) for n in retry))
266 for n, r in zip(retry, again, strict=True):
267 results[n] = r
268 return results
271def _levels() -> str:
272 return "\n".join(
273 f"{n}: {d}" for n, d in enumerate((LEVEL_UNRELATED, LEVEL_SAME_TOPIC, LEVEL_SUPPORTS, LEVEL_STATES))
274 )
277def _run_outputs(tally: _Tally, started: float, **fields: Any) -> Dict[str, Any]:
278 """The fields every traced pass records, plus the pass's own."""
279 return {
280 **fields,
281 "usage_metadata": usage_metadata_from(tally),
282 "failures": dict(tally.failures),
283 "seconds": round(time.perf_counter() - started, 2),
284 "model": tally.model,
285 }
288async def judge_pages(question: str, cands: List["RetrievedChunk"]) -> tuple[List[PageScore], _Tally]:
289 """Pass 1: one yes/no call per distinct candidate page. One traced run
290 named `page_rank`, with every verdict, relevant pages first."""
291 ts = get_settings().typesafe
292 tally = _Tally()
293 texts = await page_texts(cands)
294 inputs = {
295 "messages": [{"role": "system", "content": PAGE_QUESTION}, {"role": "user", "content": question}],
296 "pages": [p + 1 for p in sorted(texts)],
297 }
299 async def judge(page: int) -> PageScore:
300 text = texts.get(page, "")
301 if not text:
302 return PageScore(page=page, relevant=True)
303 try:
304 response = await client().system_one(
305 f"QUERY: {question}\n\nPAGE {page + 1}:\n{text}", {"relevant": PAGE}, model=ts.model
306 )
307 except TypeSafeError as exc:
308 return PageScore(page=page, relevant=True, error=tally.failed(exc))
309 tally.record(response)
310 answer = response.answers["relevant"]
311 assert isinstance(answer, NoulAnswer)
312 # Confidence is how far the probability sat from the threshold, 0 at
313 # the threshold and 1 at either end, so a page that could go either
314 # way on another call reads as near 0.
315 return PageScore(
316 page=page,
317 relevant=answer.noul >= PAGE_THRESHOLD,
318 confidence=round(abs(answer.noul - PAGE_THRESHOLD) * 2, 4),
319 )
321 started = time.perf_counter()
322 async with tracer().llm_run("page_rank", inputs, model=ts.model, extra_metadata=run_metadata()) as run:
323 pages = await _twice(sorted(texts), judge, lambda p: p.error is not None, ts.concurrency)
324 judged = [p for p in pages if texts.get(p.page)]
325 if judged and all(p.error is not None for p in judged):
326 raise RuntimeError(f"every page call failed: {dict(tally.failures)}")
327 ordered = sorted(pages, key=lambda p: (not p.relevant, p.page))
328 run.outputs = _run_outputs(
329 tally,
330 started,
331 messages=[{"role": "assistant", "content": str([p.page + 1 for p in pages if p.relevant])}],
332 pages=[{**p.model_dump(), "page": p.page + 1} for p in ordered],
333 kept=sum(1 for p in pages if p.relevant),
334 )
335 return pages, tally
338async def grade_chunks(
339 question: str, cands: List["RetrievedChunk"], kept: set[int]
340) -> tuple[List[ChunkScore], _Tally]:
341 """Pass 2: one rubric call per candidate on a kept page. One traced run
342 named `chunk_rank`, with every score, best first, failed chunks last."""
343 ts = get_settings().typesafe
344 tally = _Tally()
345 items = [(i, c) for i, c in enumerate(cands) if c.page in kept]
346 inputs = {
347 "messages": [
348 {"role": "system", "content": f"{CHUNK_QUESTION}\n{_levels()}"},
349 {"role": "user", "content": question},
350 ],
351 "chunk_ids": [c.chunk_id for _, c in items],
352 }
354 async def grade(item: tuple[int, "RetrievedChunk"]) -> ChunkScore:
355 i, c = item
356 prompt = f"QUERY: {question}\n\nCHUNK (page {c.page + 1}, {c.chunk_type}):\n{c.content}"
357 try:
358 response = await client().system_one(prompt, {"grade": CHUNK}, model=ts.model)
359 except TypeSafeError as exc:
360 return ChunkScore(index=i, chunk_id=c.chunk_id, page=c.page, error=tally.failed(exc))
361 tally.record(response)
362 answer = response.answers["grade"]
363 assert isinstance(answer, ScoreAnswer)
364 return ChunkScore(
365 index=i,
366 chunk_id=c.chunk_id,
367 page=c.page,
368 grade=min(int(answer.score + 0.5), len(CHUNK.criteria) - 1),
369 score=answer.score,
370 )
372 started = time.perf_counter()
373 async with tracer().llm_run("chunk_rank", inputs, model=ts.model, extra_metadata=run_metadata()) as run:
374 scores = await _twice(items, grade, lambda s: s.error is not None, ts.concurrency)
375 if scores and all(s.error is not None for s in scores):
376 raise RuntimeError(f"every chunk call failed: {dict(tally.failures)}")
377 graded = sorted((s for s in scores if s.error is None), key=lambda s: s.score or 0.0, reverse=True)
378 failed = [s for s in scores if s.error is not None]
379 cleared = [s for s in graded if (s.score or 0.0) >= ts.grade_cut]
380 run.outputs = _run_outputs(
381 tally,
382 started,
383 messages=[{"role": "assistant", "content": str([s.chunk_id for s in cleared])}],
384 scores=[{**s.model_dump(), "page": s.page + 1} for s in graded + failed],
385 grade_cut=ts.grade_cut,
386 cleared=len(cleared),
387 below_cut=len(graded) - len(cleared),
388 )
389 return scores, tally
392async def score(question: str, cands: List["RetrievedChunk"]) -> Ranking:
393 """Both passes over the candidates, with everything Jev returned."""
394 ts = get_settings().typesafe
395 started = time.perf_counter()
396 pages, page_tally = await judge_pages(question, cands)
397 kept = {p.page for p in pages if p.relevant}
398 scores, chunk_tally = await grade_chunks(question, cands, kept)
399 return Ranking(
400 pages=pages,
401 scores=scores,
402 grade_cut=ts.grade_cut,
403 model=chunk_tally.model or page_tally.model,
404 seconds=time.perf_counter() - started,
405 input_tokens=page_tally.input_tokens + chunk_tally.input_tokens,
406 failures=dict(page_tally.failures + chunk_tally.failures),
407 )
410async def rank(question: str, cands: List["RetrievedChunk"]) -> List[int]:
411 """Indices into `cands` in the order the answer model should read them."""
412 ranking = await score(question, cands)
413 order = ranking.order()
414 logger.debug(
415 "jev ranker: {} pages judged, {} kept; {} of {} candidates graded, {} cleared the cut of {}; "
416 "{:.1f}s, {} input tokens, failures {}",
417 len(ranking.pages),
418 len(ranking.kept_pages()),
419 len(ranking.scores),
420 len(cands),
421 ranking.cleared(),
422 ranking.grade_cut,
423 ranking.seconds,
424 ranking.input_tokens,
425 ranking.failures or "none",
426 )
427 failed_pages = [p for p in ranking.pages if p.error]
428 failed_chunks = [s for s in ranking.scores if s.error]
429 if failed_pages or failed_chunks:
430 logger.warning(
431 "jev ranker: {} page calls and {} chunk calls failed twice; pages {}; chunks {}",
432 len(failed_pages),
433 len(failed_chunks),
434 [(p.page + 1, p.error) for p in failed_pages],
435 [(s.chunk_id, s.error) for s in failed_chunks],
436 )
437 if ranking.failures.get("429"):
438 logger.warning(
439 "jev ranker: {} calls were throttled (429) past the SDK's retries; TypeSafe's limits may have changed",
440 ranking.failures["429"],
441 )
442 return order
445def rank_sync(question: str, cands: List["RetrievedChunk"]) -> List[int]:
446 return asyncio.run(rank(question, cands))