Coverage for src / quber / playground / retrieval.py: 55%
80 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Hybrid retrieval over the pgvector playground schema.
3Two stages. First, two rankings over one document's chunks are fused:
4pgvector cosine distance (meaning — a paraphrased question still lands) and
5Postgres full-text search over the same content (letter — a question that
6names a statement or line item verbatim finds it even when the whole-chunk
7embedding dilutes the match). Reciprocal rank fusion combines them without
8tuned weights; a query whose keywords match nothing degrades to the pure
9vector ranking.
11Second, a ranker orders the fused list, because rank order alone was
12measured insufficient: on the evaluation corpus the answering chunk sat as
13deep as fused rank 63 (a yield buried under same-vocabulary neighbors) or
14rank 32 (a guidance row losing to on-vocabulary sibling tables), and no
15scoring change recovered them. Which ranker runs is the playground's `ranker`
16setting, and both satisfy `ranking.ChunkRanker`.
18With `haiku`, the default, a selection agent reads a wide, source-diverse
19window of the fused ranking as one prompt and picks which chunks the question
20actually needs; the final context is composed from its picks, backfilled from
21fused order. The agent stage answered all eleven evaluation probes,
22deterministically, where fused order alone answered eight. Before the agent
23sees the window, line records are capped at two per parent table — one
24table's near-identical rows otherwise fill the window and crowd out other
25sources; each table's whole-table chunk is never capped, so the full table
26remains available even when its rows are.
28With `jev`, every page and every chunk of the whole fused list is judged in
29its own request, so neither the window nor the cap applies: both exist only
30because an LLM reads the pool as one prompt. The context is the ranker's
31order, as many as cleared its cut up to `k`, with nothing backfilled, because
32a chunk added from fused order is one Jev graded as not bearing on the
33question or one on a page Jev rejected.
34"""
36from __future__ import annotations
38import asyncio
39from dataclasses import dataclass
40from typing import Dict, List, Optional
42from loguru import logger
44from quber.playground import db
45from quber.playground.embedding import embed_query
46from quber.playground.ranking import ChunkRanker
47from quber.settings import get_settings
50@dataclass
51class RetrievedChunk:
52 chunk_id: str
53 chunk_type: str
54 page: int
55 content: str
56 score: float # fused RRF score (roughly 0..0.033); higher is better
57 parent_chunk_id: Optional[str] = None # line item -> its whole-table record
60# Fusion constants, grid-tested on the evaluation corpus (the arbor
61# preferred-stock-dividends probe: keyword rank 2, vector rank 52). A deep
62# pool lets a chunk far down one list still collect its credit from the
63# other, and the lower RRF constant weights a top keyword rank enough to
64# surface it. 60/40 left that probe at fused rank 16; 20/200 puts it at 8.
65RRF_K = 20
66POOL = 200 # candidates taken from each ranking before fusion
68# The selection agent's window into the fused ranking. 90 is the measured
69# minimum that contains every evaluation answer (the deepest sat at fused
70# rank 63); at 60 that chunk is unreachable by any downstream step.
71WINDOW = 90
72# Line records kept per parent table inside the window. Without a cap, one
73# large table's rows dominate the window (an 89-row table filled it almost
74# alone) and the agent over-selects siblings; two per table with the
75# whole-table chunk uncapped scored strictly better than no cap.
76PARENT_CAP = 2
77# The whole fused list, for the ranker that judges each chunk alone. It is
78# the union ceiling of the two POOL-sized rankings, so no chunk either
79# ranking returned is left out.
80FUSED_LIMIT = 2 * POOL
83async def retrieve(
84 doc_key: str, query: str, k: int = 10, types: Optional[List[str]] = None
85) -> List[RetrievedChunk]:
86 """The playground's retrieval: the fused list, ordered by the configured ranker.
88 Falls back to fused order if the ranker fails, so a provider outage
89 degrades ranking quality instead of breaking retrieval.
90 """
91 name = get_settings().playground.ranker
92 ranker: ChunkRanker
93 if name == "jev":
94 from quber.playground.ranking import rank
96 ranker = rank
97 pool = await _fused_window(doc_key, query, FUSED_LIMIT, types)
98 backfill = False
99 else:
100 from quber.playground.selection import select
102 ranker = select
103 pool = _cap_line_records(await _fused_window(doc_key, query, WINDOW, types), PARENT_CAP)
104 backfill = True
106 try:
107 picked = await ranker(query, pool)
108 except Exception as exc:
109 logger.warning("{} ranker failed ({}); falling back to fused order", name, exc)
110 picked, backfill = [], True
112 out: List[RetrievedChunk] = []
113 seen: set[str] = set()
114 for j in picked:
115 c = pool[j]
116 if c.chunk_id not in seen:
117 seen.add(c.chunk_id)
118 out.append(c)
119 if len(out) >= k:
120 break
121 if backfill:
122 for c in pool:
123 if len(out) >= k:
124 break
125 if c.chunk_id not in seen:
126 seen.add(c.chunk_id)
127 out.append(c)
128 # The fused positions say at a glance how deep the ranker reached.
129 position = {c.chunk_id: n + 1 for n, c in enumerate(pool)}
130 logger.debug(
131 "{} ranker sent {} of {} candidates; fused positions {}",
132 name,
133 len(out),
134 len(pool),
135 [position[c.chunk_id] for c in out],
136 )
137 return out
140def retrieve_sync(
141 doc_key: str, query: str, k: int = 10, types: Optional[List[str]] = None
142) -> List[RetrievedChunk]:
143 return asyncio.run(retrieve(doc_key, query, k, types))
146async def _fused_window(
147 doc_key: str, query: str, window: int, types: Optional[List[str]] = None
148) -> List[RetrievedChunk]:
149 """The top `window` chunks of the fused vector+keyword ranking.
151 Nothing here holds the event loop. The query embedding is a model forward
152 pass on this host, compute with nothing to await, so it runs in a worker
153 thread; the fused query is awaited on an async connection.
154 """
155 qvec = await asyncio.to_thread(embed_query, query)
156 type_filter = "AND c.chunk_type = ANY(%(types)s)" if types else ""
158 sql = f"""
159 WITH vec AS (
160 SELECT c.id, row_number() OVER (ORDER BY c.embedding <=> %(qvec)s::vector) AS rank
161 FROM ade_playground.chunks c
162 JOIN ade_playground.documents d ON d.id = c.document_id
163 WHERE d.doc_key = %(doc_key)s AND c.embedding IS NOT NULL {type_filter}
164 ORDER BY c.embedding <=> %(qvec)s::vector
165 LIMIT %(pool)s
166 ),
167 -- OR-joined lexemes of the question: AND semantics (websearch_to_tsquery)
168 -- fails whole questions, since one filler word absent from a chunk
169 -- ("find", "tell") kills the match. ts_rank_cd with log-length
170 -- normalization (flag 1) then rewards the chunk covering the most
171 -- question terms, densest first — unnormalized, a boilerplate table
172 -- repeating one query word hundreds of times (a subsidiary list
173 -- repeating the company name) outranks a short exact-coverage line.
174 oq AS (
175 SELECT to_tsquery('english',
176 COALESCE(NULLIF(array_to_string(
177 tsvector_to_array(to_tsvector('english', %(q)s)), ' | '), ''),
178 'zzznomatchzzz')) AS q
179 ),
180 kw AS (
181 SELECT c.id, row_number() OVER (
182 ORDER BY ts_rank_cd(c.content_tsv, oq.q, 1) DESC
183 ) AS rank
184 FROM ade_playground.chunks c
185 JOIN ade_playground.documents d ON d.id = c.document_id, oq
186 WHERE d.doc_key = %(doc_key)s AND c.embedding IS NOT NULL {type_filter}
187 AND c.content_tsv @@ oq.q
188 LIMIT %(pool)s
189 ),
190 fused AS (
191 SELECT COALESCE(vec.id, kw.id) AS id,
192 COALESCE(1.0 / (%(rrf)s + vec.rank), 0) +
193 COALESCE(1.0 / (%(rrf)s + kw.rank), 0) AS score
194 FROM vec FULL OUTER JOIN kw ON kw.id = vec.id
195 )
196 SELECT c.chunk_id, c.chunk_type, c.page, c.content, fused.score, c.parent_chunk_id
197 FROM fused JOIN ade_playground.chunks c ON c.id = fused.id
198 ORDER BY fused.score DESC
199 LIMIT %(cand)s
200 """
201 params = {"qvec": qvec, "doc_key": doc_key, "q": query, "pool": POOL, "rrf": RRF_K, "cand": window}
202 if types:
203 params["types"] = types
204 async with db.connect_async() as conn:
205 rows = await (await conn.execute(sql, params)).fetchall()
206 return [
207 RetrievedChunk(
208 chunk_id=r[0],
209 chunk_type=r[1],
210 page=r[2],
211 content=r[3],
212 score=float(r[4]),
213 parent_chunk_id=r[5],
214 )
215 for r in rows
216 ]
219def _cap_line_records(cands: List[RetrievedChunk], cap: int) -> List[RetrievedChunk]:
220 """Keep at most `cap` line records per parent table, in fused order.
222 Everything that is not a line record — prose, figures, and each table's
223 whole-table chunk — passes through untouched, so a capped table's full
224 content is still in the pool.
225 """
226 seen: Dict[str, int] = {}
227 out: List[RetrievedChunk] = []
228 for c in cands:
229 if c.chunk_type == "line_item" and c.parent_chunk_id:
230 n = seen.get(c.parent_chunk_id, 0)
231 if n >= cap:
232 continue
233 seen[c.parent_chunk_id] = n + 1
234 out.append(c)
235 return out