Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/mcp/server.md
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,7 @@ async def sampling_callback(
annotations=None,
meta=None,
),
meta=None,
)
]
"""
Expand Down
1 change: 0 additions & 1 deletion examples/pydantic_ai_examples/weather_agent_gradio.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,6 @@ def select_data(message: gr.SelectData) -> str:
past_messages = gr.State([])
chatbot = gr.Chatbot(
label='Packing Assistant',
type='messages',
avatar_images=(None, 'https://ai.pydantic.dev/img/logo-white.svg'),
examples=[
{'text': 'What is the weather like in Miami?'},
Expand Down
4 changes: 2 additions & 2 deletions examples/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,8 @@ dependencies = [
"rich>=13.9.2",
"uvicorn>=0.32.0",
"devtools>=0.12.2",
"gradio>=5.9.0",
"mcp[cli]>=1.4.1",
"gradio>=5.31.0",
"mcp[cli]>=1.25.0",
"modal>=1.0.4",
"duckdb>=1.3.2",
"datasets>=4.0.0",
Expand Down
10 changes: 7 additions & 3 deletions pydantic_ai_slim/pydantic_ai/_mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,11 +33,12 @@ def map_from_mcp_params(params: mcp_types.CreateMessageRequestParams) -> list[me
# TODO(Marcelo): We can reuse the `_map_tool_result_part` from the mcp module here.
if isinstance(content, mcp_types.TextContent):
user_part_content: str | Sequence[messages.UserContent] = content.text
else:
# image content
elif isinstance(content, mcp_types.ImageContent):
user_part_content = [
messages.BinaryContent(data=base64.b64decode(content.data), media_type=content.mimeType)
]
else:
raise NotImplementedError(f'Unsupported user content type: {type(content).__name__}')

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What are the new types we could support here and below? Add a comment at lease please

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the painful stuff is that content: SamplingMessageContentBlock | list[SamplingMessageContentBlock]

but I've moved the NotImplementedError into an explicit branch checking all rest-types and added an assert_never at the end

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • Could we support AudioContent easily in the ImageContent branch above?
  • Let's make it clear in the error which types are not yet supported (like in map_from_sampling_content) but relatively easily could be, vs anything that is fundamentally unsupported. Just so people know what their options are when they hit this error (e.g. contribute a fix)


request_parts.append(messages.UserPromptPart(content=user_part_content))
else:
Expand All @@ -47,7 +48,10 @@ def map_from_mcp_params(params: mcp_types.CreateMessageRequestParams) -> list[me
pai_messages.append(messages.ModelRequest(parts=request_parts))
request_parts = []

response_parts.append(map_from_sampling_content(content))
if isinstance(content, (mcp_types.TextContent, mcp_types.ImageContent, mcp_types.AudioContent)):
response_parts.append(map_from_sampling_content(content))
Comment thread
DouweM marked this conversation as resolved.
else:
raise NotImplementedError(f'Unsupported assistant content type: {type(content).__name__}')

if response_parts:
pai_messages.append(messages.ModelResponse(parts=response_parts))
Expand Down
47 changes: 39 additions & 8 deletions pydantic_ai_slim/pydantic_ai/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
from mcp.client.session import ClientSession, ElicitationFnT, LoggingFnT
from mcp.client.sse import sse_client
from mcp.client.stdio import StdioServerParameters, stdio_client
from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
from mcp.client.streamable_http import GetSessionIdCallback, streamable_http_client
from mcp.shared import exceptions as mcp_exceptions
from mcp.shared.context import RequestContext
from mcp.shared.message import SessionMessage
Expand Down Expand Up @@ -1134,8 +1134,10 @@ def _transport_client(
],
]: ...

# Only used by MCPServerSSE which has a hang bug (https://github.com/modelcontextprotocol/python-sdk/issues/1811).

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this just be on that subclass then? We may not need the _transport_client thing anymore if each class will have its own implementation of this entire thing

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeahp I wanted to keep this lean but I made that change now

# Remove pragma once https://github.com/modelcontextprotocol/python-sdk/pull/1838 is released.
@asynccontextmanager
async def client_streams(
async def client_streams( # pragma: no cover
self,
) -> AsyncIterator[
tuple[
Expand All @@ -1144,7 +1146,7 @@ async def client_streams(
]
]:
if self.http_client and self.headers:
raise ValueError('`http_client` is mutually exclusive with `headers`.') # pragma: no cover
raise ValueError('`http_client` is mutually exclusive with `headers`.')

transport_client_partial = functools.partial(
self._transport_client,
Expand All @@ -1154,7 +1156,8 @@ async def client_streams(
)

if self.http_client is not None:
# TODO: Clean up once https://github.com/modelcontextprotocol/python-sdk/pull/1177 lands.
# sse_client still uses the old httpx_client_factory API (streamable_http_client uses http_client).
# This wrapper is needed until sse_client is updated to match streamable_http_client's API.
@asynccontextmanager
async def httpx_client_factory(
headers: dict[str, str] | None = None,
Expand Down Expand Up @@ -1218,8 +1221,8 @@ def __get_pydantic_core_schema__(cls, _: Any, __: Any) -> CoreSchema:
)

@property
def _transport_client(self):
return sse_client # pragma: no cover
def _transport_client(self): # pragma: no cover
return sse_client

def __eq__(self, value: object, /) -> bool:
return super().__eq__(value) and isinstance(value, MCPServerSSE) and self.url == value.url
Expand Down Expand Up @@ -1282,8 +1285,36 @@ def __get_pydantic_core_schema__(cls, _: Any, __: Any) -> CoreSchema:
)

@property
def _transport_client(self):
return streamablehttp_client
def _transport_client(self): # pragma: no cover
return streamable_http_client

@asynccontextmanager
async def client_streams(
self,
) -> AsyncIterator[
tuple[
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
]
]:
# The new streamable_http_client API only accepts url and http_client.
# We need to create an httpx.AsyncClient with the desired timeout and headers.
if self.http_client is not None:
async with streamable_http_client(self.url, http_client=self.http_client) as (
read_stream,
write_stream,
*_,
):
yield read_stream, write_stream
else:
timeout = httpx.Timeout(self.timeout, read=self.read_timeout)
async with httpx.AsyncClient(timeout=timeout, headers=self.headers) as client:
async with streamable_http_client(self.url, http_client=client) as (
read_stream,
write_stream,
*_,
):
yield read_stream, write_stream
Comment thread
dsfaccini marked this conversation as resolved.
Outdated

def __eq__(self, value: object, /) -> bool:
return super().__eq__(value) and isinstance(value, MCPServerStreamableHTTP) and self.url == value.url
Expand Down
4 changes: 2 additions & 2 deletions pydantic_ai_slim/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -104,9 +104,9 @@ cli = [
"pyperclip>=1.9.0",
]
# MCP
mcp = ["mcp>=1.18.0"]
mcp = ["mcp>=1.25.0"]
# FastMCP
fastmcp = ["fastmcp>=2.12.0"]
fastmcp = ["fastmcp>=2.14.0"]
# Evals
evals = ["pydantic-evals=={{ version }}"]
# UI
Expand Down
7 changes: 6 additions & 1 deletion tests/test_mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -1678,7 +1678,11 @@ def test_map_from_pai_messages_with_binary_content():
assert system_prompt == ''
assert [m.model_dump(by_alias=True) for m in sampling_msgs] == snapshot(
[
{'role': 'user', 'content': {'type': 'text', 'text': 'text message', 'annotations': None, '_meta': None}},
{
'role': 'user',
'content': {'type': 'text', 'text': 'text message', 'annotations': None, '_meta': None},
'_meta': None,
},
{
'role': 'user',
'content': {
Expand All @@ -1688,6 +1692,7 @@ def test_map_from_pai_messages_with_binary_content():
'annotations': None,
'_meta': None,
},
'_meta': None,
},
]
)
Expand Down
Loading
Loading