pydantic-ai-slim 0.1.2__tar.gz → 0.1.4__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 (53) hide show
  1. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/PKG-INFO +5 -5
  2. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/_agent_graph.py +13 -1
  3. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/_utils.py +1 -10
  4. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/agent.py +16 -17
  5. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/common_tools/duckduckgo.py +0 -2
  6. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/common_tools/tavily.py +0 -2
  7. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/mcp.py +28 -1
  8. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/messages.py +2 -0
  9. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/__init__.py +8 -0
  10. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/anthropic.py +1 -0
  11. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/bedrock.py +7 -8
  12. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/gemini.py +2 -1
  13. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/groq.py +1 -0
  14. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/instrumented.py +6 -0
  15. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/openai.py +27 -20
  16. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/settings.py +10 -0
  17. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pyproject.toml +2 -2
  18. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/.gitignore +0 -0
  19. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/README.md +0 -0
  20. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/__init__.py +0 -0
  21. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/__main__.py +0 -0
  22. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/_cli.py +0 -0
  23. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/_griffe.py +0 -0
  24. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/_output.py +0 -0
  25. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/_parts_manager.py +0 -0
  26. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/_pydantic.py +0 -0
  27. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/_system_prompt.py +0 -0
  28. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/common_tools/__init__.py +0 -0
  29. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/exceptions.py +0 -0
  30. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/format_as_xml.py +0 -0
  31. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/format_prompt.py +0 -0
  32. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/_json_schema.py +0 -0
  33. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/cohere.py +0 -0
  34. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/fallback.py +0 -0
  35. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/function.py +0 -0
  36. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/mistral.py +0 -0
  37. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/test.py +0 -0
  38. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/models/wrapper.py +0 -0
  39. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/__init__.py +0 -0
  40. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/anthropic.py +0 -0
  41. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/azure.py +0 -0
  42. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/bedrock.py +0 -0
  43. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/cohere.py +0 -0
  44. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/deepseek.py +0 -0
  45. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/google_gla.py +0 -0
  46. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/google_vertex.py +0 -0
  47. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/groq.py +0 -0
  48. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/mistral.py +0 -0
  49. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/providers/openai.py +0 -0
  50. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/py.typed +0 -0
  51. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/result.py +0 -0
  52. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/pydantic_ai/tools.py +0 -0
  53. {pydantic_ai_slim-0.1.2 → pydantic_ai_slim-0.1.4}/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.1.2
3
+ Version: 0.1.4
4
4
  Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
5
5
  Author-email: Samuel Colvin <samuel@pydantic.dev>
6
6
  License-Expression: MIT
@@ -29,7 +29,7 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
29
29
  Requires-Dist: griffe>=1.3.2
30
30
  Requires-Dist: httpx>=0.27
31
31
  Requires-Dist: opentelemetry-api>=1.28.0
32
- Requires-Dist: pydantic-graph==0.1.2
32
+ Requires-Dist: pydantic-graph==0.1.4
33
33
  Requires-Dist: pydantic>=2.10
34
34
  Requires-Dist: typing-inspection>=0.4.0
35
35
  Provides-Extra: anthropic
@@ -45,17 +45,17 @@ Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == '
45
45
  Provides-Extra: duckduckgo
46
46
  Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
47
47
  Provides-Extra: evals
48
- Requires-Dist: pydantic-evals==0.1.2; extra == 'evals'
48
+ Requires-Dist: pydantic-evals==0.1.4; extra == 'evals'
49
49
  Provides-Extra: groq
50
50
  Requires-Dist: groq>=0.15.0; extra == 'groq'
51
51
  Provides-Extra: logfire
52
52
  Requires-Dist: logfire>=3.11.0; extra == 'logfire'
53
53
  Provides-Extra: mcp
54
- Requires-Dist: mcp>=1.5.0; (python_version >= '3.10') and extra == 'mcp'
54
+ Requires-Dist: mcp>=1.6.0; (python_version >= '3.10') and extra == 'mcp'
55
55
  Provides-Extra: mistral
56
56
  Requires-Dist: mistralai>=1.2.5; extra == 'mistral'
57
57
  Provides-Extra: openai
58
- Requires-Dist: openai>=1.74.0; extra == 'openai'
58
+ Requires-Dist: openai>=1.75.0; extra == 'openai'
59
59
  Provides-Extra: tavily
60
60
  Requires-Dist: tavily-python>=0.5.0; extra == 'tavily'
61
61
  Provides-Extra: vertexai
@@ -427,6 +427,18 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
427
427
  # No events are emitted during the handling of text responses, so we don't need to yield anything
428
428
  self._next_node = await self._handle_text_response(ctx, texts)
429
429
  else:
430
+ # we've got an empty response, this sometimes happens with anthropic (and perhaps other models)
431
+ # when the model has already returned text along side tool calls
432
+ # in this scenario, if text responses are allowed, we return text from the most recent model
433
+ # response, if any
434
+ if allow_text_output(ctx.deps.output_schema):
435
+ for message in reversed(ctx.state.message_history):
436
+ if isinstance(message, _messages.ModelResponse):
437
+ last_texts = [p.content for p in message.parts if isinstance(p, _messages.TextPart)]
438
+ if last_texts:
439
+ self._next_node = await self._handle_text_response(ctx, last_texts)
440
+ return
441
+
430
442
  raise exceptions.UnexpectedModelBehavior('Received empty model response')
431
443
 
432
444
  self._events_iterator = _run_stream()
@@ -530,6 +542,7 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
530
542
 
531
543
  text = '\n\n'.join(texts)
532
544
  if allow_text_output(output_schema):
545
+ # The following cast is safe because we know `str` is an allowed result type
533
546
  result_data_input = cast(NodeRunEndT, text)
534
547
  try:
535
548
  result_data = await _validate_output(result_data_input, ctx, None)
@@ -537,7 +550,6 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
537
550
  ctx.state.increment_retries(ctx.deps.max_result_retries)
538
551
  return ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry]))
539
552
  else:
540
- # The following cast is safe because we know `str` is an allowed result type
541
553
  return self._handle_final_result(ctx, result.FinalResult(result_data, None, None), [])
542
554
  else:
543
555
  ctx.state.increment_retries(ctx.deps.max_result_retries)
@@ -291,13 +291,4 @@ class PeekableAsyncStream(Generic[T]):
291
291
 
292
292
 
293
293
  def get_traceparent(x: AgentRun | AgentRunResult | GraphRun | GraphRunResult) -> str:
294
- import logfire
295
- import logfire_api
296
- from logfire.experimental.annotations import get_traceparent
297
-
298
- span: AbstractSpan | None = x._span(required=False) # type: ignore[reportPrivateUsage]
299
- if not span: # pragma: no cover
300
- return ''
301
- if isinstance(span, logfire_api.LogfireSpan): # pragma: no cover
302
- assert isinstance(span, logfire.LogfireSpan)
303
- return get_traceparent(span)
294
+ return x._traceparent(required=False) or '' # type: ignore[reportPrivateUsage]
@@ -27,7 +27,6 @@ from . import (
27
27
  result,
28
28
  usage as _usage,
29
29
  )
30
- from ._utils import AbstractSpan
31
30
  from .models.instrumented import InstrumentationSettings, InstrumentedModel
32
31
  from .result import FinalResult, OutputDataT, StreamedRunResult, ToolOutput
33
32
  from .settings import ModelSettings, merge_model_settings
@@ -659,7 +658,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
659
658
  start_node,
660
659
  state=state,
661
660
  deps=graph_deps,
662
- span=use_span(run_span, end_on_exit=True),
661
+ span=use_span(run_span, end_on_exit=True) if run_span.is_recording() else None,
663
662
  infer_name=False,
664
663
  ) as graph_run:
665
664
  yield AgentRun(graph_run)
@@ -1683,14 +1682,14 @@ class AgentRun(Generic[AgentDepsT, OutputDataT]):
1683
1682
  ]
1684
1683
 
1685
1684
  @overload
1686
- def _span(self, *, required: Literal[False]) -> AbstractSpan | None: ...
1685
+ def _traceparent(self, *, required: Literal[False]) -> str | None: ...
1687
1686
  @overload
1688
- def _span(self) -> AbstractSpan: ...
1689
- def _span(self, *, required: bool = True) -> AbstractSpan | None:
1690
- span = self._graph_run._span(required=False) # type: ignore[reportPrivateUsage]
1691
- if span is None and required: # pragma: no cover
1692
- raise AttributeError('Span is not available for this agent run')
1693
- return span
1687
+ def _traceparent(self) -> str: ...
1688
+ def _traceparent(self, *, required: bool = True) -> str | None:
1689
+ traceparent = self._graph_run._traceparent(required=False) # type: ignore[reportPrivateUsage]
1690
+ if traceparent is None and required: # pragma: no cover
1691
+ raise AttributeError('No span was created for this agent run')
1692
+ return traceparent
1694
1693
 
1695
1694
  @property
1696
1695
  def ctx(self) -> GraphRunContext[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any]]:
@@ -1729,7 +1728,7 @@ class AgentRun(Generic[AgentDepsT, OutputDataT]):
1729
1728
  graph_run_result.output.tool_name,
1730
1729
  graph_run_result.state,
1731
1730
  self._graph_run.deps.new_message_index,
1732
- self._graph_run._span(required=False), # type: ignore[reportPrivateUsage]
1731
+ self._traceparent(required=False),
1733
1732
  )
1734
1733
 
1735
1734
  def __aiter__(
@@ -1847,16 +1846,16 @@ class AgentRunResult(Generic[OutputDataT]):
1847
1846
  _output_tool_name: str | None = dataclasses.field(repr=False)
1848
1847
  _state: _agent_graph.GraphAgentState = dataclasses.field(repr=False)
1849
1848
  _new_message_index: int = dataclasses.field(repr=False)
1850
- _span_value: AbstractSpan | None = dataclasses.field(repr=False)
1849
+ _traceparent_value: str | None = dataclasses.field(repr=False)
1851
1850
 
1852
1851
  @overload
1853
- def _span(self, *, required: Literal[False]) -> AbstractSpan | None: ...
1852
+ def _traceparent(self, *, required: Literal[False]) -> str | None: ...
1854
1853
  @overload
1855
- def _span(self) -> AbstractSpan: ...
1856
- def _span(self, *, required: bool = True) -> AbstractSpan | None:
1857
- if self._span_value is None and required: # pragma: no cover
1858
- raise AttributeError('Span is not available for this agent run')
1859
- return self._span_value
1854
+ def _traceparent(self) -> str: ...
1855
+ def _traceparent(self, *, required: bool = True) -> str | None:
1856
+ if self._traceparent_value is None and required: # pragma: no cover
1857
+ raise AttributeError('No span was created for this agent run')
1858
+ return self._traceparent_value
1860
1859
 
1861
1860
  @property
1862
1861
  @deprecated('`result.data` is deprecated, use `result.output` instead.')
@@ -54,8 +54,6 @@ class DuckDuckGoSearchTool:
54
54
  """
55
55
  search = functools.partial(self.client.text, max_results=self.max_results)
56
56
  results = await anyio.to_thread.run_sync(search, query)
57
- if len(results) == 0:
58
- raise RuntimeError('No search results found.')
59
57
  return duckduckgo_ta.validate_python(results)
60
58
 
61
59
 
@@ -63,8 +63,6 @@ class TavilySearchTool:
63
63
  The search results.
64
64
  """
65
65
  results = await self.client.search(query, search_depth=search_deep, topic=topic, time_range=time_range) # type: ignore[reportUnknownMemberType]
66
- if not results['results']:
67
- raise RuntimeError('No search results found.')
68
66
  return tavily_search_ta.validate_python(results['results'])
69
67
 
70
68
 
@@ -9,7 +9,7 @@ from types import TracebackType
9
9
  from typing import Any
10
10
 
11
11
  from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
12
- from mcp.types import JSONRPCMessage
12
+ from mcp.types import JSONRPCMessage, LoggingLevel
13
13
  from typing_extensions import Self
14
14
 
15
15
  from pydantic_ai.tools import ToolDefinition
@@ -52,6 +52,11 @@ class MCPServer(ABC):
52
52
  raise NotImplementedError('MCP Server subclasses must implement this method.')
53
53
  yield
54
54
 
55
+ @abstractmethod
56
+ def _get_log_level(self) -> LoggingLevel | None:
57
+ """Get the log level for the MCP server."""
58
+ raise NotImplementedError('MCP Server subclasses must implement this method.')
59
+
55
60
  async def list_tools(self) -> list[ToolDefinition]:
56
61
  """Retrieve tools that are currently active on the server.
57
62
 
@@ -89,6 +94,8 @@ class MCPServer(ABC):
89
94
  self._client = await self._exit_stack.enter_async_context(client)
90
95
 
91
96
  await self._client.initialize()
97
+ if log_level := self._get_log_level():
98
+ await self._client.set_logging_level(log_level)
92
99
  self.is_running = True
93
100
  return self
94
101
 
@@ -150,6 +157,13 @@ class MCPServerStdio(MCPServer):
150
157
  By default the subprocess will not inherit any environment variables from the parent process.
151
158
  If you want to inherit the environment variables from the parent process, use `env=os.environ`.
152
159
  """
160
+ log_level: LoggingLevel | None = None
161
+ """The log level to set when connecting to the server, if any.
162
+
163
+ See <https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging#logging> for more details.
164
+
165
+ If `None`, no log level will be set.
166
+ """
153
167
 
154
168
  cwd: str | Path | None = None
155
169
  """The working directory to use when spawning the process."""
@@ -164,6 +178,9 @@ class MCPServerStdio(MCPServer):
164
178
  async with stdio_client(server=server) as (read_stream, write_stream):
165
179
  yield read_stream, write_stream
166
180
 
181
+ def _get_log_level(self) -> LoggingLevel | None:
182
+ return self.log_level
183
+
167
184
 
168
185
  @dataclass
169
186
  class MCPServerHTTP(MCPServer):
@@ -223,6 +240,13 @@ class MCPServerHTTP(MCPServer):
223
240
  If no new messages are received within this time, the connection will be considered stale
224
241
  and may be closed. Defaults to 5 minutes (300 seconds).
225
242
  """
243
+ log_level: LoggingLevel | None = None
244
+ """The log level to set when connecting to the server, if any.
245
+
246
+ See <https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging#logging> for more details.
247
+
248
+ If `None`, no log level will be set.
249
+ """
226
250
 
227
251
  @asynccontextmanager
228
252
  async def client_streams(
@@ -234,3 +258,6 @@ class MCPServerHTTP(MCPServer):
234
258
  url=self.url, headers=self.headers, timeout=self.timeout, sse_read_timeout=self.sse_read_timeout
235
259
  ) as (read_stream, write_stream):
236
260
  yield read_stream, write_stream
261
+
262
+ def _get_log_level(self) -> LoggingLevel | None:
263
+ return self.log_level
@@ -508,6 +508,8 @@ class ToolCallPart:
508
508
  """
509
509
  if isinstance(self.args, dict):
510
510
  return self.args
511
+ if isinstance(self.args, str) and not self.args:
512
+ return {}
511
513
  args = pydantic_core.from_json(self.args)
512
514
  assert isinstance(args, dict), 'args should be a dict'
513
515
  return cast(dict[str, Any], args)
@@ -106,6 +106,7 @@ KnownModelName = TypeAliasType(
106
106
  'google-gla:gemini-2.0-flash',
107
107
  'google-gla:gemini-2.0-flash-lite-preview-02-05',
108
108
  'google-gla:gemini-2.0-pro-exp-02-05',
109
+ 'google-gla:gemini-2.5-flash-preview-04-17',
109
110
  'google-gla:gemini-2.5-pro-exp-03-25',
110
111
  'google-gla:gemini-2.5-pro-preview-03-25',
111
112
  'google-vertex:gemini-1.0-pro',
@@ -118,6 +119,7 @@ KnownModelName = TypeAliasType(
118
119
  'google-vertex:gemini-2.0-flash',
119
120
  'google-vertex:gemini-2.0-flash-lite-preview-02-05',
120
121
  'google-vertex:gemini-2.0-pro-exp-02-05',
122
+ 'google-vertex:gemini-2.5-flash-preview-04-17',
121
123
  'google-vertex:gemini-2.5-pro-exp-03-25',
122
124
  'google-vertex:gemini-2.5-pro-preview-03-25',
123
125
  'gpt-3.5-turbo',
@@ -192,6 +194,8 @@ KnownModelName = TypeAliasType(
192
194
  'o1-mini-2024-09-12',
193
195
  'o1-preview',
194
196
  'o1-preview-2024-09-12',
197
+ 'o3',
198
+ 'o3-2025-04-16',
195
199
  'o3-mini',
196
200
  'o3-mini-2025-01-31',
197
201
  'openai:chatgpt-4o-latest',
@@ -241,8 +245,12 @@ KnownModelName = TypeAliasType(
241
245
  'openai:o1-mini-2024-09-12',
242
246
  'openai:o1-preview',
243
247
  'openai:o1-preview-2024-09-12',
248
+ 'openai:o3',
249
+ 'openai:o3-2025-04-16',
244
250
  'openai:o3-mini',
245
251
  'openai:o3-mini-2025-01-31',
252
+ 'openai:o4-mini',
253
+ 'openai:o4-mini-2025-04-16',
246
254
  'test',
247
255
  ],
248
256
  )
@@ -239,6 +239,7 @@ class AnthropicModel(Model):
239
239
  timeout=model_settings.get('timeout', NOT_GIVEN),
240
240
  metadata=model_settings.get('anthropic_metadata', NOT_GIVEN),
241
241
  extra_headers={'User-Agent': get_user_agent()},
242
+ extra_body=model_settings.get('extra_body'),
242
243
  )
243
244
  except APIStatusError as e:
244
245
  if (status_code := e.status_code) >= 400:
@@ -2,10 +2,11 @@ from __future__ import annotations
2
2
 
3
3
  import functools
4
4
  import typing
5
- from collections.abc import AsyncIterator, Iterable, Mapping
5
+ from collections.abc import AsyncIterator, Iterable, Iterator, Mapping
6
6
  from contextlib import asynccontextmanager
7
7
  from dataclasses import dataclass, field
8
8
  from datetime import datetime
9
+ from itertools import count
9
10
  from typing import TYPE_CHECKING, Any, Generic, Literal, Union, cast, overload
10
11
 
11
12
  import anyio
@@ -369,13 +370,14 @@ class BedrockConverseModel(Model):
369
370
  """Just maps a `pydantic_ai.Message` to the Bedrock `MessageUnionTypeDef`."""
370
371
  system_prompt: list[SystemContentBlockTypeDef] = []
371
372
  bedrock_messages: list[MessageUnionTypeDef] = []
373
+ document_count: Iterator[int] = count(1)
372
374
  for m in messages:
373
375
  if isinstance(m, ModelRequest):
374
376
  for part in m.parts:
375
377
  if isinstance(part, SystemPromptPart):
376
378
  system_prompt.append({'text': part.content})
377
379
  elif isinstance(part, UserPromptPart):
378
- bedrock_messages.extend(await self._map_user_prompt(part))
380
+ bedrock_messages.extend(await self._map_user_prompt(part, document_count))
379
381
  elif isinstance(part, ToolReturnPart):
380
382
  assert part.tool_call_id is not None
381
383
  bedrock_messages.append(
@@ -430,20 +432,18 @@ class BedrockConverseModel(Model):
430
432
  return system_prompt, bedrock_messages
431
433
 
432
434
  @staticmethod
433
- async def _map_user_prompt(part: UserPromptPart) -> list[MessageUnionTypeDef]:
435
+ async def _map_user_prompt(part: UserPromptPart, document_count: Iterator[int]) -> list[MessageUnionTypeDef]:
434
436
  content: list[ContentBlockUnionTypeDef] = []
435
437
  if isinstance(part.content, str):
436
438
  content.append({'text': part.content})
437
439
  else:
438
- document_count = 0
439
440
  for item in part.content:
440
441
  if isinstance(item, str):
441
442
  content.append({'text': item})
442
443
  elif isinstance(item, BinaryContent):
443
444
  format = item.format
444
445
  if item.is_document:
445
- document_count += 1
446
- name = f'Document {document_count}'
446
+ name = f'Document {next(document_count)}'
447
447
  assert format in ('pdf', 'txt', 'csv', 'doc', 'docx', 'xls', 'xlsx', 'html', 'md')
448
448
  content.append({'document': {'name': name, 'format': format, 'source': {'bytes': item.data}}})
449
449
  elif item.is_image:
@@ -464,8 +464,7 @@ class BedrockConverseModel(Model):
464
464
  content.append({'image': image})
465
465
 
466
466
  elif item.kind == 'document-url':
467
- document_count += 1
468
- name = f'Document {document_count}'
467
+ name = f'Document {next(document_count)}'
469
468
  data = response.content
470
469
  content.append({'document': {'name': name, 'format': item.format, 'source': {'bytes': data}}})
471
470
 
@@ -58,6 +58,7 @@ LatestGeminiModelNames = Literal[
58
58
  'gemini-2.0-flash',
59
59
  'gemini-2.0-flash-lite-preview-02-05',
60
60
  'gemini-2.0-pro-exp-02-05',
61
+ 'gemini-2.5-flash-preview-04-17',
61
62
  'gemini-2.5-pro-exp-03-25',
62
63
  'gemini-2.5-pro-preview-03-25',
63
64
  ]
@@ -640,7 +641,7 @@ class _GeminiTextContent(TypedDict):
640
641
 
641
642
 
642
643
  class _GeminiTools(TypedDict):
643
- function_declarations: list[Annotated[_GeminiFunction, pydantic.Field(alias='functionDeclarations')]]
644
+ function_declarations: Annotated[list[_GeminiFunction], pydantic.Field(alias='functionDeclarations')]
644
645
 
645
646
 
646
647
  class _GeminiFunction(TypedDict):
@@ -218,6 +218,7 @@ class GroqModel(Model):
218
218
  frequency_penalty=model_settings.get('frequency_penalty', NOT_GIVEN),
219
219
  logit_bias=model_settings.get('logit_bias', NOT_GIVEN),
220
220
  extra_headers={'User-Agent': get_user_agent()},
221
+ extra_body=model_settings.get('extra_body'),
221
222
  )
222
223
  except APIStatusError as e:
223
224
  if (status_code := e.status_code) >= 400:
@@ -261,9 +261,11 @@ class InstrumentedModel(WrapperModel):
261
261
  @staticmethod
262
262
  def messages_to_otel_events(messages: list[ModelMessage]) -> list[Event]:
263
263
  events: list[Event] = []
264
+ last_model_request: ModelRequest | None = None
264
265
  for message_index, message in enumerate(messages):
265
266
  message_events: list[Event] = []
266
267
  if isinstance(message, ModelRequest):
268
+ last_model_request = message
267
269
  for part in message.parts:
268
270
  if hasattr(part, 'otel_event'):
269
271
  message_events.append(part.otel_event())
@@ -275,6 +277,10 @@ class InstrumentedModel(WrapperModel):
275
277
  **(event.attributes or {}),
276
278
  }
277
279
  events.extend(message_events)
280
+ if last_model_request and last_model_request.instructions:
281
+ events.insert(
282
+ 0, Event('gen_ai.system.message', body={'content': last_model_request.instructions, 'role': 'system'})
283
+ )
278
284
  for event in events:
279
285
  event.body = InstrumentedModel.serialize_any(event.body)
280
286
  return events
@@ -57,6 +57,7 @@ try:
57
57
  )
58
58
  from openai.types.chat.chat_completion_content_part_image_param import ImageURL
59
59
  from openai.types.chat.chat_completion_content_part_input_audio_param import InputAudio
60
+ from openai.types.chat.chat_completion_content_part_param import File, FileFile
60
61
  from openai.types.responses import ComputerToolParam, FileSearchToolParam, WebSearchToolParam
61
62
  from openai.types.responses.response_input_param import FunctionCallOutput, Message
62
63
  from openai.types.shared import ReasoningEffort
@@ -284,6 +285,7 @@ class OpenAIModel(Model):
284
285
  reasoning_effort=model_settings.get('openai_reasoning_effort', NOT_GIVEN),
285
286
  user=model_settings.get('openai_user', NOT_GIVEN),
286
287
  extra_headers={'User-Agent': get_user_agent()},
288
+ extra_body=model_settings.get('extra_body'),
287
289
  )
288
290
  except APIStatusError as e:
289
291
  if (status_code := e.status_code) >= 400:
@@ -425,6 +427,16 @@ class OpenAIModel(Model):
425
427
  assert item.format in ('wav', 'mp3')
426
428
  audio = InputAudio(data=base64_encoded, format=item.format)
427
429
  content.append(ChatCompletionContentPartInputAudioParam(input_audio=audio, type='input_audio'))
430
+ elif item.is_document:
431
+ content.append(
432
+ File(
433
+ file=FileFile(
434
+ file_data=f'data:{item.media_type};base64,{base64_encoded}',
435
+ filename=f'filename.{item.format}',
436
+ ),
437
+ type='file',
438
+ )
439
+ )
428
440
  else: # pragma: no cover
429
441
  raise RuntimeError(f'Unsupported binary content type: {item.media_type}')
430
442
  elif isinstance(item, AudioUrl): # pragma: no cover
@@ -434,25 +446,18 @@ class OpenAIModel(Model):
434
446
  base64_encoded = base64.b64encode(response.content).decode('utf-8')
435
447
  audio = InputAudio(data=base64_encoded, format=response.headers.get('content-type'))
436
448
  content.append(ChatCompletionContentPartInputAudioParam(input_audio=audio, type='input_audio'))
437
- elif isinstance(item, DocumentUrl): # pragma: no cover
438
- raise NotImplementedError('DocumentUrl is not supported for OpenAI')
439
- # The following implementation should have worked, but it seems we have the following error:
440
- # pydantic_ai.exceptions.ModelHTTPError: status_code: 400, model_name: gpt-4o, body:
441
- # {
442
- # 'message': "Unknown parameter: 'messages[1].content[1].file.data'.",
443
- # 'type': 'invalid_request_error',
444
- # 'param': 'messages[1].content[1].file.data',
445
- # 'code': 'unknown_parameter'
446
- # }
447
- #
448
- # client = cached_async_http_client()
449
- # response = await client.get(item.url)
450
- # response.raise_for_status()
451
- # base64_encoded = base64.b64encode(response.content).decode('utf-8')
452
- # media_type = response.headers.get('content-type').split(';')[0]
453
- # file_data = f'data:{media_type};base64,{base64_encoded}'
454
- # file = File(file={'file_data': file_data, 'file_name': item.url, 'file_id': item.url}, type='file')
455
- # content.append(file)
449
+ elif isinstance(item, DocumentUrl):
450
+ client = cached_async_http_client()
451
+ response = await client.get(item.url)
452
+ response.raise_for_status()
453
+ base64_encoded = base64.b64encode(response.content).decode('utf-8')
454
+ media_type = response.headers.get('content-type').split(';')[0]
455
+ file_data = f'data:{media_type};base64,{base64_encoded}'
456
+ file = File(
457
+ file=FileFile(file_data=file_data, filename=f'filename.{item.format}'),
458
+ type='file',
459
+ )
460
+ content.append(file)
456
461
  elif isinstance(item, VideoUrl): # pragma: no cover
457
462
  raise NotImplementedError('VideoUrl is not supported for OpenAI')
458
463
  else:
@@ -623,6 +628,7 @@ class OpenAIResponsesModel(Model):
623
628
  reasoning=reasoning,
624
629
  user=model_settings.get('openai_user', NOT_GIVEN),
625
630
  extra_headers={'User-Agent': get_user_agent()},
631
+ extra_body=model_settings.get('extra_body'),
626
632
  )
627
633
  except APIStatusError as e:
628
634
  if (status_code := e.status_code) >= 400:
@@ -767,10 +773,11 @@ class OpenAIResponsesModel(Model):
767
773
  response = await client.get(item.url)
768
774
  response.raise_for_status()
769
775
  base64_encoded = base64.b64encode(response.content).decode('utf-8')
776
+ media_type = response.headers.get('content-type').split(';')[0]
770
777
  content.append(
771
778
  responses.ResponseInputFileParam(
772
779
  type='input_file',
773
- file_data=f'data:{item.media_type};base64,{base64_encoded}',
780
+ file_data=f'data:{media_type};base64,{base64_encoded}',
774
781
  filename=f'filename.{item.format}',
775
782
  )
776
783
  )
@@ -141,6 +141,16 @@ class ModelSettings(TypedDict, total=False):
141
141
  * Cohere
142
142
  """
143
143
 
144
+ extra_body: object
145
+ """Extra body to send to the model.
146
+
147
+ Supported by:
148
+
149
+ * OpenAI
150
+ * Anthropic
151
+ * Groq
152
+ """
153
+
144
154
 
145
155
  def merge_model_settings(base: ModelSettings | None, overrides: ModelSettings | None) -> ModelSettings | None:
146
156
  """Merge two sets of model settings, preferring the overrides.
@@ -56,7 +56,7 @@ dependencies = [
56
56
  # WARNING if you add optional groups, please update docs/install.md
57
57
  logfire = ["logfire>=3.11.0"]
58
58
  # Models
59
- openai = ["openai>=1.74.0"]
59
+ openai = ["openai>=1.75.0"]
60
60
  cohere = ["cohere>=5.13.11; platform_system != 'Emscripten'"]
61
61
  vertexai = ["google-auth>=2.36.0", "requests>=2.32.3"]
62
62
  anthropic = ["anthropic>=0.49.0"]
@@ -69,7 +69,7 @@ tavily = ["tavily-python>=0.5.0"]
69
69
  # CLI
70
70
  cli = ["rich>=13", "prompt-toolkit>=3", "argcomplete>=3.5.0"]
71
71
  # MCP
72
- mcp = ["mcp>=1.5.0; python_version >= '3.10'"]
72
+ mcp = ["mcp>=1.6.0; python_version >= '3.10'"]
73
73
  # Evals
74
74
  evals = ["pydantic-evals=={{ version }}"]
75
75