feat: Add support for toolChoice to providers

For structured output so that some providers can force tool calls

Co-authored-by: Shang Liu <[email protected]>

Author: zastrowm

src/strands/models/anthropic.py

@@ -18,7 +18,7 @@
 from ..types.content import ContentBlock, Messages
 from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
 from ..types.streaming import StreamEvent
-from ..types.tools import ToolSpec
+from ..types.tools import ToolChoice, ToolSpec
 from ._config_validation import validate_config_keys
 from .model import Model
 
@@ -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: Optional[ToolChoice] = 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,6 +228,7 @@ def format_request(
                 }
                 for tool_spec in tool_specs or []
             ],
+            **({"tool_choice": tool_choice} if tool_choice else {}),
             **({"system": system_prompt} if system_prompt else {}),
             **(self.config.get("params") or {}),
         }
@@ -350,6 +356,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: Optional[ToolChoice] = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the Anthropic model.
@@ -358,6 +365,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 +376,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 +418,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,7 +23,7 @@
     ModelThrottledException,
 )
 from ..types.streaming import CitationsDelta, StreamEvent
-from ..types.tools import ToolResult, ToolSpec
+from ..types.tools import ToolChoice, ToolResult, ToolSpec
 from ._config_validation import validate_config_keys
 from .model import Model
 
@@ -195,13 +195,15 @@ def format_request(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: Optional[ToolChoice] = 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.
@@ -224,7 +226,7 @@ def format_request(
                                 else []
                             ),
                         ],
-                        "toolChoice": {"auto": {}},
+                        **({"toolChoice": tool_choice} if tool_choice else {}),
                     }
                 }
                 if tool_specs
@@ -416,6 +418,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: Optional[ToolChoice] = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the Bedrock model.
@@ -427,6 +430,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:
@@ -445,7 +449,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:
@@ -463,6 +467,7 @@ def _stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: Optional[ToolChoice] = None,
     ) -> None:
         """Stream conversation with the Bedrock model.
 
@@ -474,14 +479,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")
@@ -738,6 +744,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,7 +14,7 @@
 
 from ..types.content import ContentBlock, Messages
 from ..types.streaming import StreamEvent
-from ..types.tools import ToolSpec
+from ..types.tools import ToolChoice, ToolSpec
 from ._config_validation import validate_config_keys
 from .openai import OpenAIModel
 
@@ -114,6 +114,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: Optional[ToolChoice] = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the LiteLLM model.
@@ -122,6 +123,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:

src/strands/models/llamaapi.py

@@ -18,7 +18,7 @@
 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 ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse
 from ._config_validation import validate_config_keys
 from .model import Model
 
@@ -330,6 +330,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: Optional[ToolChoice] = 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:

src/strands/models/mistral.py

@@ -15,7 +15,7 @@
 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 ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse
 from ._config_validation import validate_config_keys
 from .model import Model
 
@@ -397,6 +397,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: Optional[ToolChoice] = 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:

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: Optional[ToolChoice] = 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,7 +13,7 @@
 
 from ..types.content import ContentBlock, Messages
 from ..types.streaming import StopReason, StreamEvent
-from ..types.tools import ToolSpec
+from ..types.tools import ToolChoice, ToolSpec
 from ._config_validation import validate_config_keys
 from .model import Model
 
@@ -287,6 +287,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: Optional[ToolChoice] = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the Ollama model.
@@ -295,6 +296,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:

src/strands/models/openai.py

@@ -16,7 +16,7 @@
 
 from ..types.content import ContentBlock, Messages
 from ..types.streaming import StreamEvent
-from ..types.tools import ToolResult, ToolSpec, ToolUse
+from ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse
 from ._config_validation import validate_config_keys
 from .model import Model
 
@@ -174,6 +174,27 @@ def format_request_tool_message(cls, tool_result: ToolResult) -> dict[str, Any]:
             "content": [cls.format_request_message_content(content) for content in contents],
         }
 
+    @classmethod
+    def format_request_tool_choice(cls, tool_choice: ToolChoice) -> Union[str, dict[str, Any]]:
+        """Format a tool choice for OpenAI compatibility.
+
+        Args:
+            tool_choice: Tool choice configuration in Bedrock format.
+
+        Returns:
+            OpenAI compatible tool choice format.
+        """
+        match tool_choice:
+            case {"auto": _}:
+                return "auto"  # OpenAI SDK doesn't define constants for these values
+            case {"any": _}:
+                return "required"
+            case {"tool": {"name": tool_name}}:
+                return {"type": "function", "function": {"name": tool_name}}
+            case _:
+                # This should not happen with proper typing, but handle gracefully
+                return "auto"
+
     @classmethod
     def format_request_messages(cls, messages: Messages, system_prompt: Optional[str] = None) -> list[dict[str, Any]]:
         """Format an OpenAI compatible messages array.
@@ -216,14 +237,19 @@ def format_request_messages(cls, messages: Messages, system_prompt: Optional[str
         return [message for message in formatted_messages if message["content"] or "tool_calls" in message]
 
     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: Optional[ToolChoice] = None,
     ) -> dict[str, Any]:
         """Format an OpenAI compatible chat 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 OpenAI compatible chat streaming request.
@@ -248,6 +274,7 @@ def format_request(
                 }
                 for tool_spec in tool_specs or []
             ],
+            **({"tool_choice": self.format_request_tool_choice(tool_choice)} if tool_choice else {}),
             **cast(dict[str, Any], self.config.get("params", {})),
         }
 
@@ -329,6 +356,7 @@ async def stream(
         messages: Messages,
         tool_specs: Optional[list[ToolSpec]] = None,
         system_prompt: Optional[str] = None,
+        tool_choice: Optional[ToolChoice] = None,
         **kwargs: Any,
     ) -> AsyncGenerator[StreamEvent, None]:
         """Stream conversation with the OpenAI model.
@@ -337,13 +365,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("formatted request=<%s>", request)
 
         logger.debug("invoking model")