Skip to content

Commit dd97e76

Browse files
Jacksunweicopybara-github
authored andcommitted
fix(conformance): normalize tool declarations in replay verification
Co-authored-by: Wei Sun (Jack) <weisun@google.com> PiperOrigin-RevId: 931479521
1 parent 2b8c80c commit dd97e76

2 files changed

Lines changed: 111 additions & 2 deletions

File tree

src/google/adk/cli/conformance/_conformance_test_google_llm.py

Lines changed: 110 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,12 +33,117 @@ class ReplayVerificationError(Exception):
3333
"""Exception raised when replay verification fails."""
3434

3535

36+
def _normalize_type(val: Any) -> Any:
37+
if hasattr(val, 'name') and hasattr(val, 'value'):
38+
return str(val.value).lower()
39+
if isinstance(val, str) and val.startswith('Type.'):
40+
return val.split('.')[-1].lower()
41+
if isinstance(val, str) and val in (
42+
'STRING',
43+
'NUMBER',
44+
'OBJECT',
45+
'ARRAY',
46+
'INTEGER',
47+
'BOOLEAN',
48+
):
49+
return val.lower()
50+
return val
51+
52+
53+
def _resolve_refs(data: Any, defs: dict[str, Any]) -> Any:
54+
if isinstance(data, dict):
55+
if '$ref' in data:
56+
ref_path = data['$ref']
57+
if ref_path.startswith('#/$defs/'):
58+
def_name = ref_path.split('/')[-1]
59+
if def_name in defs:
60+
return _resolve_refs(defs[def_name], defs)
61+
return {k: _resolve_refs(v, defs) for k, v in data.items()}
62+
elif isinstance(data, list):
63+
return [_resolve_refs(x, defs) for x in data]
64+
else:
65+
return data
66+
67+
68+
def _normalize_schema_dict(data: Any) -> Any:
69+
if isinstance(data, dict):
70+
if '$defs' in data:
71+
defs = data['$defs']
72+
data = _resolve_refs(data, defs)
73+
data.pop('$defs', None)
74+
75+
res = {}
76+
for k, v in data.items():
77+
if k in ('title', 'default', 'description'):
78+
continue
79+
if k == 'type':
80+
res[k] = _normalize_type(v)
81+
else:
82+
res[k] = _normalize_schema_dict(v)
83+
84+
if 'anyOf' in res and isinstance(res['anyOf'], list):
85+
any_of = res['anyOf']
86+
null_schema = None
87+
non_null_schemas = []
88+
for s in any_of:
89+
if isinstance(s, dict) and s.get('type') == 'null':
90+
null_schema = s
91+
else:
92+
non_null_schemas.append(s)
93+
94+
if null_schema is not None and len(non_null_schemas) == 1:
95+
target_schema = non_null_schemas[0]
96+
if isinstance(target_schema, dict):
97+
res.update(target_schema)
98+
res['nullable'] = True
99+
res.pop('anyOf', None)
100+
101+
return res
102+
elif isinstance(data, list):
103+
return [_normalize_schema_dict(x) for x in data]
104+
else:
105+
return data
106+
107+
108+
def _normalize_tool_config(data: Any) -> Any:
109+
"""Normalize function declarations to ignore minor formatting changes."""
110+
if isinstance(data, dict):
111+
if 'name' in data and (
112+
'description' in data
113+
or 'parameters' in data
114+
or 'parameters_json_schema' in data
115+
):
116+
if data.get('name') == 'transfer_to_agent':
117+
data['description'] = 'Transfer the question to another agent.'
118+
elif 'description' in data and isinstance(data['description'], str):
119+
data['description'] = data['description'].strip()
120+
121+
params = data.pop('parameters', None)
122+
if params is not None:
123+
data['parameters_json_schema'] = params
124+
125+
if 'parameters_json_schema' in data:
126+
data['parameters_json_schema'] = _normalize_schema_dict(
127+
data['parameters_json_schema']
128+
)
129+
130+
data.pop('response', None)
131+
data.pop('response_json_schema', None)
132+
133+
return {k: _normalize_tool_config(v) for k, v in data.items()}
134+
elif isinstance(data, list):
135+
return [_normalize_tool_config(x) for x in data]
136+
else:
137+
return data
138+
139+
36140
class _ConformanceTestGemini(Gemini):
37141
"""A mocked Gemini model for conformance test replay mode.
38142
39143
This class is used to mock the Gemini model in conformance test replay mode.
40-
It is a subclass of Gemini and overrides the `generate_content_async`` method to
41-
return a mocked response from the provided recordingss.
144+
It is a subclass of Gemini and overrides the `generate_content_async` method
145+
to
146+
return a mocked response from the provided recordings.
42147
"""
43148

44149
def __init__(
@@ -114,6 +219,9 @@ def _verify_llm_request_match(
114219
exclude_none=True, exclude=excluded_fields, exclude_defaults=True
115220
)
116221

222+
recorded_dict = _normalize_tool_config(recorded_dict)
223+
current_dict = _normalize_tool_config(current_dict)
224+
117225
if recorded_dict != current_dict:
118226
raise ReplayVerificationError(
119227
f"""LLM request mismatch in turn {self._user_message_index} for agent '{self._agent_name}' (index {replay_index}):

src/google/adk/cli/conformance/_replay_validators.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ def _compare_event(
8080
"timestamp": True,
8181
"invocation_id": True,
8282
"long_running_tool_ids": True,
83+
"node_info": True,
8384
# Content fields that vary per run
8485
"content": {
8586
"parts": {

0 commit comments

Comments
 (0)