Coverage for src / quber / playground / answers / declared.py: 59%
49 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""Answering into a schema the caller declared.
3No inference. The caller states the shape it wants alongside the question and
4the agent is constrained to it — the same contract `Agent(output_type=...)`
5already offers, applied to the retrieval path. This is what a src-package
6caller does when it knows what it is asking for.
8`answer` is also the shared constrained-answer call: `planned` compiles a
9schema and hands it here, so the only difference between the two approaches is
10who authored the model.
12The last rule of the system prompt, answering only from the portfolio,
13statement, table or note the question names, exists because the context
14now reaches deep into the fused ranking and carries look-alikes: the same
15metric printed for another portfolio or period, a servicing-portfolio state
16breakdown beside the loan portfolio's, a December column beside March's. The
17ranker cannot tell those apart; it grades whether a chunk's words fit the
18item, not which portfolio the chunk belongs to. This model reads the whole
19context with its note titles and captions, so the choice sits here. Measured
20on the ABR Q1 2026 10-Q sweep of 59 questions: three answers taken from the
21servicing table when the question named the loan portfolio became correct,
22one refusal became a correct answer, one answer to a question naming no
23table in the filing became a refusal, the other 54 were unchanged.
24"""
26from __future__ import annotations
28from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Type, TypeVar
30from pydantic import BaseModel, Field
31from pydantic_ai import Agent, NativeOutput
32from pydantic_ai.settings import ModelSettings
34from quber.agents.langsmith_tracer import usage_metadata_from
35from quber.playground.agent import DEFAULT_MODEL, answer_model, build_context
36from quber.playground.answers.document import DocumentIdentity, with_document
37from quber.playground.retrieval import RetrievedChunk
38from quber.playground.tracing import run_metadata, tracer
40SYSTEM_PROMPT = """\
41You are a financial-document analyst answering a question about ONE parsed
42document, from context chunks retrieved from it.
44Answer ONLY from the provided context. Never use outside knowledge and never
45invent figures. If the context does not hold the answer, leave the answer
46fields empty and say why in `not_found`.
48Report figures exactly as the source prints them — currency sign, separators
49and all — and nothing else. "$14.47", never "$14.47 per share" and never a
50sentence. Qualifiers go in their own fields when the schema provides them.
52Populate `cited_ids` with the ids supporting your answer. Each context chunk
53begins with `[chunk id=<ID> | page <N> | type <TYPE>]`. Cite the chunk `<ID>`,
54or the finer-grained element ids inside a chunk's content, such as table cells
55tagged `<td id="t0-8-1">`. Only cite ids that appear verbatim in the context.
56Never invent or construct an id.
58When the question names a portfolio, statement, table or note, answer only
59from that one. The same item printed for a different portfolio, statement or
60table is not the answer: if the named one does not print the item, say so in
61`not_found` and do not substitute the other.
62"""
64M = TypeVar("M", bound=BaseModel)
67class DeclaredScalar(BaseModel):
68 """The shape a caller declares when it wants one figure back.
70 Carried here because it is the model the comparison hands to this path for
71 a value-seeking question, and the shape a src caller would most often want.
72 """
74 value: str = Field(
75 description=(
76 "The figure exactly as printed in the source, with its currency sign "
77 "and separators and nothing else: '$14.47', '4.6x', '(0.02)'."
78 )
79 )
80 unit: Optional[str] = Field(default=None, description="What the number counts.")
81 period: Optional[str] = Field(default=None, description="The period or as-of date, as printed.")
82 cited_ids: List[str] = Field(
83 default_factory=list, description="Chunk and/or cell ids supporting the value."
84 )
85 not_found: Optional[str] = Field(
86 default=None, description="Set only when the context lacks the answer; say what."
87 )
90class DeclaredProse(BaseModel):
91 """The shape a caller declares when it wants an explanation.
93 Sibling to `DeclaredScalar`: the two of them are what the playground's
94 expectation toggle picks between.
95 """
97 text: str = Field(description="The explanation, in prose, drawn only from the context.")
98 cited_ids: List[str] = Field(
99 default_factory=list, description="Chunk and/or cell ids supporting the explanation."
100 )
101 not_found: Optional[str] = Field(
102 default=None, description="Set only when the context lacks the answer; say what."
103 )
106def prompt_for(
107 question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None
108) -> str:
109 prompt = (
110 f"Context chunks:\n\n{build_context(chunks)}\n\n"
111 f"Question: {question}\n\n"
112 "Fill the answer fields from the context above."
113 )
114 return with_document(prompt, document)
117def build_agent(output_type: Type[M]) -> Agent[None, M]:
118 return Agent(
119 answer_model(),
120 output_type=NativeOutput(output_type),
121 system_prompt=SYSTEM_PROMPT,
122 model_settings=ModelSettings(temperature=0.0),
123 )
126def trace_inputs(prompt: str, chunks: List[RetrievedChunk]) -> Dict[str, Any]:
127 return {
128 "messages": [
129 {"role": "system", "content": SYSTEM_PROMPT},
130 {"role": "user", "content": prompt},
131 ],
132 "chunk_ids": [c.chunk_id for c in chunks],
133 }
136async def answer(
137 question: str,
138 chunks: List[RetrievedChunk],
139 output_type: Type[M],
140 document: Optional[DocumentIdentity] = None,
141) -> M:
142 prompt = prompt_for(question, chunks, document)
143 async with tracer().llm_run(
144 "answer_declared", trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata()
145 ) as run:
146 result = await build_agent(output_type).run(prompt)
147 output: M = result.output
148 run.outputs = {
149 "messages": [{"role": "assistant", "content": output.model_dump_json()}],
150 "usage_metadata": usage_metadata_from(result.usage),
151 }
152 return output
155async def answer_stream(
156 question: str,
157 chunks: List[RetrievedChunk],
158 output_type: Type[M],
159 text_field: str,
160 document: Optional[DocumentIdentity] = None,
161) -> AsyncIterator[Tuple[str, object]]:
162 """Yield ("delta", text) as one named field accumulates, then ("final", model).
164 `text_field` names the field worth watching arrive. A declared scalar has
165 none worth streaming, so callers wanting a figure skip this.
166 """
167 prompt = prompt_for(question, chunks, document)
168 async with tracer().llm_run(
169 "answer_declared", trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata()
170 ) as run:
171 async with build_agent(output_type).run_stream(prompt) as result:
172 last = ""
173 async for partial in result.stream_output(debounce_by=0.05):
174 text = getattr(partial, text_field, "") or ""
175 if text and text != last:
176 last = text
177 yield "delta", text
178 final: M = await result.get_output()
179 run.outputs = {
180 "messages": [{"role": "assistant", "content": final.model_dump_json()}],
181 "usage_metadata": usage_metadata_from(result.usage),
182 }
183 yield "final", final