[v1.x] Give recursive tool return types an object-rooted output schema (#3377)
Some checks failed
Main branch checks / checks (push) Has been cancelled
Some checks failed
Main branch checks / checks (push) Has been cancelled
This commit is contained in:
parent
4dc224fab2
commit
b22271305c
3 changed files with 71 additions and 1 deletions
|
|
@ -45,6 +45,25 @@ class StrictJsonSchema(GenerateJsonSchema):
|
||||||
raise ValueError(f"JSON schema warning: {kind} - {detail}")
|
raise ValueError(f"JSON schema warning: {kind} - {detail}")
|
||||||
|
|
||||||
|
|
||||||
|
_LOCAL_DEFS_PREFIX = "#/$defs/"
|
||||||
|
|
||||||
|
|
||||||
|
def _inline_root_ref(schema: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
"""Give a schema whose root is a bare `$ref` into `$defs` an inline root.
|
||||||
|
|
||||||
|
pydantic emits a self-referential model as `{"$defs": {...}, "$ref": "#/$defs/Model"}`, with no
|
||||||
|
`type` at the root; `Tool.outputSchema` requires `type: object` at the root. The referenced
|
||||||
|
definition is copied onto the root and `$defs` is kept, since nested references still point into
|
||||||
|
it. Root siblings of the `$ref` win over the definition's keys.
|
||||||
|
"""
|
||||||
|
ref = schema.get("$ref")
|
||||||
|
if not isinstance(ref, str) or not ref.startswith(_LOCAL_DEFS_PREFIX):
|
||||||
|
return schema
|
||||||
|
definition = cast(dict[str, Any], schema["$defs"][ref.removeprefix(_LOCAL_DEFS_PREFIX)])
|
||||||
|
siblings = {key: value for key, value in schema.items() if key != "$ref"}
|
||||||
|
return {**definition, **siblings}
|
||||||
|
|
||||||
|
|
||||||
class ArgModelBase(BaseModel):
|
class ArgModelBase(BaseModel):
|
||||||
"""A model representing the arguments to a function."""
|
"""A model representing the arguments to a function."""
|
||||||
|
|
||||||
|
|
@ -428,7 +447,7 @@ def _try_create_model_and_schema(
|
||||||
logger.info(f"Cannot create schema for type {type_expr} in {func_name}: {type(e).__name__}: {e}")
|
logger.info(f"Cannot create schema for type {type_expr} in {func_name}: {type(e).__name__}: {e}")
|
||||||
return None, None, False
|
return None, None, False
|
||||||
|
|
||||||
return model, schema, wrap_output
|
return model, _inline_root_ref(schema), wrap_output
|
||||||
|
|
||||||
return None, None, False
|
return None, None, False
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -983,6 +983,33 @@ def test_structured_output_nested_models():
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_structured_output_self_referential_model_gets_an_object_root():
|
||||||
|
"""pydantic publishes a recursive model as a bare root `$ref`; the definition is inlined onto the
|
||||||
|
root and `$defs` is kept for the nested reference."""
|
||||||
|
|
||||||
|
class Node(BaseModel):
|
||||||
|
name: str
|
||||||
|
children: list["Node"] = []
|
||||||
|
|
||||||
|
def tree() -> Node:
|
||||||
|
return Node(name="root", children=[Node(name="leaf")])
|
||||||
|
|
||||||
|
node_definition: dict[str, Any] = {
|
||||||
|
"properties": {
|
||||||
|
"name": {"title": "Name", "type": "string"},
|
||||||
|
"children": {"default": [], "items": {"$ref": "#/$defs/Node"}, "title": "Children", "type": "array"},
|
||||||
|
},
|
||||||
|
"required": ["name"],
|
||||||
|
"title": "Node",
|
||||||
|
"type": "object",
|
||||||
|
}
|
||||||
|
meta = func_metadata(tree)
|
||||||
|
assert meta.output_schema == {**node_definition, "$defs": {"Node": node_definition}}
|
||||||
|
|
||||||
|
_, structured_content = meta.convert_result(tree())
|
||||||
|
assert structured_content == {"name": "root", "children": [{"name": "leaf", "children": []}]}
|
||||||
|
|
||||||
|
|
||||||
def test_structured_output_unserializable_type_error():
|
def test_structured_output_unserializable_type_error():
|
||||||
"""Test error when structured_output=True is used with unserializable types"""
|
"""Test error when structured_output=True is used with unserializable types"""
|
||||||
from typing import NamedTuple
|
from typing import NamedTuple
|
||||||
|
|
|
||||||
|
|
@ -531,6 +531,30 @@ class TestServerTools:
|
||||||
assert isinstance(result.content[0], TextContent)
|
assert isinstance(result.content[0], TextContent)
|
||||||
assert '"name": "John Doe"' in result.content[0].text
|
assert '"name": "John Doe"' in result.content[0].text
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_tool_structured_output_self_referential_model(self):
|
||||||
|
"""A self-referential return type publishes an object-rooted outputSchema (required by the
|
||||||
|
2025-11-25 Tool shape) and its result validates client-side through the kept `$defs`."""
|
||||||
|
|
||||||
|
class Node(BaseModel):
|
||||||
|
name: str
|
||||||
|
children: list["Node"] = []
|
||||||
|
|
||||||
|
def tree() -> Node:
|
||||||
|
return Node(name="root", children=[Node(name="leaf")])
|
||||||
|
|
||||||
|
mcp = FastMCP()
|
||||||
|
mcp.add_tool(tree)
|
||||||
|
|
||||||
|
async with client_session(mcp._mcp_server) as client:
|
||||||
|
[tool] = (await client.list_tools()).tools
|
||||||
|
assert tool.outputSchema is not None
|
||||||
|
assert tool.outputSchema["type"] == "object"
|
||||||
|
assert tool.outputSchema["properties"]["children"]["items"] == {"$ref": "#/$defs/Node"}
|
||||||
|
result = await client.call_tool("tree", {})
|
||||||
|
assert result.isError is False
|
||||||
|
assert result.structuredContent == {"name": "root", "children": [{"name": "leaf", "children": []}]}
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_tool_structured_output_primitive(self):
|
async def test_tool_structured_output_primitive(self):
|
||||||
"""Test tool with structured output returning primitive type"""
|
"""Test tool with structured output returning primitive type"""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue