|
18 | 18 | from google.adk.evaluation.eval_case import IntermediateData, Invocation |
19 | 19 | from google.genai import types as genai_types |
20 | 20 |
|
21 | | -from .extraction import get_extractor, parse_json |
| 21 | +from .extraction import ( |
| 22 | + extract_agent_response_from_attrs, |
| 23 | + extract_user_text_from_attrs, |
| 24 | + get_extractor, |
| 25 | + parse_json, |
| 26 | +) |
22 | 27 | from .loader.base import Span, Trace |
23 | 28 | from .trace_attrs import ( |
24 | 29 | ADK_INVOCATION_ID, |
@@ -152,50 +157,34 @@ def _walk(span: Span, op_prefix: str, acc: list[Span]) -> None: |
152 | 157 |
|
153 | 158 |
|
154 | 159 | def _extract_user_content(first_call_llm: Span) -> genai_types.Content: |
155 | | - """Extract user input from the first call_llm span's llm_request tag.""" |
| 160 | + """Extract user input from the first call_llm span's attributes via shared extractor.""" |
| 161 | + text = extract_user_text_from_attrs(first_call_llm.tags) |
| 162 | + if text: |
| 163 | + return genai_types.Content( |
| 164 | + role="user", |
| 165 | + parts=[genai_types.Part(text=text)], |
| 166 | + ) |
156 | 167 | llm_request_raw = first_call_llm.get_tag(ADK_LLM_REQUEST, "{}") |
157 | 168 | llm_request = parse_json(llm_request_raw) |
158 | | - contents = llm_request.get("contents", []) |
159 | | - |
160 | | - for content_dict in reversed(contents): |
161 | | - if content_dict.get("role") != "user": |
162 | | - continue |
163 | | - parts = content_dict.get("parts", []) |
164 | | - # Skip function_response parts — only want actual user text messages |
165 | | - text_parts = [p for p in parts if "text" in p] |
166 | | - if text_parts: |
167 | | - return genai_types.Content( |
168 | | - role="user", |
169 | | - parts=[genai_types.Part(text=p["text"]) for p in text_parts], |
170 | | - ) |
171 | | - |
172 | | - for content_dict in contents: |
| 169 | + for content_dict in llm_request.get("contents", []): |
173 | 170 | if content_dict.get("role") == "user": |
174 | 171 | return _content_from_dict(content_dict) |
175 | | - |
176 | 172 | raise ValueError(f"call_llm span {first_call_llm.span_id}: no user content found in llm_request") |
177 | 173 |
|
178 | 174 |
|
179 | 175 | def _extract_final_response(last_call_llm: Span) -> genai_types.Content: |
180 | | - """Extract final text response from the last call_llm span's llm_response tag.""" |
| 176 | + """Extract final text response from the last call_llm span's attributes via shared extractor.""" |
| 177 | + text = extract_agent_response_from_attrs(last_call_llm.tags) |
| 178 | + if text: |
| 179 | + return genai_types.Content( |
| 180 | + role="model", |
| 181 | + parts=[genai_types.Part(text=text)], |
| 182 | + ) |
181 | 183 | llm_response_raw = last_call_llm.get_tag(ADK_LLM_RESPONSE, "{}") |
182 | 184 | llm_response = parse_json(llm_response_raw) |
183 | | - |
184 | 185 | content_dict = llm_response.get("content", {}) |
185 | 186 | if not content_dict: |
186 | 187 | raise ValueError(f"call_llm span {last_call_llm.span_id}: no content in llm_response") |
187 | | - |
188 | | - parts_dicts = content_dict.get("parts", []) |
189 | | - # Final response should have text parts, not function_call parts |
190 | | - text_parts = [p for p in parts_dicts if "text" in p] |
191 | | - if text_parts: |
192 | | - return genai_types.Content( |
193 | | - role="model", |
194 | | - parts=[genai_types.Part(text=p["text"]) for p in text_parts], |
195 | | - ) |
196 | | - |
197 | | - # If the last call_llm only has function_call parts, that's unexpected |
198 | | - # for a final response — the agent may have been cut short. |
199 | 188 | logger.warning( |
200 | 189 | "call_llm span %s: last llm_response has no text parts, may not be the actual final response", |
201 | 190 | last_call_llm.span_id, |
|
0 commit comments