Coverage for src / quber / playground / agent.py: 57%

1"""Grounded RAG agent for the ADE playground. 

2 

3Runs the best available Claude model through Pydantic AI, 

4authenticated with the Claude subscription OAuth token via the project's 

5`_oauth_gate` helper (which rewrites the system prompt into the 2-block 

6array format the subscription gate requires). 

7 

8The agent answers strictly from retrieved ADE chunks and returns the ids it 

9grounded on. Because table chunk markdown carries per-cell ids 

10(`<td id="0-8">...`), the model can cite individual cells, which the app 

11resolves to page + bounding-box overlays — the "visual reference for the 

12answer" behaviour from the Landing.ai playground. 

13""" 

14 

15from __future__ import annotations 

16 

17from functools import lru_cache 

18from typing import Any, Dict, List, Optional 

19 

20from pydantic import BaseModel, Field 

21from pydantic_ai import Agent 

22from pydantic_ai.models.anthropic import AnthropicModel 

23from pydantic_ai.providers.anthropic import AnthropicProvider 

24from pydantic_ai.settings import ModelSettings 

25 

26from quber.agents._oauth_gate import make_oauth_anthropic_model 

27from quber.agents.langsmith_tracer import usage_metadata_from 

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

29from quber.playground.retrieval import RetrievedChunk 

30from quber.playground.tracing import run_metadata, tracer 

31from quber.settings import get_settings 

32 

33# Best available Claude model for the chat/answer step. 

34DEFAULT_MODEL = "claude-opus-5" 

35 

36 

37def anthropic_model(model_id: str): 

38 """An Anthropic model over whichever credential the host presents. 

39 

40 The subscription OAuth token wins when it is set, routed through the OAuth 

41 gate as before. Without one the API key is used through the stock model, 

42 which is how the hosted playground runs: the token never leaves the 

43 developer host, and the task is given only the key. Neither set is a 

44 configuration error, reported as such. 

45 """ 

46 llm = get_settings().llm 

47 if llm.anthropic_auth_token: 

48 return make_oauth_anthropic_model(model_id, llm.anthropic_auth_token) 

49 if llm.anthropic_api_key: 

50 return AnthropicModel(model_id, provider=AnthropicProvider(api_key=llm.anthropic_api_key)) 

51 raise RuntimeError( 

52 "Neither ANTHROPIC_AUTH_TOKEN nor ANTHROPIC_API_KEY is set; the playground needs one of them." 

53 ) 

54 

55 

56def answer_model(): 

57 """The playground's answering model. 

58 

59 One place decides which model answers; the three answer paths all build 

60 from here rather than each pinning its own copy of the name. 

61 """ 

62 return anthropic_model(DEFAULT_MODEL) 

63 

64 

65SYSTEM_PROMPT = """\ 

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

67document. You are given a set of context chunks retrieved from that document. 

68 

69Rules: 

70- Answer ONLY from the provided context chunks. Never use outside knowledge 

71 and never invent figures. If the context does not contain the answer, say 

72 so plainly. 

73- Be precise and quantitative. Report exact values, units, and periods as 

74 they appear in the source. 

75- Populate `cited_ids` with the ids that support your answer. Each context 

76 chunk begins with a header line: `[chunk id=<ID> | page <N> | type <TYPE>]`. 

77 * By default, cite the `<ID>` from the header of each chunk you used. 

78 * If a chunk's content exposes finer-grained element ids — e.g. table cells 

79 tagged like `<td id="0-8">` — you MAY cite those specific cell ids instead, 

80 for tighter grounding. 

81 * Only cite ids that literally appear in the provided context (a header 

82 `<ID>` or an element id inside a chunk's content). NEVER invent, guess, or 

83 construct an id that is not present verbatim. 

84""" 

85 

86 

87class GroundedAnswer(BaseModel): 

88 answer: str = Field(description="The answer, drawn only from the context chunks.") 

89 cited_ids: List[str] = Field( 

90 default_factory=list, 

91 description="Chunk ids and/or table cell ids that support the answer.", 

92 ) 

93 

94 

95@lru_cache(maxsize=1) 

96def _agent() -> Agent[None, GroundedAnswer]: 

97 return Agent( 

98 answer_model(), 

99 output_type=GroundedAnswer, 

100 system_prompt=SYSTEM_PROMPT, 

101 model_settings=ModelSettings(temperature=0.0), 

102 ) 

103 

104 

105def build_context(chunks: List[RetrievedChunk]) -> str: 

106 blocks = [] 

107 for c in chunks: 

108 # ADE pages are 0-based; show 1-based to match the PDF viewer. 

109 header = f"[chunk id={c.chunk_id} | page {c.page + 1} | type {c.chunk_type}]" 

110 blocks.append(f"{header}\n{c.content}") 

111 return "\n\n".join(blocks) 

112 

113 

114def _prompt(question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None) -> str: 

115 context = build_context(chunks) 

116 prompt = ( 

117 f"Context chunks:\n\n{context}\n\n" 

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

119 "Answer from the context above and cite the ids you used." 

120 ) 

121 return with_document(prompt, document) 

122 

123 

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

125 return { 

126 "messages": [ 

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

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

129 ], 

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

131 } 

132 

133 

134async def answer( 

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

136) -> GroundedAnswer: 

137 prompt = _prompt(question, chunks, document) 

138 async with tracer().llm_run( 

139 "answer_grounded", _trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata() 

140 ) as run: 

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

142 output: GroundedAnswer = result.output 

143 run.outputs = { 

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

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

146 } 

147 return output 

148 

149 

150async def answer_stream( 

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

152): 

153 """Yield ("delta", <answer so far>) as the model generates, then 

154 ("final", GroundedAnswer) once the full output validates.""" 

155 prompt = _prompt(question, chunks, document) 

156 async with tracer().llm_run( 

157 "answer_grounded", _trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata() 

158 ) as run: 

159 async with _agent().run_stream(prompt) as result: 

160 last = "" 

161 async for partial in result.stream_output(debounce_by=0.05): 

162 text = getattr(partial, "answer", "") or "" 

163 if text and text != last: 

164 last = text 

165 yield "delta", text 

166 final: GroundedAnswer = await result.get_output() 

167 run.outputs = { 

168 "messages": [{"role": "assistant", "content": final.model_dump_json()}], 

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

170 } 

171 yield "final", final