Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion src/strands/tools/structured_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,9 @@ def _flatten_schema(schema: Dict[str, Any]) -> Dict[str, Any]:

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

# Copy required fields if present
if "required" in prop_value:
Expand Down Expand Up @@ -137,6 +139,10 @@ def _process_property(
if "description" in prop:
result["description"] = prop["description"]

# Need to process item refs as well (#337)
if "items" in result:
result["items"] = _process_property(result["items"], defs)

return result

# Handle direct references
Expand Down
117 changes: 117 additions & 0 deletions tests/strands/tools/test_structured_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -226,3 +226,120 @@ class EmptyDocUser(BaseModel):

tool_spec = convert_pydantic_to_tool_spec(EmptyDocUser)
assert tool_spec["description"] == "EmptyDocUser structured output tool"


def test_convert_pydantic_with_items_refs():
"""Test that no $refs exist after lists of different components."""

class Address(BaseModel):
postal_code: Optional[str] = None

class Person(BaseModel):
"""Complete person information."""

list_of_items: list[Address]
list_of_items_nullable: Optional[list[Address]]
list_of_item_or_nullable: list[Optional[Address]]

tool_spec = convert_pydantic_to_tool_spec(Person)

expected_spec = {
"description": "Complete person information.",
"inputSchema": {
"json": {
"description": "Complete person information.",
"properties": {
"list_of_item_or_nullable": {
"items": {
"anyOf": [
{
"properties": {"postal_code": {"type": ["string", "null"]}},
"title": "Address",
"type": "object",
},
{"type": "null"},
]
},
"title": "List Of Item Or Nullable",
"type": "array",
},
"list_of_items": {
"items": {
"properties": {"postal_code": {"type": ["string", "null"]}},
"title": "Address",
"type": "object",
},
"title": "List Of Items",
"type": "array",
},
"list_of_items_nullable": {
"items": {
"properties": {"postal_code": {"type": ["string", "null"]}},
"title": "Address",
"type": "object",
},
"type": ["array", "null"],
},
},
"required": ["list_of_items", "list_of_item_or_nullable"],
"title": "Person",
"type": "object",
}
},
"name": "Person",
}
assert tool_spec == expected_spec


def test_convert_pydantic_with_refs():
"""Test that no $refs exist after processing complex hierarchies."""

class Address(BaseModel):
street: str
city: str
country: str
postal_code: Optional[str] = None

class Contact(BaseModel):
address: Address

class Person(BaseModel):
"""Complete person information."""

contact: Contact = Field(description="Contact methods")

tool_spec = convert_pydantic_to_tool_spec(Person)

expected_spec = {
"description": "Complete person information.",
"inputSchema": {
"json": {
"description": "Complete person information.",
"properties": {
"contact": {
"description": "Contact methods",
"properties": {
"address": {
"properties": {
"city": {"title": "City", "type": "string"},
"country": {"title": "Country", "type": "string"},
"postal_code": {"type": ["string", "null"]},
"street": {"title": "Street", "type": "string"},
},
"required": ["street", "city", "country"],
"title": "Address",
"type": "object",
}
},
"required": ["address"],
"type": "object",
}
},
"required": ["contact"],
"title": "Person",
"type": "object",
}
},
"name": "Person",
}
assert tool_spec == expected_spec
Loading