feat: add support for Bedrock/Anthropic ToolChoice to structured_output (#720)

For structured output so that some providers can force tool calls

---------

Co-authored-by: Mackenzie Zastrow <zastrowm@users.noreply.github.com>
Co-authored-by: Shang Liu <sshangl@amazon.com>

Author: 3 people

src/strands/models/_config_validation.py renamed to src/strands/models/_validation.py

@@ -5,6 +5,8 @@
 
 from typing_extensions import get_type_hints
 
+from ..types.tools import ToolChoice
+
 
 def validate_config_keys(config_dict: Mapping[str, Any], config_class: Type) -> None:
     """Validate that config keys match the TypedDict fields.
@@ -25,3 +27,16 @@ def validate_config_keys(config_dict: Mapping[str, Any], config_class: Type) ->
             f"\nSee https://github.com/strands-agents/sdk-python/issues/815",
             stacklevel=4,
         )
+
+
+def warn_on_tool_choice_not_supported(tool_choice: ToolChoice | None) -> None:
+    """Emits a warning if a tool choice is provided but not supported by the provider.
+
+    Args:
+        tool_choice: the tool_choice provided to the provider
+    """
+    if tool_choice:
+        warnings.warn(
+            "A ToolChoice was provided to this provider but is not supported and will be ignored",
+            stacklevel=4,
+        )

src/strands/models/anthropic.py

@@ -18,8 +18,8 @@
 from ..types.content import ContentBlock, Messages
 from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
 from ..types.streaming import StreamEvent
-from ..types.tools import ToolSpec
-from ._config_validation import validate_config_keys
+from ..types.tools import ToolChoice, ToolChoiceToolDict, ToolSpec
+from ._validation import validate_config_keys
 from .model import Model
 
 logger = logging.getLogger(__name__)
@@ -195,14 +195,19 @@ def _format_request_messages(self, messages: Messages) -> list[dict[str, Any]]:
         return formatted_messages
 
     def format_request(
-        self, messages: Messages, tool_specs: Optional[list[ToolSpec]] = None, system_prompt: Optional[str] = None
+        self,
+        messages: Messages,
+        tool_specs: Optional[list[ToolSpec]] = None,
+        system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
     ) -> dict[str, Any]:
         """Format an Anthropic streaming request.
 
         Args:
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation.
 
         Returns:
             An Anthropic streaming request.
@@ -223,10 +228,25 @@ def format_request(
                 }
                 for tool_spec in tool_specs or []
             ],
+            **(self._format_tool_choice(tool_choice)),
             **({"system": system_prompt} if system_prompt else {}),
             **(self.config.get("params") or {}),
         }
 
+    @staticmethod
+    def _format_tool_choice(tool_choice: ToolChoice | None) -> dict:
+        if tool_choice is None:
+            return {}
+
+        if "any" in tool_choice:
+            return {"tool_choice": {"type": "any"}}
+        elif "auto" in tool_choice:
+            return {"tool_choice": {"type": "auto"}}
+        elif "tool" in tool_choice:
+            return {"tool_choice": {"type": "tool", "name": cast(ToolChoiceToolDict, tool_choice)["tool"]["name"]}}
+        else:
+            return {}
+
     def format_chunk(self, event: dict[str, Any]) -> StreamEvent:
         """Format the Anthropic response events into standardized message chunks.
 
@@ -350,6 +370,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the Anthropic model.
@@ -358,6 +379,7 @@ async def stream(
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation.
             **kwargs: Additional keyword arguments for future extensibility.
 
         Yields:
@@ -368,7 +390,7 @@ async def stream(
             ModelThrottledException: If the request is throttled by Anthropic.
         """
         logger.debug("formatting request")
-        request = self.format_request(messages, tool_specs, system_prompt)
+        request = self.format_request(messages, tool_specs, system_prompt, tool_choice)
         logger.debug("request=<%s>", request)
 
         logger.debug("invoking model")
@@ -410,7 +432,13 @@ async def structured_output(
         """
         tool_spec = convert_pydantic_to_tool_spec(output_model)
 
-        response = self.stream(messages=prompt, tool_specs=[tool_spec], system_prompt=system_prompt, **kwargs)
+        response = self.stream(
+            messages=prompt,
+            tool_specs=[tool_spec],
+            system_prompt=system_prompt,
+            tool_choice=cast(ToolChoice, {"any": {}}),
+            **kwargs,
+        )
         async for event in process_stream(response):
             yield event
 

src/strands/models/bedrock.py

@@ -23,8 +23,8 @@
     ModelThrottledException,
 )
 from ..types.streaming import CitationsDelta, StreamEvent
-from ..types.tools import ToolResult, ToolSpec
-from ._config_validation import validate_config_keys
+from ..types.tools import ToolChoice, ToolResult, ToolSpec
+from ._validation import validate_config_keys
 from .model import Model
 
 logger = logging.getLogger(__name__)
@@ -196,13 +196,15 @@ def format_request(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
     ) -> dict[str, Any]:
         """Format a Bedrock converse stream request.
 
         Args:
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation.
 
         Returns:
             A Bedrock converse stream request.
@@ -225,7 +227,7 @@ def format_request(
                                 else []
                             ),
                         ],
-                        "toolChoice": {"auto": {}},
+                        **({"toolChoice": tool_choice if tool_choice else {"auto": {}}}),
                     }
                 }
                 if tool_specs
@@ -417,6 +419,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the Bedrock model.
@@ -428,6 +431,7 @@ async def stream(
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation.
             **kwargs: Additional keyword arguments for future extensibility.
 
         Yields:
@@ -446,7 +450,7 @@ def callback(event: Optional[StreamEvent] = None) -> None:
         loop = asyncio.get_event_loop()
         queue: asyncio.Queue[Optional[StreamEvent]] = asyncio.Queue()
 
-        thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt)
+        thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt, tool_choice)
         task = asyncio.create_task(thread)
 
         while True:
@@ -464,6 +468,7 @@ def _stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
     ) -> None:
         """Stream conversation with the Bedrock model.
 
@@ -475,14 +480,15 @@ def _stream(
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation.
 
         Raises:
             ContextWindowOverflowException: If the input exceeds the model's context window.
             ModelThrottledException: If the model service is throttling requests.
         """
         try:
             logger.debug("formatting request")
-            request = self.format_request(messages, tool_specs, system_prompt)
+            request = self.format_request(messages, tool_specs, system_prompt, tool_choice)
             logger.debug("request=<%s>", request)
 
             logger.debug("invoking model")
@@ -739,6 +745,7 @@ async def structured_output(
             messages=prompt,
             tool_specs=[tool_spec],
             system_prompt=system_prompt,
+            tool_choice=cast(ToolChoice, {"any": {}}),
             **kwargs,
         )
         async for event in streaming.process_stream(response):

src/strands/models/litellm.py

@@ -14,8 +14,8 @@
 
 from ..types.content import ContentBlock, Messages
 from ..types.streaming import StreamEvent
-from ..types.tools import ToolSpec
-from ._config_validation import validate_config_keys
+from ..types.tools import ToolChoice, ToolSpec
+from ._validation import validate_config_keys
 from .openai import OpenAIModel
 
 logger = logging.getLogger(__name__)
@@ -114,6 +114,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the LiteLLM model.
@@ -122,13 +123,14 @@ async def stream(
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation.
             **kwargs: Additional keyword arguments for future extensibility.
 
         Yields:
             Formatted message chunks from the model.
         """
         logger.debug("formatting request")
-        request = self.format_request(messages, tool_specs, system_prompt)
+        request = self.format_request(messages, tool_specs, system_prompt, tool_choice)
         logger.debug("request=<%s>", request)
 
         logger.debug("invoking model")

src/strands/models/llamaapi.py

@@ -18,8 +18,8 @@
 from ..types.content import ContentBlock, Messages
 from ..types.exceptions import ModelThrottledException
 from ..types.streaming import StreamEvent, Usage
-from ..types.tools import ToolResult, ToolSpec, ToolUse
-from ._config_validation import validate_config_keys
+from ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse
+from ._validation import validate_config_keys, warn_on_tool_choice_not_supported
 from .model import Model
 
 logger = logging.getLogger(__name__)
@@ -330,6 +330,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the LlamaAPI model.
@@ -338,6 +339,8 @@ async def stream(
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation. **Note: This parameter is accepted for
+                interface consistency but is currently ignored for this model provider.**
             **kwargs: Additional keyword arguments for future extensibility.
 
         Yields:
@@ -346,6 +349,8 @@ async def stream(
         Raises:
             ModelThrottledException: When the model service is throttling requests from the client.
         """
+        warn_on_tool_choice_not_supported(tool_choice)
+
         logger.debug("formatting request")
         request = self.format_request(messages, tool_specs, system_prompt)
         logger.debug("request=<%s>", request)

src/strands/models/mistral.py

@@ -15,8 +15,8 @@
 from ..types.content import ContentBlock, Messages
 from ..types.exceptions import ModelThrottledException
 from ..types.streaming import StopReason, StreamEvent
-from ..types.tools import ToolResult, ToolSpec, ToolUse
-from ._config_validation import validate_config_keys
+from ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse
+from ._validation import validate_config_keys, warn_on_tool_choice_not_supported
 from .model import Model
 
 logger = logging.getLogger(__name__)
@@ -397,6 +397,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the Mistral model.
@@ -405,6 +406,8 @@ async def stream(
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation. **Note: This parameter is accepted for
+                interface consistency but is currently ignored for this model provider.**
             **kwargs: Additional keyword arguments for future extensibility.
 
         Yields:
@@ -413,6 +416,8 @@ async def stream(
         Raises:
             ModelThrottledException: When the model service is throttling requests.
         """
+        warn_on_tool_choice_not_supported(tool_choice)
+
         logger.debug("formatting request")
         request = self.format_request(messages, tool_specs, system_prompt)
         logger.debug("request=<%s>", request)

src/strands/models/model.py

@@ -8,7 +8,7 @@
 
 from ..types.content import Messages
 from ..types.streaming import StreamEvent
-from ..types.tools import ToolSpec
+from ..types.tools import ToolChoice, ToolSpec
 
 logger = logging.getLogger(__name__)
 
@@ -70,6 +70,7 @@ def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
         **kwargs: Any,
     ) -> AsyncIterable[StreamEvent]:
         """Stream conversation with the model.
@@ -84,6 +85,7 @@ def stream(
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation.
             **kwargs: Additional keyword arguments for future extensibility.
 
         Yields:

src/strands/models/ollama.py

@@ -13,8 +13,8 @@
 
 from ..types.content import ContentBlock, Messages
 from ..types.streaming import StopReason, StreamEvent
-from ..types.tools import ToolSpec
-from ._config_validation import validate_config_keys
+from ..types.tools import ToolChoice, ToolSpec
+from ._validation import validate_config_keys, warn_on_tool_choice_not_supported
 from .model import Model
 
 logger = logging.getLogger(__name__)
@@ -287,6 +287,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: ToolChoice | None = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the Ollama model.
@@ -295,11 +296,15 @@ async def stream(
             messages: List of message objects to be processed by the model.
             tool_specs: List of tool specifications to make available to the model.
             system_prompt: System prompt to provide context to the model.
+            tool_choice: Selection strategy for tool invocation. **Note: This parameter is accepted for
+                interface consistency but is currently ignored for this model provider.**
             **kwargs: Additional keyword arguments for future extensibility.
 
         Yields:
             Formatted message chunks from the model.
         """
+        warn_on_tool_choice_not_supported(tool_choice)
+
         logger.debug("formatting request")
         request = self.format_request(messages, tool_specs, system_prompt)
         logger.debug("request=<%s>", request)