Coverage for src / quber / playground / selection.py: 0%

1"""Candidate selection agent for the playground's retrieval path. 

2 

3Given a question and a wide, source-diverse pool of retrieved chunks, an agent 

4decides which chunks contain the information the question needs. Rank order 

5alone cannot do this: on the evaluation corpus, the answering chunk sat as 

6deep as fused rank 63 while near-duplicates and on-vocabulary neighbors filled 

7the top slots. The agent reads the pool and selects; it does not answer. 

8 

9Candidates are sent in full — no truncation. The 90-candidate pool measured 

10about 32k tokens on the evaluation corpus, well inside the model context, and 

11any cut risks hiding the one row of a large table that answers the question. 

12Labels are local indices that carry no information about any expected answer. 

13""" 

14 

15from __future__ import annotations 

16 

17import asyncio 

18from functools import lru_cache 

19from typing import TYPE_CHECKING, List 

20 

21from pydantic import BaseModel, Field 

22from pydantic_ai import Agent 

23from pydantic_ai.settings import ModelSettings 

24 

25from quber.agents.langsmith_tracer import usage_metadata_from 

26from quber.playground.agent import anthropic_model 

27from quber.playground.tracing import run_metadata, tracer 

28 

29if TYPE_CHECKING: 

30 from quber.playground.retrieval import RetrievedChunk 

31 

32MODEL = "claude-haiku-4-5" 

33 

34SYSTEM_PROMPT = """\ 

35You are a retrieval triage step for financial-document question answering. 

36You are given a QUESTION and a numbered list of CANDIDATE context chunks 

37retrieved from one document. Your ONLY job is to decide which candidates 

38contain the specific information needed to answer the question. 

39 

40Do NOT answer the question. Do NOT compute or state any figure. Select 

41candidates. 

42 

43Rules: 

44- Return the labels (e.g. C0, C7) of every candidate that contains information 

45 needed to answer the question, ordered MOST relevant first. 

46- A candidate is relevant if it holds the exact line, row, or statement the 

47 question asks about (matching the entity, period, metric, and units). Prefer 

48 the candidate that states the asked-for value most directly and specifically. 

49- If several near-duplicate candidates exist, include the one(s) that most 

50 precisely match the question's entity and period; you may include a few close 

51 alternates but keep the best first. 

52- Only return labels that appear in the candidate list. Never invent a label. 

53- If nothing is relevant, return an empty list. 

54""" 

55 

56 

57class Selection(BaseModel): 

58 labels: List[str] = Field( 

59 default_factory=list, 

60 description="Candidate labels (e.g. 'C3') that answer the question, most relevant first.", 

61 ) 

62 

63 

64@lru_cache(maxsize=1) 

65def _agent() -> Agent[None, Selection]: 

66 # Built over whichever credential the host presents, the same way the 

67 # answer model is. The hosted task carries only the API key; demanding the 

68 # OAuth token here meant selection never ran there and every hosted answer 

69 # came from fused order alone, with nothing but a log line to say so. 

70 return Agent( 

71 anthropic_model(MODEL), 

72 output_type=Selection, 

73 system_prompt=SYSTEM_PROMPT, 

74 model_settings=ModelSettings(temperature=0.0), 

75 ) 

76 

77 

78def _prompt(question: str, cands: List["RetrievedChunk"]) -> str: 

79 blocks = [ 

80 f"[label=C{i} | type {c.chunk_type} | page {c.page + 1}]\n{c.content}" for i, c in enumerate(cands) 

81 ] 

82 return ( 

83 f"QUESTION: {question}\n\n" 

84 f"CANDIDATES ({len(cands)}):\n\n" + "\n\n".join(blocks) + "\n\n" 

85 "Return the labels of the candidates that contain the information needed " 

86 "to answer the question, most relevant first." 

87 ) 

88 

89 

90async def select(question: str, cands: List["RetrievedChunk"]) -> List[int]: 

91 """Indices into `cands` the agent judged relevant, most relevant first. 

92 

93 Labels outside the candidate range and duplicates are discarded, so the 

94 result is always a valid, de-duplicated index list. 

95 """ 

96 prompt = _prompt(question, cands) 

97 inputs = { 

98 "messages": [ 

99 {"role": "system", "content": SYSTEM_PROMPT}, 

100 {"role": "user", "content": prompt}, 

101 ], 

102 "chunk_ids": [c.chunk_id for c in cands], 

103 } 

104 async with tracer().llm_run("select_chunks", inputs, model=MODEL, extra_metadata=run_metadata()) as run: 

105 result = await _agent().run(prompt) 

106 run.outputs = { 

107 "messages": [{"role": "assistant", "content": result.output.model_dump_json()}], 

108 "usage_metadata": usage_metadata_from(result.usage), 

109 } 

110 idxs: List[int] = [] 

111 for lab in result.output.labels: 

112 lab = lab.strip().lstrip("Cc") 

113 if lab.isdigit(): 

114 j = int(lab) 

115 if 0 <= j < len(cands) and j not in idxs: 

116 idxs.append(j) 

117 return idxs 

118 

119 

120def select_sync(question: str, cands: List["RetrievedChunk"]) -> List[int]: 

121 return asyncio.run(select(question, cands))