Coverage for src / quber / playground / answers / planned.py: 82%

1"""Answering into a schema built for the question that was asked. 

2 

3Two calls. A planner reads the question alone — never the context, never the 

4answer — and declares the fields the answer should have. Those declarations 

5are compiled into a Pydantic model, and the answering agent is constrained to 

6it through the same call `declared` uses. The only difference between the two 

7approaches is who authored the model. 

8 

9The planner emits field declarations rather than raw JSON Schema. A model 

10writing free-form JSON Schema produces dialects the validator then rejects, 

11and that surfaces as a retry loop rather than an answer; a closed set of field 

12types cannot fail that way. 

13""" 

14 

15from __future__ import annotations 

16 

17from functools import lru_cache 

18from typing import Any, Dict, List, Literal, Optional, Tuple, Type 

19 

20from pydantic import BaseModel, Field, create_model 

21from pydantic_ai import Agent, NativeOutput 

22from pydantic_ai.settings import ModelSettings 

23 

24from quber.agents.langsmith_tracer import usage_metadata_from 

25from quber.playground.agent import anthropic_model 

26from quber.playground.tracing import run_metadata, tracer 

27 

28PLANNER_MODEL = "claude-haiku-4-5" 

29 

30FieldType = Literal["string", "number", "string_list", "number_list"] 

31 

32_PY_TYPES: Dict[str, Any] = { 

33 "string": str, 

34 "number": float, 

35 "string_list": List[str], 

36 "number_list": List[float], 

37} 

38 

39PLANNER_PROMPT = """\ 

40You design the return shape for a financial-document question. You are given 

41the QUESTION only — never the document, never the answer. Do not attempt to 

42answer it. 

43 

44Decide first whether the question asks for data or for explanation. 

45 

46- If it asks you to explain, compare, summarize, or reason, set `wants_value` 

47 false and return no fields. It will be answered in prose. 

48- Otherwise set `wants_value` true and declare the fields that hold the 

49 answer, and only those. 

50 

51Rules for fields: 

52- Declare the narrowest set that answers the question. A question naming one 

53 figure gets ONE field holding that figure, plus at most the qualifiers the 

54 question itself names. 

55- Name fields in snake_case after what they hold. 

56- Each field's description must tell the answering agent to report the figure 

57 exactly as printed in the source, with no surrounding words. 

58- Never declare a field for citations or sources; those are added for you. 

59- Never declare a field whose value would be a sentence. 

60""" 

61 

62 

63class FieldSpec(BaseModel): 

64 name: str = Field(description="snake_case field name.") 

65 type: FieldType = Field(description="The field's type.") 

66 description: str = Field(description="What the answering agent must put here.") 

67 required: bool = True 

68 

69 

70class Plan(BaseModel): 

71 wants_value: bool = Field( 

72 description="True when the question asks for data; false when it asks for explanation." 

73 ) 

74 fields: List[FieldSpec] = Field( 

75 default_factory=list, description="The answer's fields. Empty when wants_value is false." 

76 ) 

77 

78 

79@lru_cache(maxsize=1) 

80def _planner() -> Agent[None, Plan]: 

81 return Agent( 

82 anthropic_model(PLANNER_MODEL), 

83 output_type=NativeOutput(Plan), 

84 system_prompt=PLANNER_PROMPT, 

85 model_settings=ModelSettings(temperature=0.0), 

86 ) 

87 

88 

89async def plan_for(question: str) -> Plan: 

90 prompt = f"QUESTION: {question}" 

91 inputs = { 

92 "messages": [ 

93 {"role": "system", "content": PLANNER_PROMPT}, 

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

95 ] 

96 } 

97 async with tracer().llm_run( 

98 "plan_answer", inputs, model=PLANNER_MODEL, extra_metadata=run_metadata() 

99 ) as run: 

100 result = await _planner().run(prompt) 

101 output: Plan = result.output 

102 run.outputs = { 

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

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

105 } 

106 return output 

107 

108 

109def compile_model(plan: Plan) -> Type[BaseModel]: 

110 """Turn a plan's field declarations into a model the answer agent fills. 

111 

112 `cited_ids` and `not_found` are appended to every compiled model, so 

113 grounding and the not-present case survive whatever the planner declared. 

114 A plan that declared a field by either of those names loses it to the 

115 appended one, which is the intended precedence. 

116 """ 

117 fields: Dict[str, Tuple[Any, Any]] = {} 

118 for spec in plan.fields: 

119 if spec.name in ("cited_ids", "not_found"): 

120 continue 

121 py = _PY_TYPES[spec.type] 

122 if spec.required: 

123 fields[spec.name] = (py, Field(description=spec.description)) 

124 else: 

125 fields[spec.name] = (Optional[py], Field(default=None, description=spec.description)) 

126 fields["cited_ids"] = ( 

127 List[str], 

128 Field(default_factory=list, description="Chunk and/or cell ids supporting the answer."), 

129 ) 

130 fields["not_found"] = ( 

131 Optional[str], 

132 Field(default=None, description="Set only when the context lacks the answer; say what."), 

133 ) 

134 return create_model("PlannedAnswer", **fields) # type: ignore[call-overload]