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

Code before:


async def test_binary_file_message() -> None:
    agent = Agent(model=FunctionModel(stream_function=simple_stream))
    run_input = create_input(
        UserMessage(id='msg_1', content=[BinaryInputContent(data='VGVzdCBEb2M=', mime_type='text/plain')])
    )
    events = await run_and_collect_events(agent, run_input)

    assert events == simple_result()


async def test_test_and_binary_file_message() -> None:
    agent = Agent(model=FunctionModel(stream_function=simple_stream))
    run_input = create_input(
        UserMessage(
            id='msg_1',
            content=[
                TextInputContent(text='Write a summary of this file:', type='text'),
                BinaryInputContent(data='VGVzdCBEb2M=', mime_type='text/plain'),
            ],
        )
    )
    events = await run_and_collect_events(agent, run_input)

    assert events == simple_result()


async def test_tool_ag_ui() -> None:
    """Test AG-UI tool call."""

    async def stream_function(
        messages: list[ModelMessage], agent_info: AgentInfo
    ) -> AsyncIterator[DeltaToolCalls | str]:
        if len(messages) == 1:
            # First call - make a tool call
            yield {0: DeltaToolCall(name='get_weather', json_args='{"location": ')}
            yield {0: DeltaToolCall(json_args='"Paris"}')}
        else:
            # Second call - return text result
            yield '{"get_weather": "Tool result"}'

    agent = Agent(
        model=FunctionModel(stream_function=stream_function),
        tools=[send_snapshot, send_custom, current_time],
    )

    thread_id = uuid_str()
    run_inputs = [
        create_input(
            UserMessage(

Review comment (suggestion):
Let's drop all of the tests above because they're not verifying that the binary content actually made it through. I'd rather expand the `test_messages` test

Code after:

    events = await run_and_collect_events(agent, *run_inputs)

    assert events == snapshot(
        [
            {
                'type': 'RUN_STARTED',
                'timestamp': IsInt(),
                'threadId': thread_id,
                'runId': (run_id := IsSameStr()),
            },
            {
                'type': 'TOOL_CALL_START',
                'timestamp': IsInt(),
                'toolCallId': (tool_call_id := IsSameStr()),
                'toolCallName': 'get_weather',
                'parentMessageId': IsStr(),
            },
            {
                'type': 'TOOL_CALL_ARGS',
                'timestamp': IsInt(),
                'toolCallId': tool_call_id,
                'delta': '{"location": ',
            },
            {'type': 'TOOL_CALL_ARGS', 'timestamp': IsInt(), 'toolCallId': tool_call_id, 'delta': '"Paris"}'},
            {'type': 'TOOL_CALL_END', 'timestamp': IsInt(), 'toolCallId': tool_call_id},
            {
                'type': 'RUN_FINISHED',
                'timestamp': IsInt(),
                'threadId': thread_id,
                'runId': run_id,
            },
            {
                'type': 'RUN_STARTED',
                'timestamp': IsInt(),
                'threadId': thread_id,
                'runId': (run_id := IsSameStr()),
            },
            {
                'type': 'TEXT_MESSAGE_START',
                'timestamp': IsInt(),
                'messageId': (message_id := IsSameStr()),
                'role': 'assistant',
            },
            {
                'type': 'TEXT_MESSAGE_CONTENT',
                'timestamp': IsInt(),
                'messageId': message_id,
                'delta': '{"get_weather": "Tool result"}',
            },
            {'type': 'TEXT_MESSAGE_END', 'timestamp': IsInt(), 'messageId': message_id},