From 9e48c5ea05056b3433827391de0b17865f89dc2e Mon Sep 17 00:00:00 2001 From: Lucas Messenger Date: Fri, 10 Apr 2026 11:14:59 -0700 Subject: [PATCH] feat(bedrock): add native structured output support via outputConfig.textFormat Add opt-in native structured output mode for BedrockModel that uses Bedrock's outputConfig.textFormat API for schema-constrained responses, replacing the tool-based workaround when enabled. Model-level: - Add `structured_output_mode` config ("tool" | "native", defaults to "tool") - Add `convert_pydantic_to_json_schema()` utility with recursive `additionalProperties: false` injection - Thread `output_config` through stream() -> _stream() -> _format_request() - Native mode parses JSON text response instead of extracting tool use args Agent-level: - When native_mode is enabled, the agent loop runs normally with tools and thinking. On end_turn, calls model.structured_output() for final formatting instead of forcing tool use (which disables thinking). Closes strands-agents/sdk-python#1652 --- src/strands/agent/agent.py | 7 +- src/strands/event_loop/event_loop.py | 67 +++++++--- src/strands/models/bedrock.py | 84 ++++++++++++- src/strands/tools/__init__.py | 3 +- .../tools/structured_output/__init__.py | 4 +- .../_structured_output_context.py | 7 +- .../structured_output_utils.py | 94 +++++++++++--- tests/strands/models/test_bedrock.py | 117 ++++++++++++++++++ tests/strands/tools/test_structured_output.py | 68 +++++++++- 9 files changed, 409 insertions(+), 42 deletions(-) diff --git a/src/strands/agent/agent.py b/src/strands/agent/agent.py index 3a23133de..f9c247251 100644 --- a/src/strands/agent/agent.py +++ b/src/strands/agent/agent.py @@ -886,9 +886,14 @@ async def _run_loop( await self._append_messages(*current_messages) + # Check if the model supports native structured output + model_config = self.model.get_config() + native_mode = isinstance(model_config, dict) and model_config.get("structured_output_mode") == "native" + structured_output_context = StructuredOutputContext( structured_output_model or self._default_structured_output_model, structured_output_prompt=structured_output_prompt or self._structured_output_prompt, + native_mode=native_mode, ) # Execute the event loop cycle with retry logic for context limits @@ -950,7 +955,7 @@ async def _execute_event_loop_cycle( # Add `Agent` to invocation_state to keep backwards-compatibility invocation_state["agent"] = self - if structured_output_context: + if structured_output_context and not structured_output_context.native_mode: structured_output_context.register_tool(self.tool_registry) try: diff --git a/src/strands/event_loop/event_loop.py b/src/strands/event_loop/event_loop.py index b4af16058..ba4e0ed97 100644 --- a/src/strands/event_loop/event_loop.py +++ b/src/strands/event_loop/event_loop.py @@ -201,25 +201,62 @@ async def event_loop_cycle( # End the cycle and return results agent.event_loop_metrics.end_cycle(cycle_start_time, cycle_trace, attributes) - # Force structured output tool call if LLM didn't use it automatically + # Handle structured output when model returns end_turn if structured_output_context.is_enabled and stop_reason == "end_turn": - if structured_output_context.force_attempted: + if structured_output_context.native_mode: + # Native mode: use model's native structured output for final formatting. + # The agent loop ran normally with tools and thinking; only this final + # step uses native structured output via model.structured_output(). + logger.debug("using native structured output for final formatting") + # Append a user message so the conversation doesn't end with an assistant + # message (some models like Opus 4.6 don't support assistant prefill). + await agent._append_messages( + {"role": "user", "content": [{"text": structured_output_context.structured_output_prompt}]} + ) + native_result = None + async for event in agent.model.structured_output( + structured_output_context.structured_output_model, + agent.messages, + system_prompt=agent.system_prompt, + ): + if "output" in event: + native_result = event["output"] + + if native_result is not None: + yield StructuredOutputEvent(structured_output=native_result) + tracer.end_event_loop_cycle_span(cycle_span, message) + yield EventLoopStopEvent( + stop_reason, + message, + agent.event_loop_metrics, + invocation_state["request_state"], + structured_output=native_result, + ) + return raise StructuredOutputException( - "The model failed to invoke the structured output tool even after it was forced." + "Native structured output mode: model did not return structured output." + ) + else: + # Tool mode: force the model to call the structured output tool + if structured_output_context.force_attempted: + raise StructuredOutputException( + "The model failed to invoke the structured output tool even after it was forced." + ) + structured_output_context.set_forced_mode() + logger.debug("Forcing structured output tool") + await agent._append_messages( + {"role": "user", "content": [{"text": structured_output_context.structured_output_prompt}]} ) - structured_output_context.set_forced_mode() - logger.debug("Forcing structured output tool") - await agent._append_messages( - {"role": "user", "content": [{"text": structured_output_context.structured_output_prompt}]} - ) - tracer.end_event_loop_cycle_span(cycle_span, message) - events = recurse_event_loop( - agent=agent, invocation_state=invocation_state, structured_output_context=structured_output_context - ) - async for typed_event in events: - yield typed_event - return + tracer.end_event_loop_cycle_span(cycle_span, message) + events = recurse_event_loop( + agent=agent, + invocation_state=invocation_state, + structured_output_context=structured_output_context, + ) + async for typed_event in events: + yield typed_event + return tracer.end_event_loop_cycle_span(cycle_span, message) yield EventLoopStopEvent(stop_reason, message, agent.event_loop_metrics, invocation_state["request_state"]) diff --git a/src/strands/models/bedrock.py b/src/strands/models/bedrock.py index 5b7a2f34e..0a4c2021b 100644 --- a/src/strands/models/bedrock.py +++ b/src/strands/models/bedrock.py @@ -21,7 +21,7 @@ from .._exception_notes import add_exception_note from ..event_loop import streaming -from ..tools import convert_pydantic_to_tool_spec +from ..tools import convert_pydantic_to_json_schema, convert_pydantic_to_tool_spec from ..tools._tool_helpers import noop_tool from ..types.content import ContentBlock, Messages, SystemContentBlock from ..types.exceptions import ( @@ -98,6 +98,9 @@ class BedrockConfig(TypedDict, total=False): Please check https://docs.aws.amazon.com/bedrock/latest/userguide/service-tiers-inference.html for supported service tiers, models, and regions stop_sequences: List of sequences that will stop generation when encountered + structured_output_mode: Mode for structured output. "tool" (default) uses tool-based approach, + "native" uses Bedrock's outputConfig.textFormat for schema-constrained responses. + Native mode requires a model that supports structured output. streaming: Flag to enable/disable streaming. Defaults to True. temperature: Controls randomness in generation (higher = more random) top_p: Controls diversity via nucleus sampling (alternative to temperature) @@ -123,6 +126,7 @@ class BedrockConfig(TypedDict, total=False): include_tool_result_status: Literal["auto"] | bool | None service_tier: str | None stop_sequences: list[str] | None + structured_output_mode: Literal["tool", "native"] | None streaming: bool | None temperature: float | None top_p: float | None @@ -218,6 +222,7 @@ def _format_request( tool_specs: list[ToolSpec] | None = None, system_prompt_content: list[SystemContentBlock] | None = None, tool_choice: ToolChoice | None = None, + output_config: dict[str, Any] | None = None, ) -> dict[str, Any]: """Format a Bedrock converse stream request. @@ -226,6 +231,7 @@ def _format_request( tool_specs: List of tool specifications to make available to the model. tool_choice: Selection strategy for tool invocation. system_prompt_content: System prompt content blocks to provide context to the model. + output_config: Output configuration for structured output (JSON schema). Returns: A Bedrock converse stream request. @@ -251,6 +257,20 @@ def _format_request( "messages": self._format_bedrock_messages(messages), "system": system_blocks, **({"serviceTier": {"type": self.config["service_tier"]}} if self.config.get("service_tier") else {}), + **( + { + "outputConfig": { + "textFormat": { + "type": "json_schema", + "structure": { + "jsonSchema": output_config, + }, + }, + } + } + if output_config + else {} + ), **( { "toolConfig": { @@ -747,6 +767,7 @@ async def stream( *, tool_choice: ToolChoice | None = None, system_prompt_content: list[SystemContentBlock] | None = None, + output_config: dict[str, Any] | None = None, **kwargs: Any, ) -> AsyncGenerator[StreamEvent, None]: """Stream conversation with the Bedrock model. @@ -760,6 +781,7 @@ async def stream( system_prompt: System prompt to provide context to the model. tool_choice: Selection strategy for tool invocation. system_prompt_content: System prompt content blocks to provide context to the model. + output_config: Output configuration for structured output (JSON schema). **kwargs: Additional keyword arguments for future extensibility. Yields: @@ -782,7 +804,9 @@ def callback(event: StreamEvent | None = None) -> None: if system_prompt and system_prompt_content is None: system_prompt_content = [{"text": system_prompt}] - thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt_content, tool_choice) + thread = asyncio.to_thread( + self._stream, callback, messages, tool_specs, system_prompt_content, tool_choice, output_config + ) task = asyncio.create_task(thread) while True: @@ -801,6 +825,7 @@ def _stream( tool_specs: list[ToolSpec] | None = None, system_prompt_content: list[SystemContentBlock] | None = None, tool_choice: ToolChoice | None = None, + output_config: dict[str, Any] | None = None, ) -> None: """Stream conversation with the Bedrock model. @@ -813,6 +838,7 @@ def _stream( tool_specs: List of tool specifications to make available to the model. system_prompt_content: System prompt content blocks to provide context to the model. tool_choice: Selection strategy for tool invocation. + output_config: Output configuration for structured output (JSON schema). Raises: ContextWindowOverflowException: If the input exceeds the model's context window. @@ -820,7 +846,7 @@ def _stream( """ try: logger.debug("formatting request") - request = self._format_request(messages, tool_specs, system_prompt_content, tool_choice) + request = self._format_request(messages, tool_specs, system_prompt_content, tool_choice, output_config) logger.debug("request=<%s>", request) logger.debug("invoking model") @@ -1032,6 +1058,10 @@ async def structured_output( ) -> AsyncGenerator[dict[str, T | Any], None]: """Get structured output from the model. + Supports two modes controlled by `structured_output_mode` config: + - "tool" (default): Converts the Pydantic model to a tool spec and forces tool use. + - "native": Uses Bedrock's outputConfig.textFormat with JSON schema for guaranteed schema compliance. + Args: output_model: The output model to use for the agent. prompt: The prompt messages to use for the agent. @@ -1041,6 +1071,21 @@ async def structured_output( Yields: Model events with the last being the structured output. """ + if self.config.get("structured_output_mode") == "native": + async for event in self._structured_output_native(output_model, prompt, system_prompt, **kwargs): + yield event + else: + async for event in self._structured_output_tool(output_model, prompt, system_prompt, **kwargs): + yield event + + async def _structured_output_tool( + self, + output_model: type[T], + prompt: Messages, + system_prompt: str | None = None, + **kwargs: Any, + ) -> AsyncGenerator[dict[str, T | Any], None]: + """Structured output using tool-based approach.""" tool_spec = convert_pydantic_to_tool_spec(output_model) response = self.stream( @@ -1073,6 +1118,39 @@ async def structured_output( yield {"output": output_model(**output_response)} + async def _structured_output_native( + self, + output_model: type[T], + prompt: Messages, + system_prompt: str | None = None, + **kwargs: Any, + ) -> AsyncGenerator[dict[str, T | Any], None]: + """Structured output using Bedrock's native outputConfig.textFormat.""" + output_config = convert_pydantic_to_json_schema(output_model) + + response = self.stream( + messages=prompt, + system_prompt=system_prompt, + output_config=output_config, + **kwargs, + ) + async for event in streaming.process_stream(response): + yield event + + _, messages, _, _ = event["stop"] + + content = messages["content"] + text_content: str | None = None + for block in content: + if "text" in block and block["text"].strip(): + text_content = block["text"] + + if text_content is None: + raise ValueError("No text content found in the Bedrock response for native structured output.") + + output_response = json.loads(text_content) + yield {"output": output_model(**output_response)} + @staticmethod def _get_default_model_with_warning(region_name: str, model_config: BedrockConfig | None = None) -> str: """Get the default Bedrock modelId based on region. diff --git a/src/strands/tools/__init__.py b/src/strands/tools/__init__.py index ada49369d..74d1ca3f0 100644 --- a/src/strands/tools/__init__.py +++ b/src/strands/tools/__init__.py @@ -4,7 +4,7 @@ """ from .decorator import tool -from .structured_output import convert_pydantic_to_tool_spec +from .structured_output import convert_pydantic_to_json_schema, convert_pydantic_to_tool_spec from .tool_provider import ToolProvider from .tools import InvalidToolUseNameException, PythonAgentTool, normalize_schema, normalize_tool_spec @@ -14,6 +14,7 @@ "InvalidToolUseNameException", "normalize_schema", "normalize_tool_spec", + "convert_pydantic_to_json_schema", "convert_pydantic_to_tool_spec", "ToolProvider", ] diff --git a/src/strands/tools/structured_output/__init__.py b/src/strands/tools/structured_output/__init__.py index a3a12d000..85df0a71f 100644 --- a/src/strands/tools/structured_output/__init__.py +++ b/src/strands/tools/structured_output/__init__.py @@ -1,6 +1,6 @@ """Structured output tools for the Strands Agents framework.""" from ._structured_output_context import DEFAULT_STRUCTURED_OUTPUT_PROMPT -from .structured_output_utils import convert_pydantic_to_tool_spec +from .structured_output_utils import convert_pydantic_to_json_schema, convert_pydantic_to_tool_spec -__all__ = ["convert_pydantic_to_tool_spec", "DEFAULT_STRUCTURED_OUTPUT_PROMPT"] +__all__ = ["convert_pydantic_to_json_schema", "convert_pydantic_to_tool_spec", "DEFAULT_STRUCTURED_OUTPUT_PROMPT"] diff --git a/src/strands/tools/structured_output/_structured_output_context.py b/src/strands/tools/structured_output/_structured_output_context.py index 9a5190d9d..227203b52 100644 --- a/src/strands/tools/structured_output/_structured_output_context.py +++ b/src/strands/tools/structured_output/_structured_output_context.py @@ -23,6 +23,7 @@ def __init__( self, structured_output_model: type[BaseModel] | None = None, structured_output_prompt: str | None = None, + native_mode: bool = False, ): """Initialize a new structured output context. @@ -30,10 +31,14 @@ def __init__( structured_output_model: Optional Pydantic model type for structured output. structured_output_prompt: Optional custom prompt message to use when forcing structured output. Defaults to "You must format the previous response as structured output." + native_mode: If True, use the model's native structured output for the final formatting step + instead of forcing tool use. The agent loop runs normally with tools and thinking; + only the final response formatting uses native structured output. """ self.results: dict[str, BaseModel] = {} self.structured_output_model: type[BaseModel] | None = structured_output_model self.structured_output_tool: StructuredOutputTool | None = None + self.native_mode: bool = native_mode self.forced_mode: bool = False self.force_attempted: bool = False self.tool_choice: ToolChoice | None = None @@ -41,7 +46,7 @@ def __init__( self.expected_tool_name: str | None = None self.structured_output_prompt: str = structured_output_prompt or DEFAULT_STRUCTURED_OUTPUT_PROMPT - if structured_output_model: + if structured_output_model and not native_mode: self.structured_output_tool = StructuredOutputTool(structured_output_model) self.expected_tool_name = self.structured_output_tool.tool_name diff --git a/src/strands/tools/structured_output/structured_output_utils.py b/src/strands/tools/structured_output/structured_output_utils.py index a78ec6195..db2f348e6 100644 --- a/src/strands/tools/structured_output/structured_output_utils.py +++ b/src/strands/tools/structured_output/structured_output_utils.py @@ -1,5 +1,6 @@ """Tools for converting Pydantic models to Bedrock tools.""" +import json from typing import Any, Union from pydantic import BaseModel @@ -257,48 +258,56 @@ def _process_nested_dict(d: dict[str, Any], defs: dict[str, Any]) -> dict[str, A return result -def convert_pydantic_to_tool_spec( +def _prepare_pydantic_schema( model: type[BaseModel], description: str | None = None, -) -> ToolSpec: - """Converts a Pydantic model to a tool description for the Amazon Bedrock Converse API. +) -> tuple[str, str, dict[str, Any]]: + """Shared pipeline for converting a Pydantic model to a flattened JSON schema. - Handles optional vs. required fields, resolves $refs, and uses docstrings. + Resolves $refs, expands nested properties, flattens the schema, and resolves the description. Args: - model: The Pydantic model class to convert - description: Optional description of the tool's purpose + model: The Pydantic model class to convert. + description: Optional description override. Returns: - ToolSpec: Dict containing the Bedrock tool specification + Tuple of (name, description, flattened_schema). """ name = model.__name__ - # Get the JSON schema input_schema = model.model_json_schema() - # Get model docstring for description if not provided model_description = description if not model_description and model.__doc__: model_description = model.__doc__.strip() - # Process all referenced models to ensure proper docstrings - # This step is important for gathering descriptions from referenced models _process_referenced_models(input_schema, model) - - # Now, let's fully expand the nested models with all their properties _expand_nested_properties(input_schema, model) - # Flatten the schema - flattened_schema = _flatten_schema(input_schema) + return name, model_description or "", _flatten_schema(input_schema) + + +def convert_pydantic_to_tool_spec( + model: type[BaseModel], + description: str | None = None, +) -> ToolSpec: + """Converts a Pydantic model to a tool description for the Amazon Bedrock Converse API. + + Handles optional vs. required fields, resolves $refs, and uses docstrings. + + Args: + model: The Pydantic model class to convert + description: Optional description of the tool's purpose - final_schema = flattened_schema + Returns: + ToolSpec: Dict containing the Bedrock tool specification + """ + name, model_description, flattened_schema = _prepare_pydantic_schema(model, description) - # Construct the tool specification return ToolSpec( name=name, description=model_description or f"{name} structured output tool", - inputSchema={"json": final_schema}, + inputSchema={"json": flattened_schema}, ) @@ -402,3 +411,52 @@ def _process_properties(schema_def: dict[str, Any], model: type[BaseModel]) -> N # Add field description if available and not already set if field and field.description and not prop_info.get("description"): prop_info["description"] = field.description + + +def _add_additional_properties_false(schema: dict[str, Any]) -> None: + """Recursively add additionalProperties: false to all object types in a JSON schema. + + Bedrock's native structured output requires additionalProperties: false at every object level. + Mutates the schema in place. + + Args: + schema: The JSON schema to process (modified in place). + """ + schema_type = schema.get("type") + if schema_type == "object" or (isinstance(schema_type, list) and "object" in schema_type): + schema["additionalProperties"] = False + + if "properties" in schema: + for value in schema["properties"].values(): + if isinstance(value, dict): + _add_additional_properties_false(value) + + if "items" in schema and isinstance(schema["items"], dict): + _add_additional_properties_false(schema["items"]) + + +def convert_pydantic_to_json_schema( + model: type[BaseModel], + description: str | None = None, +) -> dict[str, Any]: + """Convert a Pydantic model to a JSON schema dict for Bedrock native structured output. + + Returns a dict with "schema" (JSON string), "name", and "description" keys, + suitable for use in outputConfig.textFormat.structure.jsonSchema. + + Args: + model: The Pydantic model class to convert. + description: Optional description override. + + Returns: + Dict with "schema" (JSON string), "name", and "description". + """ + name, model_description, flattened_schema = _prepare_pydantic_schema(model, description) + + _add_additional_properties_false(flattened_schema) + + return { + "schema": json.dumps(flattened_schema), + "name": name, + "description": model_description or f"{name} structured output", + } diff --git a/tests/strands/models/test_bedrock.py b/tests/strands/models/test_bedrock.py index 9c565d4f4..e1691839a 100644 --- a/tests/strands/models/test_bedrock.py +++ b/tests/strands/models/test_bedrock.py @@ -1433,6 +1433,97 @@ async def test_structured_output(bedrock_client, model, test_output_model_cls, a assert tru_output == exp_output +@pytest.mark.asyncio +async def test_structured_output_native_mode(bedrock_client, model, test_output_model_cls, alist): + """Test structured_output with native mode returns parsed JSON text response.""" + model.config["structured_output_mode"] = "native" + messages = [{"role": "user", "content": [{"text": "Generate a person"}]}] + + bedrock_client.converse_stream.return_value = { + "stream": [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"start": {"text": ""}}}, + {"contentBlockDelta": {"delta": {"text": '{"name": "John", "age": 30}'}}}, + {"contentBlockStop": {}}, + {"messageStop": {"stopReason": "end_turn"}}, + ] + } + + stream = model.structured_output(test_output_model_cls, messages) + events = await alist(stream) + + tru_output = events[-1] + exp_output = {"output": test_output_model_cls(name="John", age=30)} + assert tru_output == exp_output + + +@pytest.mark.asyncio +async def test_structured_output_native_mode_sends_output_config(bedrock_client, model, test_output_model_cls, alist): + """Test that native mode sends outputConfig and no toolConfig.""" + model.config["structured_output_mode"] = "native" + messages = [{"role": "user", "content": [{"text": "Generate a person"}]}] + + bedrock_client.converse_stream.return_value = { + "stream": [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"start": {"text": ""}}}, + {"contentBlockDelta": {"delta": {"text": '{"name": "John", "age": 30}'}}}, + {"contentBlockStop": {}}, + {"messageStop": {"stopReason": "end_turn"}}, + ] + } + + stream = model.structured_output(test_output_model_cls, messages) + await alist(stream) + + call_kwargs = bedrock_client.converse_stream.call_args[1] + assert "outputConfig" in call_kwargs + assert call_kwargs["outputConfig"]["textFormat"]["type"] == "json_schema" + assert "toolConfig" not in call_kwargs + + +@pytest.mark.asyncio +async def test_structured_output_tool_mode_unchanged(bedrock_client, model, test_output_model_cls, alist): + """Test that tool mode (default) still uses the tool-based approach.""" + messages = [{"role": "user", "content": [{"text": "Generate a person"}]}] + + bedrock_client.converse_stream.return_value = { + "stream": [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"start": {"toolUse": {"toolUseId": "123", "name": "TestOutputModel"}}}}, + {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"name": "John", "age": 30}'}}}}, + {"contentBlockStop": {}}, + {"messageStop": {"stopReason": "tool_use"}}, + ] + } + + stream = model.structured_output(test_output_model_cls, messages) + events = await alist(stream) + + tru_output = events[-1] + exp_output = {"output": test_output_model_cls(name="John", age=30)} + assert tru_output == exp_output + + +@pytest.mark.asyncio +async def test_structured_output_native_mode_no_text_raises(bedrock_client, model, test_output_model_cls, alist): + """Test that native mode raises if no text content in response.""" + model.config["structured_output_mode"] = "native" + messages = [{"role": "user", "content": [{"text": "Generate a person"}]}] + + bedrock_client.converse_stream.return_value = { + "stream": [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockStop": {}}, + {"messageStop": {"stopReason": "end_turn"}}, + ] + } + + with pytest.raises(ValueError, match="No text content found"): + stream = model.structured_output(test_output_model_cls, messages) + await alist(stream) + + @pytest.mark.skipif(sys.version_info < (3, 11), reason="This test requires Python 3.11 or higher (need add_note)") @pytest.mark.asyncio async def test_add_note_on_client_error(bedrock_client, model, alist, messages): @@ -2823,3 +2914,29 @@ def test_guardrail_latest_message_disabled_does_not_wrap(model): assert "text" in formatted assert "guardContent" not in formatted + + +def test_format_request_with_output_config(model, messages, model_id): + """Test that output_config is included in request as outputConfig.textFormat.""" + output_config = { + "schema": '{"type": "object", "properties": {"name": {"type": "string"}}, "additionalProperties": false}', + "name": "TestModel", + "description": "Test description", + } + request = model._format_request(messages, output_config=output_config) + + assert request["outputConfig"] == { + "textFormat": { + "type": "json_schema", + "structure": { + "jsonSchema": output_config, + }, + }, + } + assert request["modelId"] == model_id + + +def test_format_request_without_output_config(model, messages, model_id): + """Test that outputConfig is not included when output_config is None.""" + request = model._format_request(messages) + assert "outputConfig" not in request diff --git a/tests/strands/tools/test_structured_output.py b/tests/strands/tools/test_structured_output.py index 72a53bfe6..a76f7c971 100644 --- a/tests/strands/tools/test_structured_output.py +++ b/tests/strands/tools/test_structured_output.py @@ -1,9 +1,10 @@ +import json from typing import Literal, Optional import pytest from pydantic import BaseModel, Field -from strands.tools.structured_output import convert_pydantic_to_tool_spec +from strands.tools.structured_output import convert_pydantic_to_json_schema, convert_pydantic_to_tool_spec from strands.types.tools import ToolSpec @@ -423,3 +424,68 @@ class Person(BaseModel): "name": "Person", } assert tool_spec == expected_spec + + +def test_convert_pydantic_to_json_schema_basic(): + result = convert_pydantic_to_json_schema(User) + + assert result["name"] == "User" + assert result["description"] == "User model with name and age." + assert isinstance(result["schema"], str) + + schema = json.loads(result["schema"]) + assert schema["type"] == "object" + assert schema["additionalProperties"] is False + assert "name" in schema["properties"] + assert "age" in schema["properties"] + assert schema["required"] == ["name", "age"] + + +def test_convert_pydantic_to_json_schema_nested(): + """Test nested model has additionalProperties: false at every object level and preserves structure.""" + result = convert_pydantic_to_json_schema(TwoUsersWithPlanet) + + schema = json.loads(result["schema"]) + assert schema["additionalProperties"] is False + assert "user1" in schema["properties"] + assert "user2" in schema["properties"] + + # Required nested object + user1 = schema["properties"]["user1"] + assert user1.get("additionalProperties") is False + assert user1["type"] == "object" + assert "name" in user1["properties"] + assert "age" in user1["properties"] + assert "planet" in user1["properties"] + + # Optional nested object (type includes "null") + user2 = schema["properties"]["user2"] + assert user2.get("additionalProperties") is False + assert "object" in user2["type"] + assert "null" in user2["type"] + + # List of objects: additionalProperties: false propagates into array items + list_result = convert_pydantic_to_json_schema(ListOfUsersWithPlanet) + list_schema = json.loads(list_result["schema"]) + assert list_schema["additionalProperties"] is False + + users = list_schema["properties"]["users"] + assert users["type"] == "array" + items = users["items"] + assert items.get("additionalProperties") is False + assert "name" in items["properties"] + assert "age" in items["properties"] + assert "planet" in items["properties"] + + +def test_convert_pydantic_to_json_schema_custom_description(): + result = convert_pydantic_to_json_schema(User, description="Custom desc") + assert result["description"] == "Custom desc" + + +def test_convert_pydantic_to_json_schema_no_docstring(): + class SimpleModel(BaseModel): + value: str + + result = convert_pydantic_to_json_schema(SimpleModel) + assert result["description"] == "SimpleModel structured output"