Skip to content

Update cohere and MCP, add support for MCP ResourceLink returned from tools #2094

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 39 commits into from
Jul 24, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
39 commits
Select commit Hold shift + click to select a range
c1dfca6
Update tests to be compatible with new OpenAI, MistralAI and MCP vers…
medaminezghal Jun 28, 2025
8d42e69
Fix mcp compatibility
medaminezghal Jun 28, 2025
290becb
add ResourceLink type
medaminezghal Jun 28, 2025
a60e11d
add ResourceLink type
medaminezghal Jun 28, 2025
7d98900
Merge branch 'main' into update-versins
medaminezghal Jul 5, 2025
749c71d
Fix ResourceLink MCP types
medaminezghal Jul 8, 2025
7b83106
Fix
medaminezghal Jul 9, 2025
2091142
Merge branch 'main' into update-versins
medaminezghal Jul 9, 2025
6766617
Add tests
medaminezghal Jul 10, 2025
8c32055
Merge branch 'main' into update-versins
medaminezghal Jul 10, 2025
7e1c180
Fix tests fails for new cohere version
medaminezghal Jul 11, 2025
5d3112f
Merge branch 'main' into update-versins
medaminezghal Jul 11, 2025
5d05916
Merge branch 'main' into update-versins
medaminezghal Jul 17, 2025
c23076d
Revert change to cohere and fix mcp tests naming
medaminezghal Jul 17, 2025
e0e8369
Some fixes
medaminezghal Jul 17, 2025
2289561
Merge branch 'main' into update-versins
medaminezghal Jul 17, 2025
8f4f718
Update product_name.txt
medaminezghal Jul 17, 2025
7a985a6
Update test_mcp.py
medaminezghal Jul 17, 2025
bd9856e
Add compatibility with opentelemtry-api>=1.35
medaminezghal Jul 17, 2025
0f0c9de
Merge branch 'main' into update-versins
medaminezghal Jul 17, 2025
4b76833
Add @mcp.resource to ResourceLink tests tools
medaminezghal Jul 22, 2025
2cf4fcb
Fix conflicts with main
medaminezghal Jul 22, 2025
27f1a0d
Merge branch 'main' into update-versins
medaminezghal Jul 22, 2025
30b9d0d
Fix opentelemetry compatibility with version<1.35
medaminezghal Jul 23, 2025
32d0300
Fix
medaminezghal Jul 23, 2025
0b4e81a
Fix mcp_server.py
medaminezghal Jul 23, 2025
becbc57
Merge branch 'main' into update-versins
medaminezghal Jul 23, 2025
38d3b3c
Fix uv.lock
medaminezghal Jul 23, 2025
e9ab203
Fix tests
medaminezghal Jul 23, 2025
21ac93d
Revert Otel Upgrade
medaminezghal Jul 23, 2025
ce21660
Revert Otel Upgrade
medaminezghal Jul 23, 2025
1aa5c3d
Fix lint
medaminezghal Jul 23, 2025
2feaddc
Update MCP ResourceLink tests and add VCR cassettes
DouweM Jul 23, 2025
4d44f48
Fix tests
medaminezghal Jul 24, 2025
935c711
Fix coverage
medaminezghal Jul 24, 2025
de11883
Fix coverage
medaminezghal Jul 24, 2025
4ba2c0d
Fix coverage
medaminezghal Jul 24, 2025
d4a9f51
Fix coverage
medaminezghal Jul 24, 2025
411ef24
Merge branch 'main' into update-versins
medaminezghal Jul 24, 2025
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
5 changes: 4 additions & 1 deletion docs/mcp/server.md
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,10 @@ async def sampling_callback(
SamplingMessage(
role='user',
content=TextContent(
type='text', text='write a poem about socks', annotations=None
type='text',
text='write a poem about socks',
annotations=None,
meta=None,
),
)
]
Expand Down
35 changes: 23 additions & 12 deletions pydantic_ai_slim/pydantic_ai/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ async def direct_call_tool(
except McpError as e:
raise exceptions.ModelRetry(e.error.message)

content = [self._map_tool_result_part(part) for part in result.content]
content = [await self._map_tool_result_part(part) for part in result.content]

if result.isError:
text = '\n'.join(str(part) for part in content)
Expand Down Expand Up @@ -262,8 +262,8 @@ async def _sampling_callback(
model=self.sampling_model.model_name,
)

def _map_tool_result_part(
self, part: mcp_types.Content
async def _map_tool_result_part(
self, part: mcp_types.ContentBlock
) -> str | messages.BinaryContent | dict[str, Any] | list[Any]:
# See https://github.com/jlowin/fastmcp/blob/main/docs/servers/tools.mdx#return-values

Expand All @@ -285,18 +285,29 @@ def _map_tool_result_part(
) # pragma: no cover
elif isinstance(part, mcp_types.EmbeddedResource):
resource = part.resource
if isinstance(resource, mcp_types.TextResourceContents):
return resource.text
elif isinstance(resource, mcp_types.BlobResourceContents):
return messages.BinaryContent(
data=base64.b64decode(resource.blob),
media_type=resource.mimeType or 'application/octet-stream',
)
else:
assert_never(resource)
return self._get_content(resource)
elif isinstance(part, mcp_types.ResourceLink):
resource_result: mcp_types.ReadResourceResult = await self._client.read_resource(part.uri)
return (
self._get_content(resource_result.contents[0])
if len(resource_result.contents) == 1
else [self._get_content(resource) for resource in resource_result.contents]
)
else:
assert_never(part)

def _get_content(
self, resource: mcp_types.TextResourceContents | mcp_types.BlobResourceContents
) -> str | messages.BinaryContent:
if isinstance(resource, mcp_types.TextResourceContents):
return resource.text
elif isinstance(resource, mcp_types.BlobResourceContents):
return messages.BinaryContent(
data=base64.b64decode(resource.blob), media_type=resource.mimeType or 'application/octet-stream'
)
else:
assert_never(resource)


@dataclass
class MCPServerStdio(MCPServer):
Expand Down
12 changes: 6 additions & 6 deletions pydantic_ai_slim/pydantic_ai/models/cohere.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,15 @@
AssistantChatMessageV2,
AsyncClientV2,
ChatMessageV2,
ChatResponse,
SystemChatMessageV2,
TextAssistantMessageContentItem,
TextAssistantMessageV2ContentItem,
ToolCallV2,
ToolCallV2Function,
ToolChatMessageV2,
ToolV2,
ToolV2Function,
UserChatMessageV2,
V2ChatResponse,
)
from cohere.core.api_error import ApiError
from cohere.v2.client import OMIT
Expand Down Expand Up @@ -164,7 +164,7 @@ async def _chat(
messages: list[ModelMessage],
model_settings: CohereModelSettings,
model_request_parameters: ModelRequestParameters,
) -> ChatResponse:
) -> V2ChatResponse:
tools = self._get_tools(model_request_parameters)
cohere_messages = self._map_messages(messages)
try:
Expand All @@ -185,7 +185,7 @@ async def _chat(
raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
raise # pragma: no cover

def _process_response(self, response: ChatResponse) -> ModelResponse:
def _process_response(self, response: V2ChatResponse) -> ModelResponse:
"""Process a non-streamed response, and prepare a message to return."""
parts: list[ModelResponsePart] = []
if response.message.content is not None and len(response.message.content) > 0:
Expand Down Expand Up @@ -227,7 +227,7 @@ def _map_messages(self, messages: list[ModelMessage]) -> list[ChatMessageV2]:
assert_never(item)
message_param = AssistantChatMessageV2(role='assistant')
if texts:
message_param.content = [TextAssistantMessageContentItem(text='\n\n'.join(texts))]
message_param.content = [TextAssistantMessageV2ContentItem(text='\n\n'.join(texts))]
if tool_calls:
message_param.tool_calls = tool_calls
cohere_messages.append(message_param)
Expand Down Expand Up @@ -294,7 +294,7 @@ def _map_user_message(cls, message: ModelRequest) -> Iterable[ChatMessageV2]:
assert_never(part)


def _map_usage(response: ChatResponse) -> usage.Usage:
def _map_usage(response: V2ChatResponse) -> usage.Usage:
u = response.usage
if u is None:
return usage.Usage()
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 @@ -63,7 +63,7 @@ dependencies = [
logfire = ["logfire>=3.11.0"]
# Models
openai = ["openai>=1.92.0"]
cohere = ["cohere>=5.13.11; platform_system != 'Emscripten'"]
cohere = ["cohere>=5.16.0; platform_system != 'Emscripten'"]
vertexai = ["google-auth>=2.36.0", "requests>=2.32.2"]
google = ["google-genai>=1.24.0"]
anthropic = ["anthropic>=0.52.0"]
Expand All @@ -77,7 +77,7 @@ tavily = ["tavily-python>=0.5.0"]
# CLI
cli = ["rich>=13", "prompt-toolkit>=3", "argcomplete>=3.5.0"]
# MCP
mcp = ["mcp>=1.9.4; python_version >= '3.10'"]
mcp = ["mcp>=1.10.0; python_version >= '3.10'"]
# Evals
evals = ["pydantic-evals=={{ version }}"]
# A2A
Expand Down
1 change: 1 addition & 0 deletions tests/assets/product_name.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Pydantic AI
321 changes: 321 additions & 0 deletions tests/cassettes/test_mcp/test_tool_returning_audio_resource_link.yaml

Large diffs are not rendered by default.

447 changes: 447 additions & 0 deletions tests/cassettes/test_mcp/test_tool_returning_image_resource_link.yaml

Large diffs are not rendered by default.

Loading