Coverage for src / quber / playground / answers / fixed.py: 70%

1"""Answering into a fixed envelope whose payload is a closed set of shapes. 

2 

3The agent always returns `Answer`. It chooses the payload — scalar, series, 

4grid, prose, unanswerable — from the union, so deciding whether the question 

5wants a figure or an explanation happens inside the same call that answers it. 

6No separate classification step, and no schema is invented at runtime: every 

7shape a caller can receive is written down in `shapes`. 

8""" 

9 

10from __future__ import annotations 

11 

12from functools import lru_cache 

13from typing import Any, Dict, List, Optional 

14 

15from pydantic_ai import Agent, NativeOutput 

16from pydantic_ai.settings import ModelSettings 

17 

18from quber.agents.langsmith_tracer import usage_metadata_from 

19from quber.playground.agent import DEFAULT_MODEL, answer_model, build_context 

20from quber.playground.answers.document import DocumentIdentity, with_document 

21from quber.playground.answers.shapes import Answer 

22from quber.playground.retrieval import RetrievedChunk 

23from quber.playground.tracing import run_metadata, tracer 

24 

25SYSTEM_PROMPT = """\ 

26You are a financial-document analyst answering questions about ONE parsed 

27document, from a set of context chunks retrieved from it. 

28 

29Answer ONLY from the provided context. Never use outside knowledge, never 

30invent figures. If the context does not hold the answer, return the 

31`unanswerable` payload. 

32 

33CHOOSING THE PAYLOAD 

34 

35Read what the question asks for and return that shape: 

36 

37- `scalar` — the question names ONE figure. This is the common case for 

38 questions phrased like "Return the Q1 2026 value" or "What was X as of 

39 <date>". Put the figure in `value` exactly as the source prints it, with 

40 its currency sign and separators, and NOTHING else. "$14.47" — not 

41 "$14.47 per share", not "Book value was $14.47", not "$14.47 as of 

42 March 31, 2026". The qualifiers belong in `unit`, `period`, and `label`. 

43- `series` — the question asks for one measure across several periods, 

44 segments, or categories. 

45- `grid` — the question asks for a whole table. 

46- `prose` — the question asks you to explain, compare, summarize, or 

47 reason. Only here may you write sentences. 

48- `unanswerable` — the context does not contain what was asked. 

49 

50When a question could be read either way, prefer the narrower shape: if a 

51single figure answers it, return `scalar`. 

52 

53GROUNDING 

54 

55Populate `cited_ids` with the ids supporting your answer. Each context chunk 

56begins with `[chunk id=<ID> | page <N> | type <TYPE>]`. Cite the chunk `<ID>`, 

57or — for tighter grounding — the finer-grained element ids exposed inside a 

58chunk's content, such as table cells tagged `<td id="t0-8-1">`. On a `scalar`, 

59also set `source_id` to the single cell or chunk the figure was read from, and 

60make `value` character-for-character what that cell prints. 

61 

62Only cite ids that appear verbatim in the context. Never invent or construct 

63an id. 

64""" 

65 

66 

67@lru_cache(maxsize=1) 

68def build_agent() -> Agent[None, Answer]: 

69 return Agent( 

70 answer_model(), 

71 output_type=NativeOutput(Answer), 

72 system_prompt=SYSTEM_PROMPT, 

73 model_settings=ModelSettings(temperature=0.0), 

74 ) 

75 

76 

77def prompt_for( 

78 question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None 

79) -> str: 

80 prompt = ( 

81 f"Context chunks:\n\n{build_context(chunks)}\n\n" 

82 f"Question: {question}\n\n" 

83 "Answer from the context above, choosing the payload shape the question asks for." 

84 ) 

85 return with_document(prompt, document) 

86 

87 

88def trace_inputs(prompt: str, chunks: List[RetrievedChunk]) -> Dict[str, Any]: 

89 return { 

90 "messages": [ 

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

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

93 ], 

94 "chunk_ids": [c.chunk_id for c in chunks], 

95 } 

96 

97 

98async def answer( 

99 question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None 

100) -> Answer: 

101 prompt = prompt_for(question, chunks, document) 

102 async with tracer().llm_run( 

103 "answer_fixed", trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata() 

104 ) as run: 

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

106 output: Answer = result.output 

107 run.outputs = { 

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

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

110 } 

111 return output