pydantic--pydantic-ai/tests/models/test_model_request_parameters.py

import pytest
from pydantic import TypeAdapter

from pydantic_ai.models import ModelRequestParameters, ToolDefinition
from pydantic_ai.native_tools import (
    CodeExecutionTool,
    ImageGenerationTool,
    MCPServerTool,
    MemoryTool,
    WebFetchTool,
    WebSearchTool,
    WebSearchUserLocation,
)
from pydantic_ai.output import StructuredOutputMode

from .._inline_snapshot import snapshot

ta = TypeAdapter(ModelRequestParameters)


def test_model_request_parameters_are_serializable():
    params = ModelRequestParameters(
        function_tools=[],
        native_tools=[],
        output_mode='text',
        allow_text_output=True,
        output_tools=[],
        output_object=None,
    )
    dumped = ta.dump_python(params)
    assert dumped == snapshot(
        {
            'function_tools': [],
            'native_tools': [],
            'output_mode': 'text',
            'output_object': None,
            'output_tools': [],
            'prompted_output_template': None,
            'allow_text_output': True,
            'allow_image_output': False,
            'instruction_parts': None,
            'thinking': None,
        }
    )
    assert ta.validate_python(dumped) == params

    params = ModelRequestParameters(
        function_tools=[ToolDefinition(name='test')],
        native_tools=[
            WebSearchTool(user_location=WebSearchUserLocation(city='New York', country='US')),
            CodeExecutionTool(),
            WebFetchTool(),
            ImageGenerationTool(size='1024x1024'),
            MemoryTool(),
            MCPServerTool(id='deepwiki', url='https://mcp.deepwiki.com/mcp'),
            MCPServerTool(id='github', url='https://api.githubcopilot.com/mcp'),
        ],
        output_mode='text',
        allow_text_output=True,
        output_tools=[ToolDefinition(name='final_result')],
        output_object=None,
    )
    dumped = ta.dump_python(params)
    assert dumped == snapshot(
        {
            'function_tools': [
                {
                    'name': 'test',
                    'parameters_json_schema': {'type': 'object', 'properties': {}},
                    'description': None,
                    'outer_typed_dict_key': None,
                    'strict': None,
                    'sequential': False,
                    'kind': 'function',
                    'metadata': None,
                    'timeout': None,
                    'defer_loading': False,
                    'unless_native': None,
                    'with_native': None,
                    'tool_kind': None,
                    'return_schema': None,
                    'include_return_schema': None,
                    'capability_id': None,
                }
            ],
            'native_tools': [
                {
                    'kind': 'web_search',
                    'optional': False,
                    'search_context_size': 'medium',
                    'user_location': {'city': 'New York', 'country': 'US'},
                    'blocked_domains': None,
                    'allowed_domains': None,
                    'max_uses': None,
                },
                {'kind': 'code_execution', 'optional': False, 'files': None},
                {
                    'kind': 'web_fetch',
                    'optional': False,
                    'max_uses': None,
                    'allowed_domains': None,
                    'blocked_domains': None,
                    'enable_citations': False,
                    'max_content_tokens': None,
                },
                {
                    'kind': 'image_generation',
                    'optional': False,
                    'action': 'auto',
                    'background': 'auto',
                    'input_fidelity': None,
                    'moderation': 'auto',
                    'model': None,
                    'output_compression': None,
                    'output_format': None,
                    'partial_images': 0,
                    'quality': 'auto',
                    'size': '1024x1024',
                    'aspect_ratio': None,
                },
                {'kind': 'memory', 'optional': False},
                {
                    'kind': 'mcp_server',
                    'optional': False,
                    'id': 'deepwiki',
                    'url': 'https://mcp.deepwiki.com/mcp',
                    'authorization_token': None,
                    'description': None,
                    'allowed_tools': None,
                    'headers': None,
                },
                {
                    'kind': 'mcp_server',
                    'optional': False,
                    'id': 'github',
                    'url': 'https://api.githubcopilot.com/mcp',
                    'authorization_token': None,
                    'description': None,
                    'allowed_tools': None,
                    'headers': None,
                },
            ],
            'output_mode': 'text',
            'output_object': None,
            'output_tools': [
                {
                    'name': 'final_result',
                    'parameters_json_schema': {'type': 'object', 'properties': {}},
                    'description': None,
                    'outer_typed_dict_key': None,
                    'strict': None,
                    'sequential': False,
                    'kind': 'function',
                    'metadata': None,
                    'timeout': None,
                    'defer_loading': False,
                    'unless_native': None,
                    'with_native': None,
                    'tool_kind': None,
                    'return_schema': None,
                    'include_return_schema': None,
                    'capability_id': None,
                }
            ],
            'prompted_output_template': None,
            'allow_text_output': True,
            'allow_image_output': False,
            'instruction_parts': None,
            'thinking': None,
        }
    )
    assert ta.validate_python(dumped) == params


@pytest.mark.parametrize(
    'output_mode, expected_allow_text',
    [
        ('tool', False),
        ('native', True),
        ('prompted', True),
    ],
)
def test_with_default_output_mode(output_mode: StructuredOutputMode, expected_allow_text: bool):
    params = ModelRequestParameters(output_mode='auto', allow_text_output=True)
    resolved = params.with_default_output_mode(output_mode)
    assert resolved.output_mode == output_mode
    assert resolved.allow_text_output == expected_allow_text


def test_with_default_output_mode_noop_when_not_auto():
    params = ModelRequestParameters(output_mode='tool', allow_text_output=False)
    resolved = params.with_default_output_mode('native')
    assert resolved is params


def test_with_default_output_mode_overrides_allow_text():
    params = ModelRequestParameters(output_mode='auto', allow_text_output=False)
    resolved = params.with_default_output_mode('native')
    assert resolved.output_mode == 'native'
    assert resolved.allow_text_output is True