Add MCPToolChoice

Author: seratch

src/agents/model_settings.py

@@ -46,6 +46,10 @@ def validate_from_none(value: None) -> _Omit:
 Omit = Annotated[_Omit, _OmitTypeAnnotation]
 Headers: TypeAlias = Mapping[str, Union[str, Omit]]
 
+@dataclass
+class MCPToolChoice:
+    server_label: str
+    name: str
 
 @dataclass
 class ModelSettings:
@@ -70,7 +74,7 @@ class ModelSettings:
     presence_penalty: float | None = None
     """The presence penalty to use when calling the model."""
 
-    tool_choice: Literal["auto", "required", "none"] | str | None = None
+    tool_choice: Literal["auto", "required", "none"] | str | MCPToolChoice | None = None
     """The tool choice to use when calling the model."""
 
     parallel_tool_calls: bool | None = None

src/agents/models/openai_responses.py

@@ -25,6 +25,7 @@
 from ..handoffs import Handoff
 from ..items import ItemHelpers, ModelResponse, TResponseInputItem
 from ..logger import logger
+from ..model_settings import MCPToolChoice
 from ..tool import (
     CodeInterpreterTool,
     ComputerTool,
@@ -303,19 +304,16 @@ class ConvertedTools:
 class Converter:
     @classmethod
     def convert_tool_choice(
-        cls, tool_choice: Literal["auto", "required", "none"] | str | dict[str, Any] | None
+        cls, tool_choice: Literal["auto", "required", "none"] | str | MCPToolChoice | None
     ) -> response_create_params.ToolChoice | NotGiven:
         if tool_choice is None:
             return NOT_GIVEN
-        elif isinstance(tool_choice, dict):
-            if tool_choice.get("type") == "mcp":
-                return {
-                    "server_label": tool_choice.get("server_label") or "mcp",
-                    "type": "mcp",
-                    "name": tool_choice.get("name"),
-                }
-            else:
-                raise UserError(f"Unknown tool choice: {tool_choice}")
+        elif isinstance(tool_choice, MCPToolChoice):
+            return {
+                "server_label": tool_choice.server_label,
+                "type": "mcp",
+                "name": tool_choice.name,
+            }
         elif tool_choice == "required":
             return "required"
         elif tool_choice == "auto":
@@ -343,10 +341,9 @@ def convert_tool_choice(
                 "type": "code_interpreter",
             }
         elif tool_choice == "mcp":
-            return {
-                "server_label": "mcp",
-                "type": "mcp",
-            }
+            # Note that this is still here for backwards compatibility,
+            # but migrating to MCPToolChoice is recommended.
+            return { "type": "mcp" }
         else:
             return {
                 "type": "function",

tests/model_settings/test_serialization.py

@@ -5,7 +5,7 @@
 from pydantic import TypeAdapter
 from pydantic_core import to_json
 
-from agents.model_settings import ModelSettings
+from agents.model_settings import MCPToolChoice, ModelSettings
 
 
 def verify_serialization(model_settings: ModelSettings) -> None:
@@ -29,6 +29,17 @@ def test_basic_serialization() -> None:
     verify_serialization(model_settings)
 
 
+def test_mcp_tool_choice_serialization() -> None:
+    """Tests whether ModelSettings with MCPToolChoice can be serialized to a JSON string."""
+    # First, lets create a ModelSettings instance
+    model_settings = ModelSettings(
+        temperature=0.5,
+        tool_choice=MCPToolChoice(server_label="mcp", name="mcp_tool"),
+    )
+    # Now, lets serialize the ModelSettings instance to a JSON string
+    verify_serialization(model_settings)
+
+
 def test_all_fields_serialization() -> None:
     """Tests whether ModelSettings can be serialized to a JSON string."""