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

Code before:
                model_name='function:call_tool:',
                timestamp=IsDatetime(),
                run_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    RetryPromptPart(
                        content='City not found, I only know Mexico City',
                        tool_call_id=IsStr(),
                        timestamp=IsDatetime(),
                    )
                ],
                run_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='Mexico City')],
                usage=RequestUsage(input_tokens=70, output_tokens=5),
                model_name='function:call_tool:',
                timestamp=IsDatetime(),
                run_id=IsStr(),
            ),
        ]
    )


def test_output_type_text_output_function_with_deps():
    """Test that TextOutput functions can use RunContext with custom deps type."""

    @dataclass
    class Deps:
        prefix: str

    def add_prefix_and_split(ctx: RunContext[Deps], text: str) -> list[str]:
        return f'{ctx.deps.prefix}: {text}'.split()

    def return_text(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart(content='Hello world')])

    agent = Agent(FunctionModel(return_text), deps_type=Deps, output_type=TextOutput(add_prefix_and_split))
    result = agent.run_sync('test prompt', deps=Deps(prefix='Response'))
    assert result.output == snapshot(['Response:', 'Hello', 'world'])


@pytest.mark.parametrize(
    'output_type',
    [[str, str], [str, TextOutput(upcase)], [TextOutput(upcase), TextOutput(upcase)]],
)
def test_output_type_multiple_text_output(output_type: OutputSpec[str]):
    with pytest.raises(UserError, match='Only one `str` or `TextOutput` is allowed.'):
        Agent('test', output_type=output_type)


Review comment (suggestion):
I'm pretty sure we've already tested somewhere that it _works_, this PR is just for fixing the type checking, so let's remove the new tests and add some cases to `typed_agent.py` instead.

Code after:
                model_name='function:call_tool:',
                timestamp=IsDatetime(),
                run_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    RetryPromptPart(
                        content='City not found, I only know Mexico City',
                        tool_call_id=IsStr(),
                        timestamp=IsDatetime(),
                    )
                ],
                run_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='Mexico City')],
                usage=RequestUsage(input_tokens=70, output_tokens=5),
                model_name='function:call_tool:',
                timestamp=IsDatetime(),
                run_id=IsStr(),
            ),
        ]
    )


@pytest.mark.parametrize(
    'output_type',
    [[str, str], [str, TextOutput(upcase)], [TextOutput(upcase), TextOutput(upcase)]],
)
def test_output_type_multiple_text_output(output_type: OutputSpec[str]):
    with pytest.raises(UserError, match='Only one `str` or `TextOutput` is allowed.'):
        Agent('test', output_type=output_type)


def test_output_type_text_output_invalid():
    def int_func(x: int) -> str:
        return str(int)  # pragma: no cover

    with pytest.raises(UserError, match='TextOutput must take a function taking a single `str` argument'):
        output_type: TextOutput[str] = TextOutput(int_func)  # type: ignore
        Agent('test', output_type=output_type)


def test_output_type_async_function():
    class Weather(BaseModel):
        temperature: float
        description: str

    async def get_weather(city: str) -> Weather:
        return Weather(temperature=28.7, description='sunny')