Sanitize auto-generated output tool name to support generic types (#2979)

Author: lionpeloux

pydantic_ai_slim/pydantic_ai/_output.py

@@ -2,6 +2,7 @@
 
 import inspect
 import json
+import re
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass, field
@@ -70,6 +71,7 @@
 
 DEFAULT_OUTPUT_TOOL_NAME = 'final_result'
 DEFAULT_OUTPUT_TOOL_DESCRIPTION = 'The final response which ends this conversation'
+OUTPUT_TOOL_NAME_SANITIZER = re.compile(r'[^a-zA-Z0-9-_]')
 
 
 async def execute_traced_output_function(
@@ -997,7 +999,9 @@ def build(
             if name is None:
                 name = default_name
                 if multiple:
-                    name += f'_{object_def.name}'
+                    # strip unsupported characters like "[" and "]" from generic class names
+                    safe_name = OUTPUT_TOOL_NAME_SANITIZER.sub('', object_def.name or '')
+                    name += f'_{safe_name}'
 
             i = 1
             original_name = name

tests/test_agent.py

@@ -6,7 +6,7 @@
 from collections.abc import AsyncIterable, Callable
 from dataclasses import dataclass, replace
 from datetime import timezone
-from typing import Any, Literal, Union
+from typing import Any, Generic, Literal, TypeVar, Union
 
 import httpx
 import pytest
@@ -96,6 +96,21 @@ class Person(BaseModel):
     name: str
 
 
+# Generic classes for testing tool name sanitization with generic types
+T = TypeVar('T')
+
+
+class ResultGeneric(BaseModel, Generic[T]):
+    """A generic result class."""
+
+    value: T
+    success: bool
+
+
+class StringData(BaseModel):
+    text: str
+
+
 def test_result_list_of_models_with_stringified_response():
     def return_list(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
         assert info.output_tools is not None
@@ -635,6 +650,24 @@ def validate_output(ctx: RunContext[None], o: Any) -> Any:
     assert got_tool_call_name == snapshot('final_result_Bar')
 
 
+def test_output_type_generic_class_name_sanitization():
+    """Test that generic class names with brackets are properly sanitized."""
+    # This will have a name like "ResultGeneric[StringData]" which needs sanitization
+    output_type = [ResultGeneric[StringData], ResultGeneric[int]]
+
+    m = TestModel()
+    agent = Agent(m, output_type=output_type)
+    agent.run_sync('Hello')
+
+    # The sanitizer should remove brackets from the generic type name
+    assert m.last_model_request_parameters is not None
+    assert m.last_model_request_parameters.output_tools is not None
+    assert len(m.last_model_request_parameters.output_tools) == 2
+
+    tool_names = [tool.name for tool in m.last_model_request_parameters.output_tools]
+    assert tool_names == snapshot(['final_result_ResultGenericStringData', 'final_result_ResultGenericint'])
+
+
 def test_output_type_with_two_descriptions():
     class MyOutput(BaseModel):
         """Description from docstring"""