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

Code before:

def rich_prompt_ask(prompt: str, *_args: Any, **_kwargs: Any) -> str:
    if prompt == 'Where would you like to fly from and to?':
        return 'SFO to ANC'
    elif prompt == 'What seat would you like?':
        return 'window seat with leg room'
    if prompt == 'Insert coins':
        return '1'
    elif prompt == 'Select product':
        return 'crisps'
    elif prompt == 'What is the capital of France?':  # pragma: no cover
        return 'Vichy'
    elif prompt == 'what is 1 + 1?':  # pragma: no cover
        return '2'
    else:  # pragma: no cover
        raise ValueError(f'Unexpected prompt: {prompt}')


def _create_no_op_validator() -> SchemaValidator:
    """Create a validator that accepts any input."""
    return SchemaValidator(schema=core_schema.any_schema())


class MockMCPServer(AbstractToolset[Any]):
    # Common tools that tests need
    _TOOLS = {
        'echo_deps',
        'image_generator',
        'return_fruit',
        'get_document',
        'get_company_logo',
    }

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Accept any arguments that real MCP servers might take."""
        # Store for reference but don't need to do anything with them
        self._args = args
        self._kwargs = kwargs

    @property
    def id(self) -> str | None:
        return None  # pragma: no cover

    @property
    def instructions(self) -> str | None:
        return None

    async def __aenter__(self) -> MockMCPServer:
        return self

    async def __aexit__(self, *args: Any) -> None:

Review comment (suggestion):
Instead of all of this work to keep `MCPServerStdio('python', args=['mcp_server.py'])` working, can we change the mcp-run-python example to use `mcp_server.py` as well?

Remember that the goal of this PR is to get rid of those extra downloads mcp-run-python is doing + the Deno requirement, _not_ to fully mock MCP (it's fine to call the local `tests/mcp_server.py` server which doesn't download anything)

Code after:
    if prompt == 'Insert coins':
        return '1'
    elif prompt == 'Select product':
        return 'crisps'
    elif prompt == 'What is the capital of France?':  # pragma: no cover
        return 'Vichy'
    elif prompt == 'what is 1 + 1?':  # pragma: no cover
        return '2'
    else:  # pragma: no cover
        raise ValueError(f'Unexpected prompt: {prompt}')


class MockMCPServer(AbstractToolset[Any]):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        pass

    @property
    def id(self) -> str | None:
        return None  # pragma: no cover

    @property
    def instructions(self) -> str | None:
        return None

    async def __aenter__(self) -> MockMCPServer:
        return self

    async def __aexit__(self, *args: Any) -> None:
        pass

    async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
        return {}

    async def call_tool(
        self, name: str, tool_args: dict[str, Any], ctx: RunContext[Any], tool: ToolsetTool[Any]
    ) -> Any:
        return None  # pragma: lax no cover


text_responses: dict[str, str | ToolCallPart | Sequence[ToolCallPart]] = {
    'Use the web to get the current time.': "In San Francisco, it's 8:21:41 pm PDT on Wednesday, August 6, 2025.",
    'Give me a sentence with the biggest news in AI this week.': 'Scientists have developed a universal AI detector that can identify deepfake videos.',
    'How many days between 2000-01-01 and 2025-03-18?': 'There are 9,208 days between January 1, 2000, and March 18, 2025.',
    'What is 7 plus 5?': 'The answer is 12.',
    'What is the weather like?': "It's currently raining in London.",
    'What is the weather like in West London and in Wiltshire?': (
        'The weather in West London is raining, while in Wiltshire it is sunny.'
    ),
    'What will the weather be like in Paris on Tuesday?': ToolCallPart(
        tool_name='weather_forecast', args={'location': 'Paris', 'forecast_date': '2030-01-01'}, tool_call_id='0001'
    ),