pydantic-ai 0.0.41__tar.gz → 0.0.42__tar.gz

{pydantic_ai-0.0.41 → pydantic_ai-0.0.42}/.gitignore

@@ -1,5 +1,4 @@
 site
-.python-version
 .venv
 dist
 __pycache__
@@ -16,3 +15,4 @@ examples/pydantic_ai_examples/.chat_app_messages.sqlite
 .vscode/
 /question_graph_history.json
 /docs-site/.wrangler/
+/CLAUDE.md

{pydantic_ai-0.0.41 → pydantic_ai-0.0.42}/Makefile

@@ -27,6 +27,10 @@ lint: ## Lint the code
 	uv run ruff format --check
 	uv run ruff check
 
+.PHONY: lint-js
+lint-js: ## Lint JS and TS code
+	cd mcp-run-python && npm run lint
+
 .PHONY: typecheck-pyright
 typecheck-pyright:
 	@# PYRIGHT_PYTHON_IGNORE_WARNINGS avoids the overhead of making a request to github on every invocation
@@ -62,6 +66,11 @@ testcov: test ## Run tests and generate a coverage report
 	@echo "building coverage html"
 	@uv run coverage html
 
+.PHONY: test-mrp
+test-mrp: ## Build and  tests of mcp-run-python
+	cd mcp-run-python && npm run prepare
+	uv run --package mcp-run-python pytest mcp-run-python -v
+
 .PHONY: update-examples
 update-examples: ## Update documentation examples
 	uv run -m pytest --update-examples tests/test_examples.py

{pydantic_ai-0.0.41 → pydantic_ai-0.0.42}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai
-Version: 0.0.41
+Version: 0.0.42
 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[anthropic,bedrock,cli,cohere,groq,mistral,openai,vertexai]==0.0.41
+Requires-Dist: pydantic-ai-slim[anthropic,bedrock,cli,cohere,groq,mcp,mistral,openai,vertexai]==0.0.42
 Provides-Extra: examples
-Requires-Dist: pydantic-ai-examples==0.0.41; extra == 'examples'
+Requires-Dist: pydantic-ai-examples==0.0.42; extra == 'examples'
 Provides-Extra: logfire
 Requires-Dist: logfire>=2.3; extra == 'logfire'
 Description-Content-Type: text/markdown

{pydantic_ai-0.0.41 → pydantic_ai-0.0.42}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai"
-version = "0.0.41"
+version = "0.0.42"
 description = "Agent Framework / shim to use Pydantic with LLMs"
 authors = [
     { name = "Samuel Colvin", email = "samuel@pydantic.dev" },
@@ -36,7 +36,7 @@ classifiers = [
 ]
 requires-python = ">=3.9"
 dependencies = [
-    "pydantic-ai-slim[openai,vertexai,groq,anthropic,mistral,cohere,bedrock,cli]==0.0.41",
+    "pydantic-ai-slim[openai,vertexai,groq,anthropic,mistral,cohere,bedrock,cli,mcp]==0.0.42",
 ]
 
 [project.urls]
@@ -46,16 +46,17 @@ Documentation = "https://ai.pydantic.dev"
 Changelog = "https://github.com/pydantic/pydantic-ai/releases"
 
 [project.optional-dependencies]
-examples = ["pydantic-ai-examples==0.0.41"]
+examples = ["pydantic-ai-examples==0.0.42"]
 logfire = ["logfire>=2.3"]
 
 [tool.uv.sources]
 pydantic-ai-slim = { workspace = true }
 pydantic-graph = { workspace = true }
 pydantic-ai-examples = { workspace = true }
+mcp-run-python = { workspace = true }
 
 [tool.uv.workspace]
-members = ["pydantic_ai_slim", "pydantic_graph", "examples"]
+members = ["pydantic_ai_slim", "pydantic_graph", "examples", "mcp-run-python"]
 
 [dependency-groups]
 # dev dependencies are defined in `pydantic-ai-slim/pyproject.toml` to allow for minimal testing
@@ -82,6 +83,7 @@ line-length = 120
 target-version = "py39"
 include = [
     "pydantic_ai_slim/**/*.py",
+    "mcp-run-python/**/*.py",
     "pydantic_graph/**/*.py",
     "examples/**/*.py",
     "tests/**/*.py",
@@ -116,18 +118,22 @@ quote-style = "single"
 "tests/**/*.py" = ["D"]
 "docs/**/*.py" = ["D"]
 "examples/**/*.py" = ["D101", "D103"]
+"mcp-run-python/**/*.py" = ["D", "TID251"]
 
 [tool.pyright]
+pythonVersion = "3.12"
 typeCheckingMode = "strict"
 reportMissingTypeStubs = false
 reportUnnecessaryIsInstance = false
 reportUnnecessaryTypeIgnoreComment = true
-include = ["pydantic_ai_slim", "pydantic_graph", "tests", "examples"]
+reportMissingModuleSource = false
+include = ["pydantic_ai_slim", "mcp-run-python", "pydantic_graph", "tests", "examples"]
 venvPath = ".venv"
 # see https://github.com/microsoft/pyright/issues/7771 - we don't want to error on decorated functions in tests
 # which are not otherwise used
 executionEnvironments = [{ root = "tests", reportUnusedFunction = false }]
-exclude = ["examples/pydantic_ai_examples/weather_agent_gradio.py"]
+exclude = ["examples/pydantic_ai_examples/weather_agent_gradio.py", "mcp-run-python/node_modules"]
+extraPaths = ["mcp-run-python/stubs"]
 
 [tool.mypy]
 files = "tests/typed_agent.py,tests/typed_graph.py"
@@ -139,7 +145,11 @@ xfail_strict = true
 filterwarnings = [
     "error",
     # boto3
-    "ignore::DeprecationWarning:botocore.*"
+    "ignore::DeprecationWarning:botocore.*",
+    "ignore::RuntimeWarning:pydantic_ai.mcp",
+    # uvicorn (mcp server)
+    "ignore:websockets.legacy is deprecated.*:DeprecationWarning:websockets.legacy",
+    "ignore:websockets.server.WebSocketServerProtocol is deprecated:DeprecationWarning"
 ]
 
 # https://coverage.readthedocs.io/en/latest/config.html#run

pydantic_ai-0.0.42/tests/cassettes/test_mcp/test_agent_with_stdio_server.yaml

@@ -0,0 +1,205 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '530'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+    method: POST
+    parsed_body:
+      messages:
+      - content: What is 0 degrees Celsius in Fahrenheit?
+        role: user
+      model: gpt-4o
+      n: 1
+      stream: false
+      tool_choice: auto
+      tools:
+      - function:
+          description: "Convert Celsius to Fahrenheit.\n\n    Args:\n        celsius: Temperature in Celsius\n\n    Returns:\n
+            \       Temperature in Fahrenheit\n    "
+          name: celsius_to_fahrenheit
+          parameters:
+            properties:
+              celsius:
+                title: Celsius
+                type: number
+            required:
+            - celsius
+            title: celsius_to_fahrenheitArguments
+            type: object
+        type: function
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    headers:
+      access-control-expose-headers:
+      - X-Request-ID
+      alt-svc:
+      - h3=":443"; ma=86400
+      connection:
+      - keep-alive
+      content-length:
+      - '1085'
+      content-type:
+      - application/json
+      openai-organization:
+      - pydantic-28gund
+      openai-processing-ms:
+      - '594'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      choices:
+      - finish_reason: tool_calls
+        index: 0
+        logprobs: null
+        message:
+          annotations: []
+          content: null
+          refusal: null
+          role: assistant
+          tool_calls:
+          - function:
+              arguments: '{"celsius":0}'
+              name: celsius_to_fahrenheit
+            id: call_UNesABTXfwIkYdh3HzXWw2wD
+            type: function
+      created: 1741776872
+      id: chatcmpl-BAE1IwTZc7FHM4TkNeBoPylR7rtCT
+      model: gpt-4o-2024-08-06
+      object: chat.completion
+      service_tier: default
+      system_fingerprint: fp_eb9dce56a8
+      usage:
+        completion_tokens: 19
+        completion_tokens_details:
+          accepted_prediction_tokens: 0
+          audio_tokens: 0
+          reasoning_tokens: 0
+          rejected_prediction_tokens: 0
+        prompt_tokens: 82
+        prompt_tokens_details:
+          audio_tokens: 0
+          cached_tokens: 0
+        total_tokens: 101
+    status:
+      code: 200
+      message: OK
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '879'
+      content-type:
+      - application/json
+      cookie:
+      - __cf_bm=GeO8TCYhlEUIV63eLxM4nKUU2OLlG.f8tMvM9shFTc8-1741776873-1.0.1.1-zxkkWGCAPhJIA05Uwt3Ii3DCg9da6owy45bo_yaZ1YmsoihITJCgZzpA6H4eL0xzFRDWrWkEIQYaFEXLYcrLePwDMsgwNUJbEf6sg1vm2YQ;
+        _cfuvid=AI06nwzbBcwVRHXv_BRehX1K7p9oe1qUXFkzXBWEUW0-1741776873043-0.0.1.1-604800000
+      host:
+      - api.openai.com
+    method: POST
+    parsed_body:
+      messages:
+      - content: What is 0 degrees Celsius in Fahrenheit?
+        role: user
+      - role: assistant
+        tool_calls:
+        - function:
+            arguments: '{"celsius":0}'
+            name: celsius_to_fahrenheit
+          id: call_UNesABTXfwIkYdh3HzXWw2wD
+          type: function
+      - content: '{"meta":null,"content":[{"type":"text","text":"32.0","annotations":null}],"isError":false}'
+        role: tool
+        tool_call_id: call_UNesABTXfwIkYdh3HzXWw2wD
+      model: gpt-4o
+      n: 1
+      stream: false
+      tool_choice: auto
+      tools:
+      - function:
+          description: "Convert Celsius to Fahrenheit.\n\n    Args:\n        celsius: Temperature in Celsius\n\n    Returns:\n
+            \       Temperature in Fahrenheit\n    "
+          name: celsius_to_fahrenheit
+          parameters:
+            properties:
+              celsius:
+                title: Celsius
+                type: number
+            required:
+            - celsius
+            title: celsius_to_fahrenheitArguments
+            type: object
+        type: function
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    headers:
+      access-control-expose-headers:
+      - X-Request-ID
+      alt-svc:
+      - h3=":443"; ma=86400
+      connection:
+      - keep-alive
+      content-length:
+      - '849'
+      content-type:
+      - application/json
+      openai-organization:
+      - pydantic-28gund
+      openai-processing-ms:
+      - '415'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      choices:
+      - finish_reason: stop
+        index: 0
+        logprobs: null
+        message:
+          annotations: []
+          content: 0 degrees Celsius is 32.0 degrees Fahrenheit.
+          refusal: null
+          role: assistant
+      created: 1741776873
+      id: chatcmpl-BAE1Jy3AN974xW1pziTxd6wrxliCE
+      model: gpt-4o-2024-08-06
+      object: chat.completion
+      service_tier: default
+      system_fingerprint: fp_eb9dce56a8
+      usage:
+        completion_tokens: 13
+        completion_tokens_details:
+          accepted_prediction_tokens: 0
+          audio_tokens: 0
+          reasoning_tokens: 0
+          rejected_prediction_tokens: 0
+        prompt_tokens: 139
+        prompt_tokens_details:
+          audio_tokens: 0
+          cached_tokens: 0
+        total_tokens: 152
+    status:
+      code: 200
+      message: OK
+version: 1

{pydantic_ai-0.0.41 → pydantic_ai-0.0.42}/tests/graph/test_graph.py

@@ -393,7 +393,7 @@ async def test_iter_next_error(mock_snapshot_id: object):
 
         assert isinstance(n, BaseNode)
         n = await run.next()
-        assert n == snapshot(End(None))
+        assert n == snapshot(End(data=None))
 
         with pytest.raises(TypeError, match=r'`next` must be called with a `BaseNode` instance, got End\(data=None\).'):
             await run.next()

{pydantic_ai-0.0.41 → pydantic_ai-0.0.42}/tests/graph/test_persistence.py

@@ -287,7 +287,7 @@ async def test_rerun_node(mock_snapshot_id: object):
     node = Foo()
     async with graph.iter(node, persistence=sp) as run:
         end = await run.next()
-        assert end == snapshot(End(123))
+        assert end == snapshot(End(data=123))
 
     msg = "Incorrect snapshot status 'success', must be 'created' or 'pending'."
     with pytest.raises(GraphNodeStatusError, match=msg):

pydantic_ai-0.0.42/tests/mcp_server.py

@@ -0,0 +1,19 @@
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP('PydanticAI MCP Server')
+
+
+@mcp.tool()
+async def celsius_to_fahrenheit(celsius: float) -> float:
+    """Convert Celsius to Fahrenheit.
+
+    Args:
+        celsius: Temperature in Celsius
+
+    Returns:
+        Temperature in Fahrenheit
+    """
+    return (celsius * 9 / 5) + 32
+
+
+mcp.run()

{pydantic_ai-0.0.41 → pydantic_ai-0.0.42}/tests/test_examples.py

@@ -59,7 +59,7 @@ def find_filter_examples() -> Iterable[CodeExample]:
 
 
 @pytest.mark.parametrize('example', find_filter_examples(), ids=str)
-def test_docs_examples(
+def test_docs_examples(  # noqa: C901
     example: CodeExample,
     eval_example: EvalExample,
     mocker: MockerFixture,
@@ -78,6 +78,10 @@ def test_docs_examples(
     mocker.patch('random.randint', return_value=4)
     mocker.patch('rich.prompt.Prompt.ask', side_effect=rich_prompt_ask)
 
+    if sys.version_info >= (3, 10):
+        mocker.patch('pydantic_ai.mcp.MCPServerHTTP', return_value=MockMCPServer())
+        mocker.patch('mcp.server.fastmcp.FastMCP')
+
     env.set('OPENAI_API_KEY', 'testing')
     env.set('GEMINI_API_KEY', 'testing')
     env.set('GROQ_API_KEY', 'testing')
@@ -136,10 +140,13 @@ def test_docs_examples(
     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__'
         if eval_example.update_examples:  # pragma: no cover
-            module_dict = eval_example.run_print_update(example, call=call_name)
+            module_dict = 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_dict = eval_example.run_print_check(example, call=call_name, module_globals=test_globals)
 
         os.chdir(cwd)
         if title := opt_title:
@@ -182,7 +189,22 @@ def rich_prompt_ask(prompt: str, *_args: Any, **_kwargs: Any) -> str:
         raise ValueError(f'Unexpected prompt: {prompt}')
 
 
+class MockMCPServer:
+    is_running = True
+
+    async def __aenter__(self) -> MockMCPServer:
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        pass
+
+    @staticmethod
+    async def list_tools() -> list[None]:
+        return []
+
+
 text_responses: dict[str, str | ToolCallPart] = {
+    '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 the weather like in West London and in Wiltshire?': (
         'The weather in West London is raining, while in Wiltshire it is sunny.'
     ),

pydantic_ai-0.0.42/tests/test_mcp.py

@@ -0,0 +1,93 @@
+"""Tests for the MCP (Model Context Protocol) server implementation."""
+
+import pytest
+from dirty_equals import IsInstance
+from inline_snapshot import snapshot
+
+from pydantic_ai.agent import Agent
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, ToolCallPart, ToolReturnPart, UserPromptPart
+
+from .conftest import IsDatetime, try_import
+
+with try_import() as imports_successful:
+    from mcp.types import CallToolResult, TextContent
+
+    from pydantic_ai.mcp import MCPServerHTTP, MCPServerStdio
+    from pydantic_ai.models.openai import OpenAIModel
+    from pydantic_ai.providers.openai import OpenAIProvider
+
+
+pytestmark = [
+    pytest.mark.skipif(not imports_successful(), reason='mcp and openai not installed'),
+    pytest.mark.anyio,
+    pytest.mark.vcr,
+]
+
+
+async def test_stdio_server():
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    async with server:
+        tools = await server.list_tools()
+        assert len(tools) == 1
+        assert tools[0].name == 'celsius_to_fahrenheit'
+        assert tools[0].description.startswith('Convert Celsius to Fahrenheit.')
+
+        # Test calling the temperature conversion tool
+        result = await server.call_tool('celsius_to_fahrenheit', {'celsius': 0})
+        assert result.content == snapshot([TextContent(type='text', text='32.0')])
+
+
+def test_sse_server():
+    sse_server = MCPServerHTTP(url='http://localhost:8000/sse')
+    assert sse_server.url == 'http://localhost:8000/sse'
+
+
+async def test_agent_with_stdio_server(allow_model_requests: None, openai_api_key: str):
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    model = OpenAIModel('gpt-4o', provider=OpenAIProvider(api_key=openai_api_key))
+    agent = Agent(model, mcp_servers=[server])
+    async with agent.run_mcp_servers():
+        result = await agent.run('What is 0 degrees Celsius in Fahrenheit?')
+        assert result.data == snapshot('0 degrees Celsius is 32.0 degrees Fahrenheit.')
+        assert result.all_messages() == snapshot(
+            [
+                ModelRequest(
+                    parts=[UserPromptPart(content='What is 0 degrees Celsius in Fahrenheit?', timestamp=IsDatetime())]
+                ),
+                ModelResponse(
+                    parts=[
+                        ToolCallPart(
+                            tool_name='celsius_to_fahrenheit',
+                            args='{"celsius":0}',
+                            tool_call_id='call_UNesABTXfwIkYdh3HzXWw2wD',
+                        )
+                    ],
+                    model_name='gpt-4o-2024-08-06',
+                    timestamp=IsDatetime(),
+                ),
+                ModelRequest(
+                    parts=[
+                        ToolReturnPart(
+                            tool_name='celsius_to_fahrenheit',
+                            content=IsInstance(CallToolResult),
+                            tool_call_id='call_UNesABTXfwIkYdh3HzXWw2wD',
+                            timestamp=IsDatetime(),
+                        )
+                    ]
+                ),
+                ModelResponse(
+                    parts=[TextPart(content='0 degrees Celsius is 32.0 degrees Fahrenheit.')],
+                    model_name='gpt-4o-2024-08-06',
+                    timestamp=IsDatetime(),
+                ),
+            ]
+        )
+
+
+async def test_agent_with_server_not_running(openai_api_key: str):
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    model = OpenAIModel('gpt-4o', provider=OpenAIProvider(api_key=openai_api_key))
+    agent = Agent(model, mcp_servers=[server])
+    with pytest.raises(UserError, match='MCP server is not running'):
+        await agent.run('What is 0 degrees Celsius in Fahrenheit?')

{pydantic_ai-0.0.41 → pydantic_ai-0.0.42}/tests/test_tools.py

@@ -6,8 +6,9 @@ import pydantic_core
 import pytest
 from _pytest.logging import LogCaptureFixture
 from inline_snapshot import snapshot
-from pydantic import BaseModel, Field
-from pydantic_core import PydanticSerializationError
+from pydantic import BaseModel, Field, WithJsonSchema
+from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
+from pydantic_core import PydanticSerializationError, core_schema
 
 from pydantic_ai import Agent, RunContext, Tool, UserError
 from pydantic_ai.messages import (
@@ -100,8 +101,8 @@ def test_docstring_google(docstring_format: Literal['google', 'auto']):
             'description': 'Do foobar stuff, a lot.',
             'parameters_json_schema': {
                 'properties': {
-                    'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
-                    'bar': {'description': 'The bar thing.', 'title': 'Bar', 'type': 'string'},
+                    'foo': {'description': 'The foo thing.', 'type': 'integer'},
+                    'bar': {'description': 'The bar thing.', 'type': 'string'},
                 },
                 'required': ['foo', 'bar'],
                 'type': 'object',
@@ -138,7 +139,7 @@ def test_docstring_sphinx(docstring_format: Literal['sphinx', 'auto']):
             'name': 'sphinx_style_docstring',
             'description': 'Sphinx style docstring.',
             'parameters_json_schema': {
-                'properties': {'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'}},
+                'properties': {'foo': {'description': 'The foo thing.', 'type': 'integer'}},
                 'required': ['foo'],
                 'type': 'object',
                 'additionalProperties': False,
@@ -174,8 +175,8 @@ def test_docstring_numpy(docstring_format: Literal['numpy', 'auto']):
             'description': 'Numpy style docstring.',
             'parameters_json_schema': {
                 'properties': {
-                    'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
-                    'bar': {'description': 'The bar thing.', 'title': 'Bar', 'type': 'string'},
+                    'foo': {'description': 'The foo thing.', 'type': 'integer'},
+                    'bar': {'description': 'The bar thing.', 'type': 'string'},
                 },
                 'required': ['foo', 'bar'],
                 'type': 'object',
@@ -234,8 +235,8 @@ def test_docstring_google_no_body(docstring_format: Literal['google', 'auto']):
             'description': '',
             'parameters_json_schema': {
                 'properties': {
-                    'foo': {'description': 'The foo thing.', 'title': 'Foo', 'type': 'integer'},
-                    'bar': {'description': 'from fields', 'title': 'Bar', 'type': 'string'},
+                    'foo': {'description': 'The foo thing.', 'type': 'integer'},
+                    'bar': {'description': 'from fields', 'type': 'string'},
                 },
                 'required': ['foo', 'bar'],
                 'type': 'object',
@@ -266,8 +267,8 @@ def test_takes_just_model():
             'description': None,
             'parameters_json_schema': {
                 'properties': {
-                    'x': {'title': 'X', 'type': 'integer'},
-                    'y': {'title': 'Y', 'type': 'string'},
+                    'x': {'type': 'integer'},
+                    'y': {'type': 'string'},
                 },
                 'required': ['x', 'y'],
                 'title': 'Foo',
@@ -298,8 +299,8 @@ def test_takes_model_and_int():
                 '$defs': {
                     'Foo': {
                         'properties': {
-                            'x': {'title': 'X', 'type': 'integer'},
-                            'y': {'title': 'Y', 'type': 'string'},
+                            'x': {'type': 'integer'},
+                            'y': {'type': 'string'},
                         },
                         'required': ['x', 'y'],
                         'title': 'Foo',
@@ -308,7 +309,7 @@ def test_takes_model_and_int():
                 },
                 'properties': {
                     'model': {'$ref': '#/$defs/Foo'},
-                    'z': {'title': 'Z', 'type': 'integer'},
+                    'z': {'type': 'integer'},
                 },
                 'required': ['model', 'z'],
                 'type': 'object',
@@ -645,7 +646,7 @@ def test_json_schema_required_parameters(set_event_loop: None):
                 'outer_typed_dict_key': None,
                 'parameters_json_schema': {
                     'additionalProperties': False,
-                    'properties': {'a': {'title': 'A', 'type': 'integer'}, 'b': {'title': 'B', 'type': 'integer'}},
+                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                     'required': ['a'],
                     'type': 'object',
                 },
@@ -656,7 +657,7 @@ def test_json_schema_required_parameters(set_event_loop: None):
                 'outer_typed_dict_key': None,
                 'parameters_json_schema': {
                     'additionalProperties': False,
-                    'properties': {'a': {'title': 'A', 'type': 'integer'}, 'b': {'title': 'B', 'type': 'integer'}},
+                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                     'required': ['b'],
                     'type': 'object',
                 },
@@ -706,3 +707,47 @@ def test_call_tool_without_unrequired_parameters(set_event_loop: None):
         ]
     )
     assert tool_returns == snapshot([15, 17, 51, 68])
+
+
+def test_schema_generator():
+    class MyGenerateJsonSchema(GenerateJsonSchema):
+        def typed_dict_schema(self, schema: core_schema.TypedDictSchema) -> JsonSchemaValue:
+            # Add useless property titles just to show we can
+            s = super().typed_dict_schema(schema)
+            for p in s.get('properties', {}):
+                s['properties'][p]['title'] = f'{s["properties"][p].get("title")} title'
+            return s
+
+    agent = Agent(FunctionModel(get_json_schema))
+
+    def my_tool(x: Annotated[Union[str, None], WithJsonSchema({'type': 'string'})] = None, **kwargs: Any):
+        return x  # pragma: no cover
+
+    agent.tool_plain(name='my_tool_1')(my_tool)
+    agent.tool_plain(name='my_tool_2', schema_generator=MyGenerateJsonSchema)(my_tool)
+
+    result = agent.run_sync('Hello')
+    json_schema = json.loads(result.data)
+    assert json_schema == snapshot(
+        [
+            {
+                'description': '',
+                'name': 'my_tool_1',
+                'outer_typed_dict_key': None,
+                'parameters_json_schema': {
+                    'additionalProperties': True,
+                    'properties': {'x': {'type': 'string'}},
+                    'type': 'object',
+                },
+            },
+            {
+                'description': '',
+                'name': 'my_tool_2',
+                'outer_typed_dict_key': None,
+                'parameters_json_schema': {
+                    'properties': {'x': {'type': 'string', 'title': 'X title'}},
+                    'type': 'object',
+                },
+            },
+        ]
+    )