pydantic-ai 0.3.1__tar.gz → 0.3.2__tar.gz

{pydantic_ai-0.3.1 → pydantic_ai-0.3.2}/.gitignore

@@ -19,3 +19,4 @@ examples/pydantic_ai_examples/.chat_app_messages.sqlite
 node_modules/
 **.idea/
 .coverage*
+/test_tmp/

{pydantic_ai-0.3.1 → pydantic_ai-0.3.2}/Makefile

@@ -64,7 +64,7 @@ test: ## Run tests and collect coverage data
 	@uv run coverage report
 
 .PHONY: test-fast
-test-fast: ## Same as test except no coverage. ~1/4th the time depending on hardware.
+test-fast: ## Same as test except no coverage and 4x faster depending on hardware
 	uv run pytest -n auto --dist=loadgroup
 
 .PHONY: test-all-python
@@ -78,12 +78,12 @@ test-all-python: ## Run tests on Python 3.9 to 3.13
 	@uv run coverage report
 
 .PHONY: testcov
-testcov: test ## Run tests and generate a coverage report
+testcov: test ## Run tests and generate an HTML coverage report
 	@echo "building coverage html"
 	@uv run coverage html
 
 .PHONY: test-mrp
-test-mrp: ## Build and  tests of mcp-run-python
+test-mrp: ## Build and tests of mcp-run-python
 	cd mcp-run-python && deno task build
 	uv run --package mcp-run-python pytest mcp-run-python -v
 

{pydantic_ai-0.3.1 → pydantic_ai-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai
-Version: 0.3.1
+Version: 0.3.2
 Summary: Agent Framework / shim to use Pydantic with LLMs
 Project-URL: Homepage, https://ai.pydantic.dev
 Project-URL: Source, https://github.com/pydantic/pydantic-ai
@@ -28,9 +28,9 @@ Classifier: Topic :: Internet
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
-Requires-Dist: pydantic-ai-slim[a2a,anthropic,bedrock,cli,cohere,evals,google,groq,mcp,mistral,openai,vertexai]==0.3.1
+Requires-Dist: pydantic-ai-slim[a2a,anthropic,bedrock,cli,cohere,evals,google,groq,mcp,mistral,openai,vertexai]==0.3.2
 Provides-Extra: examples
-Requires-Dist: pydantic-ai-examples==0.3.1; extra == 'examples'
+Requires-Dist: pydantic-ai-examples==0.3.2; extra == 'examples'
 Provides-Extra: logfire
 Requires-Dist: logfire>=3.11.0; extra == 'logfire'
 Description-Content-Type: text/markdown

pydantic_ai-0.3.2/tests/example_modules/mcp_server.py

@@ -0,0 +1,26 @@
+from typing import Any
+
+from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.session import ServerSessionT
+from mcp.shared.context import LifespanContextT, RequestT
+
+mcp = FastMCP('PydanticAI MCP Server')
+
+
+@mcp.tool()
+async def echo_deps(ctx: Context[ServerSessionT, LifespanContextT, RequestT]) -> dict[str, Any]:
+    """Echo the run context.
+
+    Args:
+        ctx: Context object containing request and session information.
+
+    Returns:
+        Dictionary with an echo message and the deps.
+    """
+
+    deps: Any = getattr(ctx.request_context.meta, 'deps')
+    return {'echo': 'This is an echo message', 'deps': deps}
+
+
+if __name__ == '__main__':
+    mcp.run()

{pydantic_ai-0.3.1 → pydantic_ai-0.3.2}/tests/mcp_server.py

@@ -5,7 +5,7 @@ from typing import Any
 from mcp.server.fastmcp import Context, FastMCP, Image
 from mcp.server.session import ServerSessionT
 from mcp.shared.context import LifespanContextT, RequestT
-from mcp.types import BlobResourceContents, EmbeddedResource, TextResourceContents
+from mcp.types import BlobResourceContents, EmbeddedResource, SamplingMessage, TextContent, TextResourceContents
 from pydantic import AnyUrl
 
 mcp = FastMCP('PydanticAI MCP Server')
@@ -136,6 +136,23 @@ async def echo_deps(ctx: Context[ServerSessionT, LifespanContextT, RequestT]) ->
     return {'echo': 'This is an echo message', 'deps': deps}
 
 
+@mcp.tool()
+async def use_sampling(ctx: Context, foo: str) -> str:  # type: ignore
+    """Use sampling callback."""
+
+    result = await ctx.session.create_message(
+        [
+            SamplingMessage(role='assistant', content=TextContent(type='text', text='')),
+            SamplingMessage(role='user', content=TextContent(type='text', text=foo)),
+        ],
+        max_tokens=1_024,
+        system_prompt='this is a test of MCP sampling',
+        temperature=0.5,
+        stop_sequences=['potato'],
+    )
+    return result.model_dump_json(indent=2)
+
+
 @mcp._mcp_server.set_logging_level()  # pyright: ignore[reportPrivateUsage]
 async def set_logging_level(level: str) -> None:
     global log_level

{pydantic_ai-0.3.1 → pydantic_ai-0.3.2}/tests/models/cassettes/test_cohere/test_cohere_model_thinking_part.yaml

@@ -258,7 +258,7 @@ interactions:
                - **Bridge or Ferry:** If available, use a bridge or a ferry service. These are typically the safest and most reliable methods for crossing a river.
 
             3. **Prepare and Pack Essential Items:**
-               - **Life Jacket/Personal Floatation Device (PFD):** Always wear a life jacket or PFD when crossing a river, especially if swimming or using a boat.
+               - **Life Jacket/Personal Flotation Device (PFD):** Always wear a life jacket or PFD when crossing a river, especially if swimming or using a boat.
                - **First-Aid Kit:** Carry a basic first-aid kit to handle any minor injuries that might occur during the crossing.
                - **Map and Compass:** Navigate the river and its surroundings with the help of a map and compass, especially if you're in an unfamiliar area.
                - **Communication Device:** Have a means of communication, such as a satellite phone or a personal locator beacon, especially in remote areas.

{pydantic_ai-0.3.1 → pydantic_ai-0.3.2}/tests/models/test_cohere.py

@@ -512,7 +512,7 @@ Crossing a river can be a different challenge compared to crossing a street, and
    - **Bridge or Ferry:** If available, use a bridge or a ferry service. These are typically the safest and most reliable methods for crossing a river.
 
 3. **Prepare and Pack Essential Items:**
-   - **Life Jacket/Personal Floatation Device (PFD):** Always wear a life jacket or PFD when crossing a river, especially if swimming or using a boat.
+   - **Life Jacket/Personal Flotation Device (PFD):** Always wear a life jacket or PFD when crossing a river, especially if swimming or using a boat.
    - **First-Aid Kit:** Carry a basic first-aid kit to handle any minor injuries that might occur during the crossing.
    - **Map and Compass:** Navigate the river and its surroundings with the help of a map and compass, especially if you're in an unfamiliar area.
    - **Communication Device:** Have a means of communication, such as a satellite phone or a personal locator beacon, especially in remote areas.

pydantic_ai-0.3.2/tests/models/test_mcp_sampling.py

@@ -0,0 +1,135 @@
+import base64
+from dataclasses import dataclass
+from datetime import timezone
+from typing import Any
+from unittest.mock import AsyncMock
+
+import pytest
+from inline_snapshot import snapshot
+
+from pydantic_ai.agent import Agent
+from pydantic_ai.exceptions import UnexpectedModelBehavior
+from pydantic_ai.messages import BinaryContent, ModelRequest, ModelResponse, SystemPromptPart, TextPart, UserPromptPart
+from pydantic_ai.usage import Usage
+
+from ..conftest import IsNow, try_import
+
+with try_import() as imports_successful:
+    from mcp import CreateMessageResult
+    from mcp.types import TextContent
+
+    from pydantic_ai.models.mcp_sampling import MCPSamplingModel
+
+pytestmark = pytest.mark.skipif(not imports_successful(), reason='mcp package not installed')
+
+
+@dataclass
+class FakeSession:
+    create_message: Any
+
+
+def fake_session(create_message: Any) -> Any:
+    return FakeSession(create_message)
+
+
+def test_mcp_sampling_model():
+    model = MCPSamplingModel(fake_session(AsyncMock()))
+    assert model.model_name == 'mcp-sampling'
+    assert model.system == 'MCP'
+
+
+def test_assistant_text():
+    result = CreateMessageResult(
+        role='assistant', content=TextContent(type='text', text='text content'), model='test-model'
+    )
+    create_message = AsyncMock(return_value=result)
+    agent = Agent(model=MCPSamplingModel(fake_session(create_message)))
+
+    result = agent.run_sync('Hello')
+    assert result.output == snapshot('text content')
+    assert result.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='Hello',
+                        timestamp=IsNow(tz=timezone.utc),
+                    )
+                ]
+            ),
+            ModelResponse(
+                parts=[TextPart(content='text content')],
+                usage=Usage(requests=1),
+                model_name='test-model',
+                timestamp=IsNow(tz=timezone.utc),
+            ),
+        ]
+    )
+
+
+def test_user_text():
+    result = CreateMessageResult(role='user', content=TextContent(type='text', text='text content'), model='test-model')
+    create_message = AsyncMock(return_value=result)
+    agent = Agent(model=MCPSamplingModel(fake_session(create_message)))
+
+    expected_match = 'Unexpected result from MCP sampling, expected "assistant" role, got user.'
+    with pytest.raises(UnexpectedModelBehavior, match=expected_match):
+        agent.run_sync('Hello')
+
+
+def test_assistant_text_history():
+    result = CreateMessageResult(
+        role='assistant', content=TextContent(type='text', text='text content'), model='test-model'
+    )
+    create_message = AsyncMock(return_value=result)
+    agent = Agent(model=MCPSamplingModel(fake_session(create_message)), instructions='testing')
+
+    result = agent.run_sync('1')
+    result = agent.run_sync('2', message_history=result.all_messages())
+
+    assert result.output == snapshot('text content')
+    assert result.all_messages() == snapshot(
+        [
+            ModelRequest(parts=[UserPromptPart(content='1', timestamp=IsNow(tz=timezone.utc))], instructions='testing'),
+            ModelResponse(
+                parts=[TextPart(content='text content')],
+                usage=Usage(requests=1),
+                model_name='test-model',
+                timestamp=IsNow(tz=timezone.utc),
+            ),
+            ModelRequest(parts=[UserPromptPart(content='2', timestamp=IsNow(tz=timezone.utc))], instructions='testing'),
+            ModelResponse(
+                parts=[TextPart(content='text content')],
+                usage=Usage(requests=1),
+                model_name='test-model',
+                timestamp=IsNow(tz=timezone.utc),
+            ),
+        ]
+    )
+
+
+def test_assistant_text_history_complex():
+    history = [
+        ModelRequest(
+            parts=[
+                UserPromptPart(content='1'),
+                UserPromptPart(
+                    content=['a string', BinaryContent(data=base64.b64encode(b'data'), media_type='image/jpeg')]
+                ),
+                SystemPromptPart(content='system content'),
+            ]
+        ),
+        ModelResponse(
+            parts=[TextPart(content='text content')],
+            usage=Usage(requests=1),
+            model_name='test-model',
+        ),
+    ]
+
+    result = CreateMessageResult(
+        role='assistant', content=TextContent(type='text', text='text content'), model='test-model'
+    )
+    create_message = AsyncMock(return_value=result)
+    agent = Agent(model=MCPSamplingModel(fake_session(create_message)))
+    result = agent.run_sync('1', message_history=history)
+    assert result.output == snapshot('text content')

{pydantic_ai-0.3.1 → pydantic_ai-0.3.2}/tests/test_examples.py

@@ -3,13 +3,13 @@ from __future__ import annotations as _annotations
 import json
 import os
 import re
+import shutil
 import sys
 from collections.abc import AsyncIterator, Iterable, Sequence
 from dataclasses import dataclass
 from inspect import FrameInfo
 from io import StringIO
 from pathlib import Path
-from types import ModuleType
 from typing import Any
 
 import httpx
@@ -64,32 +64,44 @@ pytestmark = [
         reason='google-auth or logfire or google-provider not installed',
     ),
 ]
+code_examples: dict[str, CodeExample] = {}
 
 
 def find_filter_examples() -> Iterable[ParameterSet]:
     # Ensure this is run from the package root regardless of where/how the tests are run
-    os.chdir(Path(__file__).parent.parent)
+    root_dir = Path(__file__).parent.parent
+    os.chdir(root_dir)
 
     for ex in find_examples('docs', 'pydantic_ai_slim', 'pydantic_graph', 'pydantic_evals'):
         if ex.path.name != '_utils.py':
             try:
-                path = ex.path.relative_to(Path.cwd())
+                path = ex.path.relative_to(root_dir)
             except ValueError:
                 path = ex.path
             test_id = f'{path}:{ex.start_line}'
             prefix_settings = ex.prefix_settings()
-            if opt_title := prefix_settings.get('title'):
-                test_id += f':{opt_title}'
+            if title := prefix_settings.get('title'):
+                if title.endswith('.py'):
+                    code_examples[title] = ex
+                test_id += f':{title}'
             yield pytest.param(ex, id=test_id)
 
 
 @pytest.fixture
-def reset_cwd():
-    original_cwd = os.getcwd()
+def tmp_path_cwd(tmp_path: Path):
+    cwd = os.getcwd()
+
+    root_dir = Path(__file__).parent.parent
+    for file in (root_dir / 'tests' / 'example_modules').glob('*.py'):
+        shutil.copy(file, tmp_path)
+    sys.path.append(str(tmp_path))
+    os.chdir(tmp_path)
+
     try:
-        yield
+        yield tmp_path
     finally:
-        os.chdir(original_cwd)
+        os.chdir(cwd)
+        sys.path.remove(str(tmp_path))
 
 
 @pytest.mark.xdist_group(name='doc_tests')
@@ -101,8 +113,7 @@ def test_docs_examples(  # noqa: C901
     client_with_handler: ClientWithHandler,
     allow_model_requests: None,
     env: TestEnv,
-    tmp_path: Path,
-    reset_cwd: None,
+    tmp_path_cwd: Path,
 ):
     mocker.patch('pydantic_ai.agent.models.infer_model', side_effect=mock_infer_model)
     mocker.patch('pydantic_ai._utils.group_by_temporal', side_effect=mock_group_by_temporal)
@@ -142,14 +153,13 @@ def test_docs_examples(  # noqa: C901
     env.set('AWS_SECRET_ACCESS_KEY', 'testing')
     env.set('AWS_DEFAULT_REGION', 'us-east-1')
 
-    sys.path.append('tests/example_modules')
-
     prefix_settings = example.prefix_settings()
-    opt_title = prefix_settings.get('title')
     opt_test = prefix_settings.get('test', '')
     opt_lint = prefix_settings.get('lint', '')
     noqa = prefix_settings.get('noqa', '')
-    python_version = prefix_settings.get('py', None)
+    python_version = prefix_settings.get('py')
+    dunder_name = prefix_settings.get('dunder_name', '__main__')
+    requires = prefix_settings.get('requires')
 
     if python_version:
         python_version_info = tuple(int(v) for v in python_version.split('.'))
@@ -159,14 +169,12 @@ def test_docs_examples(  # noqa: C901
     if opt_test.startswith('skip') and opt_lint.startswith('skip'):
         pytest.skip('both running code and lint skipped')
 
-    if opt_title in {
-        'ai_q_and_a_run.py',
-        'count_down_from_persistence.py',
-        'generate_dataset_example.py',
-        'generate_dataset_example_json.py',
-        'save_load_dataset_example.py',
-    }:
-        os.chdir(tmp_path)
+    if requires:
+        for req in requires.split(','):
+            if ex := code_examples.get(req):
+                (tmp_path_cwd / req).write_text(ex.source)
+            else:  # pragma: no cover
+                raise KeyError(f'Example {req} not found, check the `requires` header of this example.')
 
     ruff_ignore: list[str] = ['D', 'Q001']
     # `from bank_database import DatabaseConn` wrongly sorted in imports
@@ -196,19 +204,12 @@ def test_docs_examples(  # noqa: C901
     if opt_test.startswith('skip'):
         print(opt_test[4:].lstrip(' -') or 'running code skipped')
     else:
-        test_globals: dict[str, str] = {}
-        if opt_title == 'mcp_client.py':
-            test_globals['__name__'] = '__test__'
+        test_globals: dict[str, str] = {'__name__': dunder_name}
+
         if eval_example.update_examples:  # pragma: lax no cover
-            module_dict = eval_example.run_print_update(example, call=call_name, module_globals=test_globals)
+            eval_example.run_print_update(example, call=call_name, module_globals=test_globals)
         else:
-            module_dict = eval_example.run_print_check(example, call=call_name, module_globals=test_globals)
-
-        if title := opt_title:
-            if title.endswith('.py'):
-                module_name = title[:-3]
-                sys.modules[module_name] = module = ModuleType(module_name)
-                module.__dict__.update(module_dict)
+            eval_example.run_print_check(example, call=call_name, module_globals=test_globals)
 
 
 def print_callback(s: str) -> str:
@@ -432,6 +433,10 @@ text_responses: dict[str, str | ToolCallPart] = {
             'explanation': 'I am not equipped to provide travel information, such as flights from Amsterdam to Mexico City.'
         },
     ),
+    'Create an image of a robot in a punk style.': ToolCallPart(
+        tool_name='image_generator', args={'subject': 'robot', 'style': 'punk'}, tool_call_id='0001'
+    ),
+    "subject='robot' style='punk'": '<svg/>',
 }
 
 tool_responses: dict[tuple[str, str], str] = {
@@ -676,6 +681,8 @@ async def model_logic(  # noqa: C901
                 )
             ]
         )
+    elif isinstance(m, ToolReturnPart) and m.tool_name == 'image_generator':
+        return ModelResponse(parts=[TextPart('Image file written to robot_punk.svg.')])
     else:
         sys.stdout.write(str(debug.format(messages, info)))
         raise RuntimeError(f'Unexpected message: {m}')

{pydantic_ai-0.3.1 → pydantic_ai-0.3.2}/tests/test_mcp.py

@@ -1,21 +1,24 @@
 """Tests for the MCP (Model Context Protocol) server implementation."""
 
+import base64
 import re
+from datetime import timezone
 from pathlib import Path
-from typing import Any, Final
+from typing import Any
 from unittest.mock import AsyncMock, patch
 
 import pytest
 from inline_snapshot import snapshot
 
 from pydantic_ai.agent import Agent
-from pydantic_ai.exceptions import ModelRetry, UserError
+from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError
 from pydantic_ai.messages import (
     BinaryContent,
     ModelRequest,
     ModelResponse,
     RetryPromptPart,
     TextPart,
+    ThinkingPart,
     ToolCallPart,
     ToolReturnPart,
     UserPromptPart,
@@ -24,19 +27,19 @@ from pydantic_ai.models.test import TestModel
 from pydantic_ai.tools import RunContext
 from pydantic_ai.usage import Usage
 
-from .conftest import IsDatetime, IsStr, try_import
+from .conftest import IsDatetime, IsNow, IsStr, try_import
 
 with try_import() as imports_successful:
-    from mcp import ErrorData, McpError
+    from mcp import ErrorData, McpError, SamplingMessage
+    from mcp.types import CreateMessageRequestParams, ImageContent, TextContent
 
+    from pydantic_ai._mcp import map_from_mcp_params, map_from_model_response
     from pydantic_ai.mcp import CallToolFunc, MCPServerSSE, MCPServerStdio, ToolResult
     from pydantic_ai.models.google import GoogleModel
     from pydantic_ai.models.openai import OpenAIModel
     from pydantic_ai.providers.google import GoogleProvider
     from pydantic_ai.providers.openai import OpenAIProvider
 
-TOOL_COUNT: Final[int] = 12
-
 pytestmark = [
     pytest.mark.skipif(not imports_successful(), reason='mcp and openai not installed'),
     pytest.mark.anyio,
@@ -55,7 +58,7 @@ async def test_stdio_server():
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
     async with server:
         tools = await server.list_tools()
-        assert len(tools) == TOOL_COUNT
+        assert len(tools) == snapshot(13)
         assert tools[0].name == 'celsius_to_fahrenheit'
         assert tools[0].description.startswith('Convert Celsius to Fahrenheit.')
 
@@ -64,6 +67,13 @@ async def test_stdio_server():
         assert result == snapshot('32.0')
 
 
+async def test_reentrant_context_manager():
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    async with server:
+        async with server:
+            pass
+
+
 async def test_stdio_server_with_tool_prefix():
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'], tool_prefix='foo')
     async with server:
@@ -76,7 +86,7 @@ async def test_stdio_server_with_cwd():
     server = MCPServerStdio('python', ['mcp_server.py'], cwd=test_dir)
     async with server:
         tools = await server.list_tools()
-        assert len(tools) == TOOL_COUNT
+        assert len(tools) == snapshot(13)
 
 
 async def test_process_tool_call() -> None:
@@ -104,7 +114,7 @@ async def test_process_tool_call() -> None:
 def test_sse_server():
     sse_server = MCPServerSSE(url='http://localhost:8000/sse')
     assert sse_server.url == 'http://localhost:8000/sse'
-    assert sse_server._get_log_level() is None  # pyright: ignore[reportPrivateUsage]
+    assert sse_server.log_level is None
 
 
 def test_sse_server_with_header_and_timeout():
@@ -119,7 +129,7 @@ def test_sse_server_with_header_and_timeout():
     assert sse_server.headers is not None and sse_server.headers['my-custom-header'] == 'my-header-value'
     assert sse_server.timeout == 10
     assert sse_server.sse_read_timeout == 100
-    assert sse_server._get_log_level() == 'info'  # pyright: ignore[reportPrivateUsage]
+    assert sse_server.log_level == 'info'
 
 
 @pytest.mark.vcr()
@@ -240,10 +250,10 @@ async def test_agent_with_server_not_running(openai_api_key: str):
 
 async def test_log_level_unset():
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
-    assert server._get_log_level() is None  # pyright: ignore[reportPrivateUsage]
+    assert server.log_level is None
     async with server:
         tools = await server.list_tools()
-        assert len(tools) == TOOL_COUNT
+        assert len(tools) == snapshot(13)
         assert tools[10].name == 'get_log_level'
 
         result = await server.call_tool('get_log_level', {})
@@ -252,7 +262,7 @@ async def test_log_level_unset():
 
 async def test_log_level_set():
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'], log_level='info')
-    assert server._get_log_level() == 'info'  # pyright: ignore[reportPrivateUsage]
+    assert server.log_level == 'info'
     async with server:
         result = await server.call_tool('get_log_level', {})
         assert result == snapshot('info')
@@ -963,6 +973,30 @@ async def test_tool_returning_multiple_items(allow_model_requests: None, agent:
         )
 
 
+async def test_client_sampling():
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    server.sampling_model = TestModel(custom_output_text='sampling model response')
+    async with server:
+        result = await server.call_tool('use_sampling', {'foo': 'bar'})
+        assert result == snapshot(
+            {
+                'meta': None,
+                'role': 'assistant',
+                'content': {'type': 'text', 'text': 'sampling model response', 'annotations': None},
+                'model': 'test',
+                'stopReason': None,
+            }
+        )
+
+
+async def test_client_sampling_disabled():
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'], allow_sampling=False)
+    server.sampling_model = TestModel(custom_output_text='sampling model response')
+    async with server:
+        with pytest.raises(ModelRetry, match='Error executing tool use_sampling: Sampling not supported'):
+            await server.call_tool('use_sampling', {'foo': 'bar'})
+
+
 async def test_mcp_server_raises_mcp_error(allow_model_requests: None, agent: Agent) -> None:
     server = agent._mcp_servers[0]  # pyright: ignore[reportPrivateUsage]
 
@@ -976,3 +1010,52 @@ async def test_mcp_server_raises_mcp_error(allow_model_requests: None, agent: Ag
         ):
             with pytest.raises(ModelRetry, match='Test MCP error conversion'):
                 await server.call_tool('test_tool', {})
+
+
+def test_map_from_mcp_params_model_request():
+    params = CreateMessageRequestParams(
+        messages=[
+            SamplingMessage(role='user', content=TextContent(type='text', text='xx')),
+            SamplingMessage(
+                role='user',
+                content=ImageContent(type='image', data=base64.b64encode(b'img').decode(), mimeType='image/png'),
+            ),
+        ],
+        maxTokens=8,
+    )
+    pai_messages = map_from_mcp_params(params)
+    assert pai_messages == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(content='xx', timestamp=IsNow(tz=timezone.utc)),
+                    UserPromptPart(
+                        content=[BinaryContent(data=b'img', media_type='image/png')], timestamp=IsNow(tz=timezone.utc)
+                    ),
+                ]
+            )
+        ]
+    )
+
+
+def test_map_from_mcp_params_model_response():
+    params = CreateMessageRequestParams(
+        messages=[
+            SamplingMessage(role='assistant', content=TextContent(type='text', text='xx')),
+        ],
+        maxTokens=8,
+    )
+    pai_messages = map_from_mcp_params(params)
+    assert pai_messages == snapshot(
+        [
+            ModelResponse(
+                parts=[TextPart(content='xx')],
+                timestamp=IsNow(tz=timezone.utc),
+            )
+        ]
+    )
+
+
+def test_map_from_model_response():
+    with pytest.raises(UnexpectedModelBehavior, match='Unexpected part type: ThinkingPart, expected TextPart'):
+        map_from_model_response(ModelResponse(parts=[ThinkingPart(content='Thinking...')]))