pydantic-ai-slim 0.2.16__tar.gz → 0.2.18__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of pydantic-ai-slim might be problematic. Click here for more details.

Files changed (76) hide show
  1. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/PKG-INFO +5 -5
  2. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/_agent_graph.py +10 -0
  3. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/mcp.py +187 -53
  4. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/messages.py +43 -13
  5. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/__init__.py +91 -4
  6. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/anthropic.py +12 -16
  7. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/bedrock.py +23 -15
  8. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/gemini.py +15 -13
  9. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/google.py +11 -11
  10. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/instrumented.py +98 -32
  11. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/openai.py +56 -36
  12. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/openai.py +9 -2
  13. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/google.py +1 -1
  14. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pyproject.toml +1 -1
  15. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/.gitignore +0 -0
  16. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/LICENSE +0 -0
  17. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/README.md +0 -0
  18. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/__init__.py +0 -0
  19. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/__main__.py +0 -0
  20. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/_a2a.py +0 -0
  21. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/_cli.py +0 -0
  22. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/_function_schema.py +0 -0
  23. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/_griffe.py +0 -0
  24. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/_output.py +0 -0
  25. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/_parts_manager.py +0 -0
  26. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/_system_prompt.py +0 -0
  27. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/_utils.py +0 -0
  28. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/agent.py +0 -0
  29. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/common_tools/__init__.py +0 -0
  30. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/common_tools/duckduckgo.py +0 -0
  31. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/common_tools/tavily.py +0 -0
  32. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/direct.py +0 -0
  33. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/exceptions.py +0 -0
  34. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/ext/__init__.py +0 -0
  35. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/ext/langchain.py +0 -0
  36. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/format_as_xml.py +0 -0
  37. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/format_prompt.py +0 -0
  38. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/cohere.py +0 -0
  39. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/fallback.py +0 -0
  40. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/function.py +0 -0
  41. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/groq.py +0 -0
  42. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/mistral.py +0 -0
  43. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/test.py +0 -0
  44. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/models/wrapper.py +0 -0
  45. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/__init__.py +0 -0
  46. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/_json_schema.py +0 -0
  47. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/amazon.py +0 -0
  48. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/anthropic.py +0 -0
  49. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/cohere.py +0 -0
  50. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/deepseek.py +0 -0
  51. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/google.py +0 -0
  52. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/grok.py +0 -0
  53. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/meta.py +0 -0
  54. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/mistral.py +0 -0
  55. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/profiles/qwen.py +0 -0
  56. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/__init__.py +0 -0
  57. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/anthropic.py +0 -0
  58. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/azure.py +0 -0
  59. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/bedrock.py +0 -0
  60. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/cohere.py +0 -0
  61. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/deepseek.py +0 -0
  62. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/fireworks.py +0 -0
  63. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/google_gla.py +0 -0
  64. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/google_vertex.py +0 -0
  65. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/grok.py +0 -0
  66. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/groq.py +0 -0
  67. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/heroku.py +0 -0
  68. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/mistral.py +0 -0
  69. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/openai.py +0 -0
  70. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/openrouter.py +0 -0
  71. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/providers/together.py +0 -0
  72. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/py.typed +0 -0
  73. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/result.py +0 -0
  74. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/settings.py +0 -0
  75. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/tools.py +0 -0
  76. {pydantic_ai_slim-0.2.16 → pydantic_ai_slim-0.2.18}/pydantic_ai/usage.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pydantic-ai-slim
3
- Version: 0.2.16
3
+ Version: 0.2.18
4
4
  Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
5
5
  Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>
6
6
  License-Expression: MIT
@@ -30,11 +30,11 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
30
30
  Requires-Dist: griffe>=1.3.2
31
31
  Requires-Dist: httpx>=0.27
32
32
  Requires-Dist: opentelemetry-api>=1.28.0
33
- Requires-Dist: pydantic-graph==0.2.16
33
+ Requires-Dist: pydantic-graph==0.2.18
34
34
  Requires-Dist: pydantic>=2.10
35
35
  Requires-Dist: typing-inspection>=0.4.0
36
36
  Provides-Extra: a2a
37
- Requires-Dist: fasta2a==0.2.16; extra == 'a2a'
37
+ Requires-Dist: fasta2a==0.2.18; extra == 'a2a'
38
38
  Provides-Extra: anthropic
39
39
  Requires-Dist: anthropic>=0.52.0; extra == 'anthropic'
40
40
  Provides-Extra: bedrock
@@ -48,7 +48,7 @@ Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == '
48
48
  Provides-Extra: duckduckgo
49
49
  Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
50
50
  Provides-Extra: evals
51
- Requires-Dist: pydantic-evals==0.2.16; extra == 'evals'
51
+ Requires-Dist: pydantic-evals==0.2.18; extra == 'evals'
52
52
  Provides-Extra: google
53
53
  Requires-Dist: google-genai>=1.15.0; extra == 'google'
54
54
  Provides-Extra: groq
@@ -56,7 +56,7 @@ Requires-Dist: groq>=0.15.0; extra == 'groq'
56
56
  Provides-Extra: logfire
57
57
  Requires-Dist: logfire>=3.11.0; extra == 'logfire'
58
58
  Provides-Extra: mcp
59
- Requires-Dist: mcp>=1.9.0; (python_version >= '3.10') and extra == 'mcp'
59
+ Requires-Dist: mcp>=1.9.4; (python_version >= '3.10') and extra == 'mcp'
60
60
  Provides-Extra: mistral
61
61
  Requires-Dist: mistralai>=1.2.5; extra == 'mistral'
62
62
  Provides-Extra: openai
@@ -183,6 +183,16 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
183
183
 
184
184
  if user_prompt is not None:
185
185
  parts.append(_messages.UserPromptPart(user_prompt))
186
+ elif (
187
+ len(parts) == 0
188
+ and message_history
189
+ and (last_message := message_history[-1])
190
+ and isinstance(last_message, _messages.ModelRequest)
191
+ ):
192
+ # Drop last message that came from history and reuse its parts
193
+ messages.pop()
194
+ parts.extend(last_message.parts)
195
+
186
196
  return messages, _messages.ModelRequest(parts, instructions=instructions)
187
197
 
188
198
  async def _reevaluate_dynamic_prompts(
@@ -1,27 +1,32 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  import base64
4
+ import functools
4
5
  import json
5
6
  from abc import ABC, abstractmethod
6
7
  from collections.abc import AsyncIterator, Sequence
7
- from contextlib import AsyncExitStack, asynccontextmanager
8
+ from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
8
9
  from dataclasses import dataclass
9
10
  from pathlib import Path
10
11
  from types import TracebackType
11
- from typing import Any
12
+ from typing import Any, Callable
12
13
 
13
14
  import anyio
15
+ import httpx
14
16
  from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
17
+ from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
18
+ from mcp.shared.message import SessionMessage
15
19
  from mcp.types import (
20
+ AudioContent,
16
21
  BlobResourceContents,
22
+ Content,
17
23
  EmbeddedResource,
18
24
  ImageContent,
19
- JSONRPCMessage,
20
25
  LoggingLevel,
21
26
  TextContent,
22
27
  TextResourceContents,
23
28
  )
24
- from typing_extensions import Self, assert_never
29
+ from typing_extensions import Self, assert_never, deprecated
25
30
 
26
31
  from pydantic_ai.exceptions import ModelRetry
27
32
  from pydantic_ai.messages import BinaryContent
@@ -37,7 +42,7 @@ except ImportError as _import_error:
37
42
  'you can use the `mcp` optional group — `pip install "pydantic-ai-slim[mcp]"`'
38
43
  ) from _import_error
39
44
 
40
- __all__ = 'MCPServer', 'MCPServerStdio', 'MCPServerHTTP'
45
+ __all__ = 'MCPServer', 'MCPServerStdio', 'MCPServerHTTP', 'MCPServerSSE', 'MCPServerStreamableHTTP'
41
46
 
42
47
 
43
48
  class MCPServer(ABC):
@@ -56,8 +61,8 @@ class MCPServer(ABC):
56
61
  """
57
62
 
58
63
  _client: ClientSession
59
- _read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
60
- _write_stream: MemoryObjectSendStream[JSONRPCMessage]
64
+ _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
65
+ _write_stream: MemoryObjectSendStream[SessionMessage]
61
66
  _exit_stack: AsyncExitStack
62
67
 
63
68
  @abstractmethod
@@ -66,8 +71,8 @@ class MCPServer(ABC):
66
71
  self,
67
72
  ) -> AsyncIterator[
68
73
  tuple[
69
- MemoryObjectReceiveStream[JSONRPCMessage | Exception],
70
- MemoryObjectSendStream[JSONRPCMessage],
74
+ MemoryObjectReceiveStream[SessionMessage | Exception],
75
+ MemoryObjectSendStream[SessionMessage],
71
76
  ]
72
77
  ]:
73
78
  """Create the streams for the MCP server."""
@@ -158,9 +163,7 @@ class MCPServer(ABC):
158
163
  await self._exit_stack.aclose()
159
164
  self.is_running = False
160
165
 
161
- def _map_tool_result_part(
162
- self, part: TextContent | ImageContent | EmbeddedResource
163
- ) -> str | BinaryContent | dict[str, Any] | list[Any]:
166
+ def _map_tool_result_part(self, part: Content) -> str | BinaryContent | dict[str, Any] | list[Any]:
164
167
  # See https://github.com/jlowin/fastmcp/blob/main/docs/servers/tools.mdx#return-values
165
168
 
166
169
  if isinstance(part, TextContent):
@@ -173,6 +176,10 @@ class MCPServer(ABC):
173
176
  return text
174
177
  elif isinstance(part, ImageContent):
175
178
  return BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType)
179
+ elif isinstance(part, AudioContent):
180
+ # NOTE: The FastMCP server doesn't support audio content.
181
+ # See <https://github.com/modelcontextprotocol/python-sdk/issues/952> for more details.
182
+ return BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType) # pragma: no cover
176
183
  elif isinstance(part, EmbeddedResource):
177
184
  resource = part.resource
178
185
  if isinstance(resource, TextResourceContents):
@@ -266,8 +273,8 @@ class MCPServerStdio(MCPServer):
266
273
  self,
267
274
  ) -> AsyncIterator[
268
275
  tuple[
269
- MemoryObjectReceiveStream[JSONRPCMessage | Exception],
270
- MemoryObjectSendStream[JSONRPCMessage],
276
+ MemoryObjectReceiveStream[SessionMessage | Exception],
277
+ MemoryObjectSendStream[SessionMessage],
271
278
  ]
272
279
  ]:
273
280
  server = StdioServerParameters(command=self.command, args=list(self.args), env=self.env, cwd=self.cwd)
@@ -285,47 +292,40 @@ class MCPServerStdio(MCPServer):
285
292
 
286
293
 
287
294
  @dataclass
288
- class MCPServerHTTP(MCPServer):
289
- """An MCP server that connects over streamable HTTP connections.
295
+ class _MCPServerHTTP(MCPServer):
296
+ url: str
297
+ """The URL of the endpoint on the MCP server."""
290
298
 
291
- This class implements the SSE transport from the MCP specification.
292
- See <https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse> for more information.
299
+ headers: dict[str, Any] | None = None
300
+ """Optional HTTP headers to be sent with each request to the endpoint.
293
301
 
294
- The name "HTTP" is used since this implemented will be adapted in future to use the new
295
- [Streamable HTTP](https://github.com/modelcontextprotocol/specification/pull/206) currently in development.
302
+ These headers will be passed directly to the underlying `httpx.AsyncClient`.
303
+ Useful for authentication, custom headers, or other HTTP-specific configurations.
296
304
 
297
305
  !!! note
298
- Using this class as an async context manager will create a new pool of HTTP connections to connect
299
- to a server which should already be running.
306
+ You can either pass `headers` or `http_client`, but not both.
300
307
 
301
- Example:
302
- ```python {py="3.10"}
303
- from pydantic_ai import Agent
304
- from pydantic_ai.mcp import MCPServerHTTP
308
+ See [`MCPServerHTTP.http_client`][pydantic_ai.mcp.MCPServerHTTP.http_client] for more information.
309
+ """
305
310
 
306
- server = MCPServerHTTP('http://localhost:3001/sse') # (1)!
307
- agent = Agent('openai:gpt-4o', mcp_servers=[server])
311
+ http_client: httpx.AsyncClient | None = None
312
+ """An `httpx.AsyncClient` to use with the endpoint.
308
313
 
309
- async def main():
310
- async with agent.run_mcp_servers(): # (2)!
311
- ...
312
- ```
314
+ This client may be configured to use customized connection parameters like self-signed certificates.
313
315
 
314
- 1. E.g. you might be connecting to a server run with [`mcp-run-python`](../mcp/run-python.md).
315
- 2. This will connect to a server running on `localhost:3001`.
316
- """
316
+ !!! note
317
+ You can either pass `headers` or `http_client`, but not both.
317
318
 
318
- url: str
319
- """The URL of the SSE endpoint on the MCP server.
319
+ If you want to use both, you can pass the headers to the `http_client` instead.
320
320
 
321
- For example for a server running locally, this might be `http://localhost:3001/sse`.
322
- """
321
+ ```python {py="3.10" test="skip"}
322
+ import httpx
323
323
 
324
- headers: dict[str, Any] | None = None
325
- """Optional HTTP headers to be sent with each request to the SSE endpoint.
324
+ from pydantic_ai.mcp import MCPServerSSE
326
325
 
327
- These headers will be passed directly to the underlying `httpx.AsyncClient`.
328
- Useful for authentication, custom headers, or other HTTP-specific configurations.
326
+ http_client = httpx.AsyncClient(headers={'Authorization': 'Bearer ...'})
327
+ server = MCPServerSSE('http://localhost:3001/sse', http_client=http_client)
328
+ ```
329
329
  """
330
330
 
331
331
  timeout: float = 5
@@ -342,10 +342,11 @@ class MCPServerHTTP(MCPServer):
342
342
  If no new messages are received within this time, the connection will be considered stale
343
343
  and may be closed. Defaults to 5 minutes (300 seconds).
344
344
  """
345
+
345
346
  log_level: LoggingLevel | None = None
346
347
  """The log level to set when connecting to the server, if any.
347
348
 
348
- See <https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging#logging> for more details.
349
+ See <https://modelcontextprotocol.io/introduction#logging> for more details.
349
350
 
350
351
  If `None`, no log level will be set.
351
352
  """
@@ -358,28 +359,161 @@ class MCPServerHTTP(MCPServer):
358
359
  For example, if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
359
360
  """
360
361
 
362
+ @property
363
+ @abstractmethod
364
+ def _transport_client(
365
+ self,
366
+ ) -> Callable[
367
+ ...,
368
+ AbstractAsyncContextManager[
369
+ tuple[
370
+ MemoryObjectReceiveStream[SessionMessage | Exception],
371
+ MemoryObjectSendStream[SessionMessage],
372
+ GetSessionIdCallback,
373
+ ],
374
+ ]
375
+ | AbstractAsyncContextManager[
376
+ tuple[
377
+ MemoryObjectReceiveStream[SessionMessage | Exception],
378
+ MemoryObjectSendStream[SessionMessage],
379
+ ]
380
+ ],
381
+ ]: ...
382
+
361
383
  @asynccontextmanager
362
384
  async def client_streams(
363
385
  self,
364
386
  ) -> AsyncIterator[
365
- tuple[
366
- MemoryObjectReceiveStream[JSONRPCMessage | Exception],
367
- MemoryObjectSendStream[JSONRPCMessage],
368
- ]
387
+ tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]]
369
388
  ]: # pragma: no cover
370
- async with sse_client(
389
+ if self.http_client and self.headers:
390
+ raise ValueError('`http_client` is mutually exclusive with `headers`.')
391
+
392
+ transport_client_partial = functools.partial(
393
+ self._transport_client,
371
394
  url=self.url,
372
- headers=self.headers,
373
395
  timeout=self.timeout,
374
396
  sse_read_timeout=self.sse_read_timeout,
375
- ) as (read_stream, write_stream):
376
- yield read_stream, write_stream
397
+ )
398
+
399
+ if self.http_client is not None:
400
+
401
+ def httpx_client_factory(
402
+ headers: dict[str, str] | None = None,
403
+ timeout: httpx.Timeout | None = None,
404
+ auth: httpx.Auth | None = None,
405
+ ) -> httpx.AsyncClient:
406
+ assert self.http_client is not None
407
+ return self.http_client
408
+
409
+ async with transport_client_partial(httpx_client_factory=httpx_client_factory) as (
410
+ read_stream,
411
+ write_stream,
412
+ *_,
413
+ ):
414
+ yield read_stream, write_stream
415
+ else:
416
+ async with transport_client_partial(headers=self.headers) as (read_stream, write_stream, *_):
417
+ yield read_stream, write_stream
377
418
 
378
419
  def _get_log_level(self) -> LoggingLevel | None:
379
420
  return self.log_level
380
421
 
381
422
  def __repr__(self) -> str: # pragma: no cover
382
- return f'MCPServerHTTP(url={self.url!r}, tool_prefix={self.tool_prefix!r})'
423
+ return f'{self.__class__.__name__}(url={self.url!r}, tool_prefix={self.tool_prefix!r})'
383
424
 
384
425
  def _get_client_initialize_timeout(self) -> float: # pragma: no cover
385
426
  return self.timeout
427
+
428
+
429
+ @dataclass
430
+ class MCPServerSSE(_MCPServerHTTP):
431
+ """An MCP server that connects over streamable HTTP connections.
432
+
433
+ This class implements the SSE transport from the MCP specification.
434
+ See <https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse> for more information.
435
+
436
+ !!! note
437
+ Using this class as an async context manager will create a new pool of HTTP connections to connect
438
+ to a server which should already be running.
439
+
440
+ Example:
441
+ ```python {py="3.10"}
442
+ from pydantic_ai import Agent
443
+ from pydantic_ai.mcp import MCPServerSSE
444
+
445
+ server = MCPServerSSE('http://localhost:3001/sse') # (1)!
446
+ agent = Agent('openai:gpt-4o', mcp_servers=[server])
447
+
448
+ async def main():
449
+ async with agent.run_mcp_servers(): # (2)!
450
+ ...
451
+ ```
452
+
453
+ 1. E.g. you might be connecting to a server run with [`mcp-run-python`](../mcp/run-python.md).
454
+ 2. This will connect to a server running on `localhost:3001`.
455
+ """
456
+
457
+ @property
458
+ def _transport_client(self):
459
+ return sse_client # pragma: no cover
460
+
461
+
462
+ @deprecated('The `MCPServerHTTP` class is deprecated, use `MCPServerSSE` instead.')
463
+ @dataclass
464
+ class MCPServerHTTP(MCPServerSSE):
465
+ """An MCP server that connects over HTTP using the old SSE transport.
466
+
467
+ This class implements the SSE transport from the MCP specification.
468
+ See <https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse> for more information.
469
+
470
+ !!! note
471
+ Using this class as an async context manager will create a new pool of HTTP connections to connect
472
+ to a server which should already be running.
473
+
474
+ Example:
475
+ ```python {py="3.10" test="skip"}
476
+ from pydantic_ai import Agent
477
+ from pydantic_ai.mcp import MCPServerHTTP
478
+
479
+ server = MCPServerHTTP('http://localhost:3001/sse') # (1)!
480
+ agent = Agent('openai:gpt-4o', mcp_servers=[server])
481
+
482
+ async def main():
483
+ async with agent.run_mcp_servers(): # (2)!
484
+ ...
485
+ ```
486
+
487
+ 1. E.g. you might be connecting to a server run with [`mcp-run-python`](../mcp/run-python.md).
488
+ 2. This will connect to a server running on `localhost:3001`.
489
+ """
490
+
491
+
492
+ @dataclass
493
+ class MCPServerStreamableHTTP(_MCPServerHTTP):
494
+ """An MCP server that connects over HTTP using the Streamable HTTP transport.
495
+
496
+ This class implements the Streamable HTTP transport from the MCP specification.
497
+ See <https://modelcontextprotocol.io/introduction#streamable-http> for more information.
498
+
499
+ !!! note
500
+ Using this class as an async context manager will create a new pool of HTTP connections to connect
501
+ to a server which should already be running.
502
+
503
+ Example:
504
+ ```python {py="3.10"}
505
+ from pydantic_ai import Agent
506
+ from pydantic_ai.mcp import MCPServerStreamableHTTP
507
+
508
+ server = MCPServerStreamableHTTP('http://localhost:8000/mcp') # (1)!
509
+ agent = Agent('openai:gpt-4o', mcp_servers=[server])
510
+
511
+ async def main():
512
+ async with agent.run_mcp_servers(): # (2)!
513
+ ...
514
+ ```
515
+ """
516
+
517
+ @property
518
+ def _transport_client(self):
519
+ return streamablehttp_client # pragma: no cover
@@ -2,6 +2,7 @@ from __future__ import annotations as _annotations
2
2
 
3
3
  import base64
4
4
  import uuid
5
+ from abc import ABC, abstractmethod
5
6
  from collections.abc import Sequence
6
7
  from dataclasses import dataclass, field, replace
7
8
  from datetime import datetime
@@ -80,8 +81,35 @@ class SystemPromptPart:
80
81
 
81
82
 
82
83
  @dataclass(repr=False)
83
- class VideoUrl:
84
- """A URL to an video."""
84
+ class FileUrl(ABC):
85
+ """Abstract base class for any URL-based file."""
86
+
87
+ url: str
88
+ """The URL of the file."""
89
+
90
+ force_download: bool = False
91
+ """If the model supports it:
92
+
93
+ * If True, the file is downloaded and the data is sent to the model as bytes.
94
+ * If False, the URL is sent directly to the model and no download is performed.
95
+ """
96
+
97
+ @property
98
+ @abstractmethod
99
+ def media_type(self) -> str:
100
+ """Return the media type of the file, based on the url."""
101
+
102
+ @property
103
+ @abstractmethod
104
+ def format(self) -> str:
105
+ """The file format."""
106
+
107
+ __repr__ = _utils.dataclasses_no_defaults_repr
108
+
109
+
110
+ @dataclass(repr=False)
111
+ class VideoUrl(FileUrl):
112
+ """A URL to a video."""
85
113
 
86
114
  url: str
87
115
  """The URL of the video."""
@@ -108,9 +136,19 @@ class VideoUrl:
108
136
  return 'video/x-ms-wmv'
109
137
  elif self.url.endswith('.three_gp'):
110
138
  return 'video/3gpp'
139
+ # Assume that YouTube videos are mp4 because there would be no extension
140
+ # to infer from. This should not be a problem, as Gemini disregards media
141
+ # type for YouTube URLs.
142
+ elif self.is_youtube:
143
+ return 'video/mp4'
111
144
  else:
112
145
  raise ValueError(f'Unknown video file extension: {self.url}')
113
146
 
147
+ @property
148
+ def is_youtube(self) -> bool:
149
+ """True if the URL has a YouTube domain."""
150
+ return self.url.startswith(('https://youtu.be/', 'https://youtube.com/', 'https://www.youtube.com/'))
151
+
114
152
  @property
115
153
  def format(self) -> VideoFormat:
116
154
  """The file format of the video.
@@ -119,11 +157,9 @@ class VideoUrl:
119
157
  """
120
158
  return _video_format_lookup[self.media_type]
121
159
 
122
- __repr__ = _utils.dataclasses_no_defaults_repr
123
-
124
160
 
125
161
  @dataclass(repr=False)
126
- class AudioUrl:
162
+ class AudioUrl(FileUrl):
127
163
  """A URL to an audio file."""
128
164
 
129
165
  url: str
@@ -147,11 +183,9 @@ class AudioUrl:
147
183
  """The file format of the audio file."""
148
184
  return _audio_format_lookup[self.media_type]
149
185
 
150
- __repr__ = _utils.dataclasses_no_defaults_repr
151
-
152
186
 
153
187
  @dataclass(repr=False)
154
- class ImageUrl:
188
+ class ImageUrl(FileUrl):
155
189
  """A URL to an image."""
156
190
 
157
191
  url: str
@@ -182,11 +216,9 @@ class ImageUrl:
182
216
  """
183
217
  return _image_format_lookup[self.media_type]
184
218
 
185
- __repr__ = _utils.dataclasses_no_defaults_repr
186
-
187
219
 
188
220
  @dataclass(repr=False)
189
- class DocumentUrl:
221
+ class DocumentUrl(FileUrl):
190
222
  """The URL of the document."""
191
223
 
192
224
  url: str
@@ -215,8 +247,6 @@ class DocumentUrl:
215
247
  except KeyError as e:
216
248
  raise ValueError(f'Unknown document media type: {media_type}') from e
217
249
 
218
- __repr__ = _utils.dataclasses_no_defaults_repr
219
-
220
250
 
221
251
  @dataclass(repr=False)
222
252
  class BinaryContent:
@@ -6,21 +6,23 @@ specific LLM being used.
6
6
 
7
7
  from __future__ import annotations as _annotations
8
8
 
9
+ import base64
9
10
  from abc import ABC, abstractmethod
10
11
  from collections.abc import AsyncIterator, Iterator
11
12
  from contextlib import asynccontextmanager, contextmanager
12
13
  from dataclasses import dataclass, field, replace
13
14
  from datetime import datetime
14
15
  from functools import cache, cached_property
16
+ from typing import Generic, TypeVar, overload
15
17
 
16
18
  import httpx
17
- from typing_extensions import Literal, TypeAliasType
19
+ from typing_extensions import Literal, TypeAliasType, TypedDict
18
20
 
19
21
  from pydantic_ai.profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
20
22
 
21
23
  from .._parts_manager import ModelResponsePartsManager
22
24
  from ..exceptions import UserError
23
- from ..messages import ModelMessage, ModelRequest, ModelResponse, ModelResponseStreamEvent
25
+ from ..messages import FileUrl, ModelMessage, ModelRequest, ModelResponse, ModelResponseStreamEvent, VideoUrl
24
26
  from ..profiles._json_schema import JsonSchemaTransformer
25
27
  from ..settings import ModelSettings
26
28
  from ..tools import ToolDefinition
@@ -553,9 +555,9 @@ def infer_model(model: Model | KnownModelName | str) -> Model:
553
555
 
554
556
  return OpenAIModel(model_name, provider=provider)
555
557
  elif provider in ('google-gla', 'google-vertex'):
556
- from .gemini import GeminiModel
558
+ from .google import GoogleModel
557
559
 
558
- return GeminiModel(model_name, provider=provider)
560
+ return GoogleModel(model_name, provider=provider)
559
561
  elif provider == 'groq':
560
562
  from .groq import GroqModel
561
563
 
@@ -611,6 +613,91 @@ def _cached_async_http_transport() -> httpx.AsyncHTTPTransport:
611
613
  return httpx.AsyncHTTPTransport()
612
614
 
613
615
 
616
+ DataT = TypeVar('DataT', str, bytes)
617
+
618
+
619
+ class DownloadedItem(TypedDict, Generic[DataT]):
620
+ """The downloaded data and its type."""
621
+
622
+ data: DataT
623
+ """The downloaded data."""
624
+
625
+ data_type: str
626
+ """The type of data that was downloaded.
627
+
628
+ Extracted from header "content-type", but defaults to the media type inferred from the file URL if content-type is "application/octet-stream".
629
+ """
630
+
631
+
632
+ @overload
633
+ async def download_item(
634
+ item: FileUrl,
635
+ data_format: Literal['bytes'],
636
+ type_format: Literal['mime', 'extension'] = 'mime',
637
+ ) -> DownloadedItem[bytes]: ...
638
+
639
+
640
+ @overload
641
+ async def download_item(
642
+ item: FileUrl,
643
+ data_format: Literal['base64', 'base64_uri', 'text'],
644
+ type_format: Literal['mime', 'extension'] = 'mime',
645
+ ) -> DownloadedItem[str]: ...
646
+
647
+
648
+ async def download_item(
649
+ item: FileUrl,
650
+ data_format: Literal['bytes', 'base64', 'base64_uri', 'text'] = 'bytes',
651
+ type_format: Literal['mime', 'extension'] = 'mime',
652
+ ) -> DownloadedItem[str] | DownloadedItem[bytes]:
653
+ """Download an item by URL and return the content as a bytes object or a (base64-encoded) string.
654
+
655
+ Args:
656
+ item: The item to download.
657
+ data_format: The format to return the content in:
658
+ - `bytes`: The raw bytes of the content.
659
+ - `base64`: The base64-encoded content.
660
+ - `base64_uri`: The base64-encoded content as a data URI.
661
+ - `text`: The content as a string.
662
+ type_format: The format to return the media type in:
663
+ - `mime`: The media type as a MIME type.
664
+ - `extension`: The media type as an extension.
665
+
666
+ Raises:
667
+ UserError: If the URL points to a YouTube video or its protocol is gs://.
668
+ """
669
+ if item.url.startswith('gs://'):
670
+ raise UserError('Downloading from protocol "gs://" is not supported.')
671
+ elif isinstance(item, VideoUrl) and item.is_youtube:
672
+ raise UserError('Downloading YouTube videos is not supported.')
673
+
674
+ client = cached_async_http_client()
675
+ response = await client.get(item.url, follow_redirects=True)
676
+ response.raise_for_status()
677
+
678
+ if content_type := response.headers.get('content-type'):
679
+ content_type = content_type.split(';')[0]
680
+ if content_type == 'application/octet-stream':
681
+ content_type = None
682
+
683
+ media_type = content_type or item.media_type
684
+
685
+ data_type = media_type
686
+ if type_format == 'extension':
687
+ data_type = data_type.split('/')[1]
688
+
689
+ data = response.content
690
+ if data_format in ('base64', 'base64_uri'):
691
+ data = base64.b64encode(data).decode('utf-8')
692
+ if data_format == 'base64_uri':
693
+ data = f'data:{media_type};base64,{data}'
694
+ return DownloadedItem[str](data=data, data_type=data_type)
695
+ elif data_format == 'text':
696
+ return DownloadedItem[str](data=data.decode('utf-8'), data_type=data_type)
697
+ else:
698
+ return DownloadedItem[bytes](data=data, data_type=data_type)
699
+
700
+
614
701
  @cache
615
702
  def get_user_agent() -> str:
616
703
  """Get the user agent string for the HTTP client."""