Repository: pydantic/pydantic-ai
File: tests/test_agent.py (Python)

Code before:
                    'type': 'object',
                },
                kind='output',
            ),
            ToolDefinition(
                name='final_result_Bar',
                description='This is a bar model.',
                parameters_json_schema={
                    'properties': {'b': {'type': 'string'}},
                    'required': ['b'],
                    'title': 'Bar',
                    'type': 'object',
                },
                kind='output',
            ),
        ]
    )

    result = agent.run_sync('Hello', model=TestModel(seed=1))
    assert result.output == mod.Bar(b='b')
    assert got_tool_call_name == snapshot('final_result_Bar')


def test_output_type_generic_class_name_sanitization(create_module: Callable[[str], Any]):
    """Test that generic class names with brackets are properly sanitized."""
    module_code = '''
from pydantic import BaseModel
from typing import Generic, TypeVar

T = TypeVar('T')

class Result(BaseModel, Generic[T]):
    """A generic result class."""
    value: T
    success: bool

class StringData(BaseModel):
    text: str

# This will have a name like "Result[StringData]" which needs sanitization
OutputType = [Result[StringData], Result[int]]
    '''

    mod = create_module(module_code)

    m = TestModel()
    agent = Agent(m, output_type=mod.OutputType)
    agent.run_sync('Hello')

    # The sanitizer should remove brackets from the generic type name
    assert m.last_model_request_parameters is not None

Review comment (suggestion):
We can inline this as regular code -- we're only using a stringified module above because we're injecting code `{union_code}`.

Code after:
    assert len(m.last_model_request_parameters.output_tools) == 2

    assert m.last_model_request_parameters.output_tools == snapshot(
        [
            ToolDefinition(
                name='final_result_Foo',
                description='Foo: The final response which ends this conversation',
                parameters_json_schema={
                    'properties': {
                        'a': {'type': 'integer'},
                        'b': {'type': 'string'},
                    },
                    'required': ['a', 'b'],
                    'title': 'Foo',
                    'type': 'object',
                },
                kind='output',
            ),
            ToolDefinition(
                name='final_result_Bar',
                description='This is a bar model.',
                parameters_json_schema={
                    'properties': {'b': {'type': 'string'}},
                    'required': ['b'],
                    'title': 'Bar',
                    'type': 'object',
                },
                kind='output',
            ),
        ]
    )

    result = agent.run_sync('Hello', model=TestModel(seed=1))
    assert result.output == mod.Bar(b='b')
    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