Skip to content

Commit 4cb9025

Browse files
zastrowmdbschmigelski
authored andcommitted
Expand additional $refs for structured_output (strands-agents#439)
Addresses issue#337 Previously lists of items that were optional were not correctly expanding $refs. Derived classes also weren't having their $refs expanded as the subclass already had a "properties" object which bypassed $ref expansion
1 parent 5dcf6e1 commit 4cb9025

File tree

2 files changed

+124
-1
lines changed

2 files changed

+124
-1
lines changed

src/strands/tools/structured_output.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,9 @@ def _flatten_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
5454

5555
# Process each nested property
5656
for nested_prop_name, nested_prop_value in prop_value["properties"].items():
57-
processed_prop["properties"][nested_prop_name] = nested_prop_value
57+
is_required = "required" in prop_value and nested_prop_name in prop_value["required"]
58+
sub_property = _process_property(nested_prop_value, schema.get("$defs", {}), is_required)
59+
processed_prop["properties"][nested_prop_name] = sub_property
5860

5961
# Copy required fields if present
6062
if "required" in prop_value:
@@ -137,6 +139,10 @@ def _process_property(
137139
if "description" in prop:
138140
result["description"] = prop["description"]
139141

142+
# Need to process item refs as well (#337)
143+
if "items" in result:
144+
result["items"] = _process_property(result["items"], defs)
145+
140146
return result
141147

142148
# Handle direct references

tests/strands/tools/test_structured_output.py

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,3 +226,120 @@ class EmptyDocUser(BaseModel):
226226

227227
tool_spec = convert_pydantic_to_tool_spec(EmptyDocUser)
228228
assert tool_spec["description"] == "EmptyDocUser structured output tool"
229+
230+
231+
def test_convert_pydantic_with_items_refs():
232+
"""Test that no $refs exist after lists of different components."""
233+
234+
class Address(BaseModel):
235+
postal_code: Optional[str] = None
236+
237+
class Person(BaseModel):
238+
"""Complete person information."""
239+
240+
list_of_items: list[Address]
241+
list_of_items_nullable: Optional[list[Address]]
242+
list_of_item_or_nullable: list[Optional[Address]]
243+
244+
tool_spec = convert_pydantic_to_tool_spec(Person)
245+
246+
expected_spec = {
247+
"description": "Complete person information.",
248+
"inputSchema": {
249+
"json": {
250+
"description": "Complete person information.",
251+
"properties": {
252+
"list_of_item_or_nullable": {
253+
"items": {
254+
"anyOf": [
255+
{
256+
"properties": {"postal_code": {"type": ["string", "null"]}},
257+
"title": "Address",
258+
"type": "object",
259+
},
260+
{"type": "null"},
261+
]
262+
},
263+
"title": "List Of Item Or Nullable",
264+
"type": "array",
265+
},
266+
"list_of_items": {
267+
"items": {
268+
"properties": {"postal_code": {"type": ["string", "null"]}},
269+
"title": "Address",
270+
"type": "object",
271+
},
272+
"title": "List Of Items",
273+
"type": "array",
274+
},
275+
"list_of_items_nullable": {
276+
"items": {
277+
"properties": {"postal_code": {"type": ["string", "null"]}},
278+
"title": "Address",
279+
"type": "object",
280+
},
281+
"type": ["array", "null"],
282+
},
283+
},
284+
"required": ["list_of_items", "list_of_item_or_nullable"],
285+
"title": "Person",
286+
"type": "object",
287+
}
288+
},
289+
"name": "Person",
290+
}
291+
assert tool_spec == expected_spec
292+
293+
294+
def test_convert_pydantic_with_refs():
295+
"""Test that no $refs exist after processing complex hierarchies."""
296+
297+
class Address(BaseModel):
298+
street: str
299+
city: str
300+
country: str
301+
postal_code: Optional[str] = None
302+
303+
class Contact(BaseModel):
304+
address: Address
305+
306+
class Person(BaseModel):
307+
"""Complete person information."""
308+
309+
contact: Contact = Field(description="Contact methods")
310+
311+
tool_spec = convert_pydantic_to_tool_spec(Person)
312+
313+
expected_spec = {
314+
"description": "Complete person information.",
315+
"inputSchema": {
316+
"json": {
317+
"description": "Complete person information.",
318+
"properties": {
319+
"contact": {
320+
"description": "Contact methods",
321+
"properties": {
322+
"address": {
323+
"properties": {
324+
"city": {"title": "City", "type": "string"},
325+
"country": {"title": "Country", "type": "string"},
326+
"postal_code": {"type": ["string", "null"]},
327+
"street": {"title": "Street", "type": "string"},
328+
},
329+
"required": ["street", "city", "country"],
330+
"title": "Address",
331+
"type": "object",
332+
}
333+
},
334+
"required": ["address"],
335+
"type": "object",
336+
}
337+
},
338+
"required": ["contact"],
339+
"title": "Person",
340+
"type": "object",
341+
}
342+
},
343+
"name": "Person",
344+
}
345+
assert tool_spec == expected_spec

0 commit comments

Comments
 (0)