pydantic-ai-slim 0.0.36__tar.gz → 0.0.38__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 (45) hide show
  1. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/PKG-INFO +2 -2
  2. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/__init__.py +2 -1
  3. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/agent.py +18 -4
  4. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/messages.py +115 -12
  5. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/__init__.py +5 -0
  6. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/anthropic.py +33 -2
  7. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/bedrock.py +55 -7
  8. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/cohere.py +5 -0
  9. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/fallback.py +7 -1
  10. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/gemini.py +13 -18
  11. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/groq.py +8 -0
  12. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/instrumented.py +44 -19
  13. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/mistral.py +7 -0
  14. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/openai.py +26 -1
  15. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/providers/google_vertex.py +49 -4
  16. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/usage.py +1 -1
  17. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pyproject.toml +2 -2
  18. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/.gitignore +0 -0
  19. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/README.md +0 -0
  20. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/_agent_graph.py +0 -0
  21. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/_cli.py +0 -0
  22. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/_griffe.py +0 -0
  23. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/_parts_manager.py +0 -0
  24. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/_pydantic.py +0 -0
  25. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/_result.py +0 -0
  26. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/_system_prompt.py +0 -0
  27. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/_utils.py +0 -0
  28. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/common_tools/__init__.py +0 -0
  29. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/common_tools/duckduckgo.py +0 -0
  30. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/common_tools/tavily.py +0 -0
  31. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/exceptions.py +0 -0
  32. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/format_as_xml.py +0 -0
  33. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/function.py +0 -0
  34. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/test.py +0 -0
  35. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/vertexai.py +0 -0
  36. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/models/wrapper.py +0 -0
  37. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/providers/__init__.py +0 -0
  38. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/providers/bedrock.py +0 -0
  39. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/providers/deepseek.py +0 -0
  40. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/providers/google_gla.py +0 -0
  41. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/providers/openai.py +0 -0
  42. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/py.typed +0 -0
  43. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/result.py +0 -0
  44. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/settings.py +0 -0
  45. {pydantic_ai_slim-0.0.36 → pydantic_ai_slim-0.0.38}/pydantic_ai/tools.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pydantic-ai-slim
3
- Version: 0.0.36
3
+ Version: 0.0.38
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.0.36
32
+ Requires-Dist: pydantic-graph==0.0.38
33
33
  Requires-Dist: pydantic>=2.10
34
34
  Requires-Dist: typing-inspection>=0.4.0
35
35
  Provides-Extra: anthropic
@@ -10,7 +10,7 @@ from .exceptions import (
10
10
  UsageLimitExceeded,
11
11
  UserError,
12
12
  )
13
- from .messages import AudioUrl, BinaryContent, ImageUrl
13
+ from .messages import AudioUrl, BinaryContent, DocumentUrl, ImageUrl
14
14
  from .tools import RunContext, Tool
15
15
 
16
16
  __all__ = (
@@ -33,6 +33,7 @@ __all__ = (
33
33
  # messages
34
34
  'ImageUrl',
35
35
  'AudioUrl',
36
+ 'DocumentUrl',
36
37
  'BinaryContent',
37
38
  # tools
38
39
  'Tool',
@@ -922,6 +922,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
922
922
  self,
923
923
  /,
924
924
  *,
925
+ name: str | None = None,
925
926
  retries: int | None = None,
926
927
  prepare: ToolPrepareFunc[AgentDepsT] | None = None,
927
928
  docstring_format: DocstringFormat = 'auto',
@@ -933,6 +934,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
933
934
  func: ToolFuncContext[AgentDepsT, ToolParams] | None = None,
934
935
  /,
935
936
  *,
937
+ name: str | None = None,
936
938
  retries: int | None = None,
937
939
  prepare: ToolPrepareFunc[AgentDepsT] | None = None,
938
940
  docstring_format: DocstringFormat = 'auto',
@@ -969,6 +971,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
969
971
 
970
972
  Args:
971
973
  func: The tool function to register.
974
+ name: The name of the tool, defaults to the function name.
972
975
  retries: The number of retries to allow for this tool, defaults to the agent's default retries,
973
976
  which defaults to 1.
974
977
  prepare: custom method to prepare the tool definition for each step, return `None` to omit this
@@ -984,13 +987,17 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
984
987
  func_: ToolFuncContext[AgentDepsT, ToolParams],
985
988
  ) -> ToolFuncContext[AgentDepsT, ToolParams]:
986
989
  # noinspection PyTypeChecker
987
- self._register_function(func_, True, retries, prepare, docstring_format, require_parameter_descriptions)
990
+ self._register_function(
991
+ func_, True, name, retries, prepare, docstring_format, require_parameter_descriptions
992
+ )
988
993
  return func_
989
994
 
990
995
  return tool_decorator
991
996
  else:
992
997
  # noinspection PyTypeChecker
993
- self._register_function(func, True, retries, prepare, docstring_format, require_parameter_descriptions)
998
+ self._register_function(
999
+ func, True, name, retries, prepare, docstring_format, require_parameter_descriptions
1000
+ )
994
1001
  return func
995
1002
 
996
1003
  @overload
@@ -1001,6 +1008,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
1001
1008
  self,
1002
1009
  /,
1003
1010
  *,
1011
+ name: str | None = None,
1004
1012
  retries: int | None = None,
1005
1013
  prepare: ToolPrepareFunc[AgentDepsT] | None = None,
1006
1014
  docstring_format: DocstringFormat = 'auto',
@@ -1012,6 +1020,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
1012
1020
  func: ToolFuncPlain[ToolParams] | None = None,
1013
1021
  /,
1014
1022
  *,
1023
+ name: str | None = None,
1015
1024
  retries: int | None = None,
1016
1025
  prepare: ToolPrepareFunc[AgentDepsT] | None = None,
1017
1026
  docstring_format: DocstringFormat = 'auto',
@@ -1048,6 +1057,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
1048
1057
 
1049
1058
  Args:
1050
1059
  func: The tool function to register.
1060
+ name: The name of the tool, defaults to the function name.
1051
1061
  retries: The number of retries to allow for this tool, defaults to the agent's default retries,
1052
1062
  which defaults to 1.
1053
1063
  prepare: custom method to prepare the tool definition for each step, return `None` to omit this
@@ -1062,19 +1072,22 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
1062
1072
  def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
1063
1073
  # noinspection PyTypeChecker
1064
1074
  self._register_function(
1065
- func_, False, retries, prepare, docstring_format, require_parameter_descriptions
1075
+ func_, False, name, retries, prepare, docstring_format, require_parameter_descriptions
1066
1076
  )
1067
1077
  return func_
1068
1078
 
1069
1079
  return tool_decorator
1070
1080
  else:
1071
- self._register_function(func, False, retries, prepare, docstring_format, require_parameter_descriptions)
1081
+ self._register_function(
1082
+ func, False, name, retries, prepare, docstring_format, require_parameter_descriptions
1083
+ )
1072
1084
  return func
1073
1085
 
1074
1086
  def _register_function(
1075
1087
  self,
1076
1088
  func: ToolFuncEither[AgentDepsT, ToolParams],
1077
1089
  takes_ctx: bool,
1090
+ name: str | None,
1078
1091
  retries: int | None,
1079
1092
  prepare: ToolPrepareFunc[AgentDepsT] | None,
1080
1093
  docstring_format: DocstringFormat,
@@ -1085,6 +1098,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
1085
1098
  tool = Tool[AgentDepsT](
1086
1099
  func,
1087
1100
  takes_ctx=takes_ctx,
1101
+ name=name,
1088
1102
  max_retries=retries_,
1089
1103
  prepare=prepare,
1090
1104
  docstring_format=docstring_format,
@@ -4,6 +4,7 @@ import uuid
4
4
  from collections.abc import Sequence
5
5
  from dataclasses import dataclass, field, replace
6
6
  from datetime import datetime
7
+ from mimetypes import guess_type
7
8
  from typing import Annotated, Any, Literal, Union, cast, overload
8
9
 
9
10
  import pydantic
@@ -83,9 +84,57 @@ class ImageUrl:
83
84
  else:
84
85
  raise ValueError(f'Unknown image file extension: {self.url}')
85
86
 
87
+ @property
88
+ def format(self) -> ImageFormat:
89
+ """The file format of the image.
90
+
91
+ The choice of supported formats were based on the Bedrock Converse API. Other APIs don't require to use a format.
92
+ """
93
+ return _image_format(self.media_type)
94
+
95
+
96
+ @dataclass
97
+ class DocumentUrl:
98
+ """The URL of the document."""
99
+
100
+ url: str
101
+ """The URL of the document."""
102
+
103
+ kind: Literal['document-url'] = 'document-url'
104
+ """Type identifier, this is available on all parts as a discriminator."""
105
+
106
+ @property
107
+ def media_type(self) -> str:
108
+ """Return the media type of the document, based on the url."""
109
+ type_, _ = guess_type(self.url)
110
+ if type_ is None:
111
+ raise RuntimeError(f'Unknown document file extension: {self.url}')
112
+ return type_
113
+
114
+ @property
115
+ def format(self) -> DocumentFormat:
116
+ """The file format of the document.
117
+
118
+ The choice of supported formats were based on the Bedrock Converse API. Other APIs don't require to use a format.
119
+ """
120
+ return _document_format(self.media_type)
121
+
86
122
 
87
123
  AudioMediaType: TypeAlias = Literal['audio/wav', 'audio/mpeg']
88
124
  ImageMediaType: TypeAlias = Literal['image/jpeg', 'image/png', 'image/gif', 'image/webp']
125
+ DocumentMediaType: TypeAlias = Literal[
126
+ 'application/pdf',
127
+ 'text/plain',
128
+ 'text/csv',
129
+ 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
130
+ 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
131
+ 'text/html',
132
+ 'text/markdown',
133
+ 'application/vnd.ms-excel',
134
+ ]
135
+ AudioFormat: TypeAlias = Literal['wav', 'mp3']
136
+ ImageFormat: TypeAlias = Literal['jpeg', 'png', 'gif', 'webp']
137
+ DocumentFormat: TypeAlias = Literal['csv', 'doc', 'docx', 'html', 'md', 'pdf', 'txt', 'xls', 'xlsx']
89
138
 
90
139
 
91
140
  @dataclass
@@ -95,7 +144,7 @@ class BinaryContent:
95
144
  data: bytes
96
145
  """The binary data."""
97
146
 
98
- media_type: AudioMediaType | ImageMediaType | str
147
+ media_type: AudioMediaType | ImageMediaType | DocumentMediaType | str
99
148
  """The media type of the binary data."""
100
149
 
101
150
  kind: Literal['binary'] = 'binary'
@@ -112,17 +161,69 @@ class BinaryContent:
112
161
  return self.media_type.startswith('image/')
113
162
 
114
163
  @property
115
- def audio_format(self) -> Literal['mp3', 'wav']:
116
- """Return the audio format given the media type."""
117
- if self.media_type == 'audio/mpeg':
118
- return 'mp3'
119
- elif self.media_type == 'audio/wav':
120
- return 'wav'
121
- else:
122
- raise ValueError(f'Unknown audio media type: {self.media_type}')
164
+ def is_document(self) -> bool:
165
+ """Return `True` if the media type is a document type."""
166
+ return self.media_type in {
167
+ 'application/pdf',
168
+ 'text/plain',
169
+ 'text/csv',
170
+ 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
171
+ 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
172
+ 'text/html',
173
+ 'text/markdown',
174
+ 'application/vnd.ms-excel',
175
+ }
123
176
 
124
-
125
- UserContent: TypeAlias = 'str | ImageUrl | AudioUrl | BinaryContent'
177
+ @property
178
+ def format(self) -> str:
179
+ """The file format of the binary content."""
180
+ if self.is_audio:
181
+ if self.media_type == 'audio/mpeg':
182
+ return 'mp3'
183
+ elif self.media_type == 'audio/wav':
184
+ return 'wav'
185
+ elif self.is_image:
186
+ return _image_format(self.media_type)
187
+ elif self.is_document:
188
+ return _document_format(self.media_type)
189
+ raise ValueError(f'Unknown media type: {self.media_type}')
190
+
191
+
192
+ UserContent: TypeAlias = 'str | ImageUrl | AudioUrl | DocumentUrl | BinaryContent'
193
+
194
+
195
+ def _document_format(media_type: str) -> DocumentFormat:
196
+ if media_type == 'application/pdf':
197
+ return 'pdf'
198
+ elif media_type == 'text/plain':
199
+ return 'txt'
200
+ elif media_type == 'text/csv':
201
+ return 'csv'
202
+ elif media_type == 'application/vnd.openxmlformats-officedocument.wordprocessingml.document':
203
+ return 'docx'
204
+ elif media_type == 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet':
205
+ return 'xlsx'
206
+ elif media_type == 'text/html':
207
+ return 'html'
208
+ elif media_type == 'text/markdown':
209
+ return 'md'
210
+ elif media_type == 'application/vnd.ms-excel':
211
+ return 'xls'
212
+ else:
213
+ raise ValueError(f'Unknown document media type: {media_type}')
214
+
215
+
216
+ def _image_format(media_type: str) -> ImageFormat:
217
+ if media_type == 'image/jpeg':
218
+ return 'jpeg'
219
+ elif media_type == 'image/png':
220
+ return 'png'
221
+ elif media_type == 'image/gif':
222
+ return 'gif'
223
+ elif media_type == 'image/webp':
224
+ return 'webp'
225
+ else:
226
+ raise ValueError(f'Unknown image media type: {media_type}')
126
227
 
127
228
 
128
229
  @dataclass
@@ -395,7 +496,9 @@ class ModelResponse:
395
496
  ModelMessage = Annotated[Union[ModelRequest, ModelResponse], pydantic.Discriminator('kind')]
396
497
  """Any message sent to or returned by a model."""
397
498
 
398
- ModelMessagesTypeAdapter = pydantic.TypeAdapter(list[ModelMessage], config=pydantic.ConfigDict(defer_build=True))
499
+ ModelMessagesTypeAdapter = pydantic.TypeAdapter(
500
+ list[ModelMessage], config=pydantic.ConfigDict(defer_build=True, ser_json_bytes='base64')
501
+ )
399
502
  """Pydantic [`TypeAdapter`][pydantic.type_adapter.TypeAdapter] for (de)serializing messages."""
400
503
 
401
504
 
@@ -266,6 +266,11 @@ class Model(ABC):
266
266
  """The system / model provider, ex: openai."""
267
267
  raise NotImplementedError()
268
268
 
269
+ @property
270
+ def base_url(self) -> str | None:
271
+ """The base URL for the provider API, if available."""
272
+ return None
273
+
269
274
 
270
275
  @dataclass
271
276
  class StreamedResponse(ABC):
@@ -9,6 +9,7 @@ from datetime import datetime, timezone
9
9
  from json import JSONDecodeError, loads as json_loads
10
10
  from typing import Any, Literal, Union, cast, overload
11
11
 
12
+ from anthropic.types import DocumentBlockParam
12
13
  from httpx import AsyncClient as AsyncHTTPClient
13
14
  from typing_extensions import assert_never
14
15
 
@@ -16,6 +17,7 @@ from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
16
17
  from .._utils import guard_tool_call_id as _guard_tool_call_id
17
18
  from ..messages import (
18
19
  BinaryContent,
20
+ DocumentUrl,
19
21
  ImageUrl,
20
22
  ModelMessage,
21
23
  ModelRequest,
@@ -42,11 +44,13 @@ from . import (
42
44
  try:
43
45
  from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropic, AsyncStream
44
46
  from anthropic.types import (
47
+ Base64PDFSourceParam,
45
48
  ContentBlock,
46
49
  ImageBlockParam,
47
50
  Message as AnthropicMessage,
48
51
  MessageParam,
49
52
  MetadataParam,
53
+ PlainTextSourceParam,
50
54
  RawContentBlockDeltaEvent,
51
55
  RawContentBlockStartEvent,
52
56
  RawContentBlockStopEvent,
@@ -143,6 +147,10 @@ class AnthropicModel(Model):
143
147
  else:
144
148
  self.client = AsyncAnthropic(api_key=api_key, http_client=cached_async_http_client())
145
149
 
150
+ @property
151
+ def base_url(self) -> str:
152
+ return str(self.client.base_url)
153
+
146
154
  async def request(
147
155
  self,
148
156
  messages: list[ModelMessage],
@@ -284,7 +292,9 @@ class AnthropicModel(Model):
284
292
  anthropic_messages: list[MessageParam] = []
285
293
  for m in messages:
286
294
  if isinstance(m, ModelRequest):
287
- user_content_params: list[ToolResultBlockParam | TextBlockParam | ImageBlockParam] = []
295
+ user_content_params: list[
296
+ ToolResultBlockParam | TextBlockParam | ImageBlockParam | DocumentBlockParam
297
+ ] = []
288
298
  for request_part in m.parts:
289
299
  if isinstance(request_part, SystemPromptPart):
290
300
  system_prompt += request_part.content
@@ -330,7 +340,9 @@ class AnthropicModel(Model):
330
340
  return system_prompt, anthropic_messages
331
341
 
332
342
  @staticmethod
333
- async def _map_user_prompt(part: UserPromptPart) -> AsyncGenerator[ImageBlockParam | TextBlockParam]:
343
+ async def _map_user_prompt(
344
+ part: UserPromptPart,
345
+ ) -> AsyncGenerator[ImageBlockParam | TextBlockParam | DocumentBlockParam]:
334
346
  if isinstance(part.content, str):
335
347
  yield TextBlockParam(text=part.content, type='text')
336
348
  else:
@@ -375,6 +387,25 @@ class AnthropicModel(Model):
375
387
  )
376
388
  else: # pragma: no cover
377
389
  raise RuntimeError(f'Unsupported image type: {mime_type}')
390
+ elif isinstance(item, DocumentUrl):
391
+ response = await cached_async_http_client().get(item.url)
392
+ response.raise_for_status()
393
+ if item.media_type == 'application/pdf':
394
+ yield DocumentBlockParam(
395
+ source=Base64PDFSourceParam(
396
+ data=io.BytesIO(response.content),
397
+ media_type=item.media_type,
398
+ type='base64',
399
+ ),
400
+ type='document',
401
+ )
402
+ elif item.media_type == 'text/plain':
403
+ yield DocumentBlockParam(
404
+ source=PlainTextSourceParam(data=response.text, media_type=item.media_type, type='text'),
405
+ type='document',
406
+ )
407
+ else: # pragma: no cover
408
+ raise RuntimeError(f'Unsupported media type: {item.media_type}')
378
409
  else:
379
410
  raise RuntimeError(f'Unsupported content type: {type(item)}')
380
411
 
@@ -10,10 +10,15 @@ from typing import TYPE_CHECKING, Generic, Literal, Union, cast, overload
10
10
 
11
11
  import anyio
12
12
  import anyio.to_thread
13
+ from mypy_boto3_bedrock_runtime.type_defs import ImageBlockTypeDef
13
14
  from typing_extensions import ParamSpec, assert_never
14
15
 
15
16
  from pydantic_ai import _utils, result
16
17
  from pydantic_ai.messages import (
18
+ AudioUrl,
19
+ BinaryContent,
20
+ DocumentUrl,
21
+ ImageUrl,
17
22
  ModelMessage,
18
23
  ModelRequest,
19
24
  ModelResponse,
@@ -26,7 +31,7 @@ from pydantic_ai.messages import (
26
31
  ToolReturnPart,
27
32
  UserPromptPart,
28
33
  )
29
- from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
34
+ from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse, cached_async_http_client
30
35
  from pydantic_ai.providers import Provider, infer_provider
31
36
  from pydantic_ai.settings import ModelSettings
32
37
  from pydantic_ai.tools import ToolDefinition
@@ -37,6 +42,7 @@ if TYPE_CHECKING:
37
42
  from mypy_boto3_bedrock_runtime import BedrockRuntimeClient
38
43
  from mypy_boto3_bedrock_runtime.type_defs import (
39
44
  ContentBlockOutputTypeDef,
45
+ ContentBlockUnionTypeDef,
40
46
  ConverseResponseTypeDef,
41
47
  ConverseStreamMetadataEventTypeDef,
42
48
  ConverseStreamOutputTypeDef,
@@ -162,6 +168,10 @@ class BedrockConverseModel(Model):
162
168
  }
163
169
  }
164
170
 
171
+ @property
172
+ def base_url(self) -> str:
173
+ return str(self.client.meta.endpoint_url)
174
+
165
175
  async def request(
166
176
  self,
167
177
  messages: list[ModelMessage],
@@ -240,7 +250,7 @@ class BedrockConverseModel(Model):
240
250
  else:
241
251
  tool_choice = {'auto': {}}
242
252
 
243
- system_prompt, bedrock_messages = self._map_message(messages)
253
+ system_prompt, bedrock_messages = await self._map_message(messages)
244
254
  inference_config = self._map_inference_config(model_settings)
245
255
 
246
256
  params = {
@@ -281,7 +291,7 @@ class BedrockConverseModel(Model):
281
291
 
282
292
  return inference_config
283
293
 
284
- def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[MessageUnionTypeDef]]:
294
+ async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[MessageUnionTypeDef]]:
285
295
  """Just maps a `pydantic_ai.Message` to the Bedrock `MessageUnionTypeDef`."""
286
296
  system_prompt: str = ''
287
297
  bedrock_messages: list[MessageUnionTypeDef] = []
@@ -291,10 +301,7 @@ class BedrockConverseModel(Model):
291
301
  if isinstance(part, SystemPromptPart):
292
302
  system_prompt += part.content
293
303
  elif isinstance(part, UserPromptPart):
294
- if isinstance(part.content, str):
295
- bedrock_messages.append({'role': 'user', 'content': [{'text': part.content}]})
296
- else:
297
- raise NotImplementedError('User prompt can only be a string for now.')
304
+ bedrock_messages.extend(await self._map_user_prompt(part))
298
305
  elif isinstance(part, ToolReturnPart):
299
306
  assert part.tool_call_id is not None
300
307
  bedrock_messages.append(
@@ -344,6 +351,47 @@ class BedrockConverseModel(Model):
344
351
  assert_never(m)
345
352
  return system_prompt, bedrock_messages
346
353
 
354
+ @staticmethod
355
+ async def _map_user_prompt(part: UserPromptPart) -> list[MessageUnionTypeDef]:
356
+ content: list[ContentBlockUnionTypeDef] = []
357
+ if isinstance(part.content, str):
358
+ content.append({'text': part.content})
359
+ else:
360
+ document_count = 0
361
+ for item in part.content:
362
+ if isinstance(item, str):
363
+ content.append({'text': item})
364
+ elif isinstance(item, BinaryContent):
365
+ format = item.format
366
+ if item.is_document:
367
+ document_count += 1
368
+ name = f'Document {document_count}'
369
+ assert format in ('pdf', 'txt', 'csv', 'doc', 'docx', 'xls', 'xlsx', 'html', 'md')
370
+ content.append({'document': {'name': name, 'format': format, 'source': {'bytes': item.data}}})
371
+ elif item.is_image:
372
+ assert format in ('jpeg', 'png', 'gif', 'webp')
373
+ content.append({'image': {'format': format, 'source': {'bytes': item.data}}})
374
+ else:
375
+ raise NotImplementedError('Binary content is not supported yet.')
376
+ elif isinstance(item, (ImageUrl, DocumentUrl)):
377
+ response = await cached_async_http_client().get(item.url)
378
+ response.raise_for_status()
379
+ if item.kind == 'image-url':
380
+ format = item.media_type.split('/')[1]
381
+ assert format in ('jpeg', 'png', 'gif', 'webp'), f'Unsupported image format: {format}'
382
+ image: ImageBlockTypeDef = {'format': format, 'source': {'bytes': response.content}}
383
+ content.append({'image': image})
384
+ elif item.kind == 'document-url':
385
+ document_count += 1
386
+ name = f'Document {document_count}'
387
+ data = response.content
388
+ content.append({'document': {'name': name, 'format': item.format, 'source': {'bytes': data}}})
389
+ elif isinstance(item, AudioUrl): # pragma: no cover
390
+ raise NotImplementedError('Audio is not supported yet.')
391
+ else:
392
+ assert_never(item)
393
+ return [{'role': 'user', 'content': content}]
394
+
347
395
  @staticmethod
348
396
  def _map_tool_call(t: ToolCallPart) -> ContentBlockOutputTypeDef:
349
397
  assert t.tool_call_id is not None
@@ -127,6 +127,11 @@ class CohereModel(Model):
127
127
  else:
128
128
  self.client = AsyncClientV2(api_key=api_key, httpx_client=http_client)
129
129
 
130
+ @property
131
+ def base_url(self) -> str:
132
+ client_wrapper = self.client._client_wrapper # type: ignore
133
+ return str(client_wrapper.get_base_url())
134
+
130
135
  async def request(
131
136
  self,
132
137
  messages: list[ModelMessage],
@@ -61,7 +61,9 @@ class FallbackModel(Model):
61
61
 
62
62
  for model in self.models:
63
63
  try:
64
- return await model.request(messages, model_settings, model_request_parameters)
64
+ response, usage = await model.request(messages, model_settings, model_request_parameters)
65
+ response.model_used = model # type: ignore
66
+ return response, usage
65
67
  except Exception as exc:
66
68
  if self._fallback_on(exc):
67
69
  exceptions.append(exc)
@@ -106,6 +108,10 @@ class FallbackModel(Model):
106
108
  """The system / model provider, n/a for fallback models."""
107
109
  return None
108
110
 
111
+ @property
112
+ def base_url(self) -> str | None:
113
+ return self.models[0].base_url
114
+
109
115
 
110
116
  def _default_fallback_condition_factory(exceptions: tuple[type[Exception], ...]) -> Callable[[Exception], bool]:
111
117
  """Create a default fallback condition for the given exceptions."""
@@ -21,6 +21,7 @@ from .. import ModelHTTPError, UnexpectedModelBehavior, UserError, _utils, usage
21
21
  from ..messages import (
22
22
  AudioUrl,
23
23
  BinaryContent,
24
+ DocumentUrl,
24
25
  ImageUrl,
25
26
  ModelMessage,
26
27
  ModelRequest,
@@ -143,6 +144,7 @@ class GeminiModel(Model):
143
144
  else:
144
145
  self._system = provider.name
145
146
  self.client = provider.client
147
+ self._url = str(self.client.base_url)
146
148
  else:
147
149
  if api_key is None:
148
150
  if env_api_key := os.getenv('GEMINI_API_KEY'):
@@ -159,7 +161,7 @@ class GeminiModel(Model):
159
161
  return self._auth
160
162
 
161
163
  @property
162
- def url(self) -> str:
164
+ def base_url(self) -> str:
163
165
  assert self._url is not None, 'URL not initialized'
164
166
  return self._url
165
167
 
@@ -257,7 +259,7 @@ class GeminiModel(Model):
257
259
  'User-Agent': get_user_agent(),
258
260
  }
259
261
  if self._provider is None: # pragma: no cover
260
- url = self.url + ('streamGenerateContent' if streamed else 'generateContent')
262
+ url = self.base_url + ('streamGenerateContent' if streamed else 'generateContent')
261
263
  headers.update(await self.auth.headers())
262
264
  else:
263
265
  url = f'/{self._model_name}:{"streamGenerateContent" if streamed else "generateContent"}'
@@ -361,22 +363,15 @@ class GeminiModel(Model):
361
363
  content.append(
362
364
  _GeminiInlineDataPart(inline_data={'data': base64_encoded, 'mime_type': item.media_type})
363
365
  )
364
- elif isinstance(item, (AudioUrl, ImageUrl)):
365
- try:
366
- content.append(
367
- _GeminiFileDataPart(file_data={'file_uri': item.url, 'mime_type': item.media_type})
368
- )
369
- except ValueError:
370
- # Download the file if can't find the mime type.
371
- client = cached_async_http_client()
372
- response = await client.get(item.url, follow_redirects=True)
373
- response.raise_for_status()
374
- base64_encoded = base64.b64encode(response.content).decode('utf-8')
375
- content.append(
376
- _GeminiInlineDataPart(
377
- inline_data={'data': base64_encoded, 'mime_type': response.headers['Content-Type']}
378
- )
379
- )
366
+ elif isinstance(item, (AudioUrl, ImageUrl, DocumentUrl)):
367
+ client = cached_async_http_client()
368
+ response = await client.get(item.url, follow_redirects=True)
369
+ response.raise_for_status()
370
+ mime_type = response.headers['Content-Type'].split(';')[0]
371
+ inline_data = _GeminiInlineDataPart(
372
+ inline_data={'data': base64.b64encode(response.content).decode('utf-8'), 'mime_type': mime_type}
373
+ )
374
+ content.append(inline_data)
380
375
  else:
381
376
  assert_never(item)
382
377
  return content
@@ -15,6 +15,7 @@ from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
15
15
  from .._utils import guard_tool_call_id as _guard_tool_call_id
16
16
  from ..messages import (
17
17
  BinaryContent,
18
+ DocumentUrl,
18
19
  ImageUrl,
19
20
  ModelMessage,
20
21
  ModelRequest,
@@ -123,6 +124,10 @@ class GroqModel(Model):
123
124
  else:
124
125
  self.client = AsyncGroq(api_key=api_key, http_client=cached_async_http_client())
125
126
 
127
+ @property
128
+ def base_url(self) -> str:
129
+ return str(self.client.base_url)
130
+
126
131
  async def request(
127
132
  self,
128
133
  messages: list[ModelMessage],
@@ -338,8 +343,11 @@ class GroqModel(Model):
338
343
  content.append(chat.ChatCompletionContentPartImageParam(image_url=image_url, type='image_url'))
339
344
  else:
340
345
  raise RuntimeError('Only images are supported for binary content in Groq.')
346
+ elif isinstance(item, DocumentUrl): # pragma: no cover
347
+ raise RuntimeError('DocumentUrl is not supported in Groq.')
341
348
  else: # pragma: no cover
342
349
  raise RuntimeError(f'Unsupported content type: {type(item)}')
350
+
343
351
  return chat.ChatCompletionUserMessageParam(role='user', content=content)
344
352
 
345
353
 
@@ -5,6 +5,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping
5
5
  from contextlib import asynccontextmanager, contextmanager
6
6
  from dataclasses import dataclass, field
7
7
  from typing import Any, Callable, Literal
8
+ from urllib.parse import urlparse
8
9
 
9
10
  from opentelemetry._events import Event, EventLogger, EventLoggerProvider, get_event_logger_provider
10
11
  from opentelemetry.trace import Span, Tracer, TracerProvider, get_tracer_provider
@@ -87,6 +88,10 @@ class InstrumentationSettings:
87
88
  self.event_mode = event_mode
88
89
 
89
90
 
91
+ GEN_AI_SYSTEM_ATTRIBUTE = 'gen_ai.system'
92
+ GEN_AI_REQUEST_MODEL_ATTRIBUTE = 'gen_ai.request.model'
93
+
94
+
90
95
  @dataclass
91
96
  class InstrumentedModel(WrapperModel):
92
97
  """Model which is instrumented with OpenTelemetry."""
@@ -137,19 +142,13 @@ class InstrumentedModel(WrapperModel):
137
142
  model_settings: ModelSettings | None,
138
143
  ) -> Iterator[Callable[[ModelResponse, Usage], None]]:
139
144
  operation = 'chat'
140
- model_name = self.model_name
141
- span_name = f'{operation} {model_name}'
142
- system = getattr(self.wrapped, 'system', '') or self.wrapped.__class__.__name__.removesuffix('Model').lower()
143
- system = {'google-gla': 'gemini', 'google-vertex': 'vertex_ai', 'mistral': 'mistral_ai'}.get(system, system)
145
+ span_name = f'{operation} {self.model_name}'
144
146
  # TODO Missing attributes:
145
- # - server.address: requires a Model.base_url abstract method or similar
146
- # - server.port: to parse from the base_url
147
147
  # - error.type: unclear if we should do something here or just always rely on span exceptions
148
148
  # - gen_ai.request.stop_sequences/top_k: model_settings doesn't include these
149
149
  attributes: dict[str, AttributeValue] = {
150
150
  'gen_ai.operation.name': operation,
151
- 'gen_ai.system': system,
152
- 'gen_ai.request.model': model_name,
151
+ **self.model_attributes(self.wrapped),
153
152
  }
154
153
 
155
154
  if model_settings:
@@ -175,21 +174,26 @@ class InstrumentedModel(WrapperModel):
175
174
  },
176
175
  )
177
176
  )
178
- span.set_attributes(
179
- {
180
- # TODO finish_reason (https://github.com/open-telemetry/semantic-conventions/issues/1277), id
181
- # https://github.com/pydantic/pydantic-ai/issues/886
182
- 'gen_ai.response.model': response.model_name or model_name,
183
- **usage.opentelemetry_attributes(),
177
+ new_attributes: dict[str, AttributeValue] = usage.opentelemetry_attributes() # type: ignore
178
+ if model_used := getattr(response, 'model_used', None):
179
+ # FallbackModel sets model_used on the response so that we can report the attributes
180
+ # of the model that was actually used.
181
+ new_attributes.update(self.model_attributes(model_used))
182
+ attributes.update(new_attributes)
183
+ request_model = attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE]
184
+ new_attributes['gen_ai.response.model'] = response.model_name or request_model
185
+ span.set_attributes(new_attributes)
186
+ span.update_name(f'{operation} {request_model}')
187
+ for event in events:
188
+ event.attributes = {
189
+ GEN_AI_SYSTEM_ATTRIBUTE: attributes[GEN_AI_SYSTEM_ATTRIBUTE],
190
+ **(event.attributes or {}),
184
191
  }
185
- )
186
- self._emit_events(system, span, events)
192
+ self._emit_events(span, events)
187
193
 
188
194
  yield finish
189
195
 
190
- def _emit_events(self, system: str, span: Span, events: list[Event]) -> None:
191
- for event in events:
192
- event.attributes = {'gen_ai.system': system, **(event.attributes or {})}
196
+ def _emit_events(self, span: Span, events: list[Event]) -> None:
193
197
  if self.options.event_mode == 'logs':
194
198
  for event in events:
195
199
  self.options.event_logger.emit(event)
@@ -207,6 +211,27 @@ class InstrumentedModel(WrapperModel):
207
211
  }
208
212
  )
209
213
 
214
+ @staticmethod
215
+ def model_attributes(model: Model):
216
+ system = getattr(model, 'system', '') or model.__class__.__name__.removesuffix('Model').lower()
217
+ system = {'google-gla': 'gemini', 'google-vertex': 'vertex_ai', 'mistral': 'mistral_ai'}.get(system, system)
218
+ attributes: dict[str, AttributeValue] = {
219
+ GEN_AI_SYSTEM_ATTRIBUTE: system,
220
+ GEN_AI_REQUEST_MODEL_ATTRIBUTE: model.model_name,
221
+ }
222
+ if base_url := model.base_url:
223
+ try:
224
+ parsed = urlparse(base_url)
225
+ except Exception: # pragma: no cover
226
+ pass
227
+ else:
228
+ if parsed.hostname:
229
+ attributes['server.address'] = parsed.hostname
230
+ if parsed.port:
231
+ attributes['server.port'] = parsed.port
232
+
233
+ return attributes
234
+
210
235
  @staticmethod
211
236
  def event_to_dict(event: Event) -> dict[str, Any]:
212
237
  if not event.body:
@@ -17,6 +17,7 @@ from .. import ModelHTTPError, UnexpectedModelBehavior, _utils
17
17
  from .._utils import now_utc as _now_utc
18
18
  from ..messages import (
19
19
  BinaryContent,
20
+ DocumentUrl,
20
21
  ImageUrl,
21
22
  ModelMessage,
22
23
  ModelRequest,
@@ -140,6 +141,10 @@ class MistralModel(Model):
140
141
  api_key = os.getenv('MISTRAL_API_KEY') if api_key is None else api_key
141
142
  self.client = Mistral(api_key=api_key, async_client=http_client or cached_async_http_client())
142
143
 
144
+ @property
145
+ def base_url(self) -> str:
146
+ return str(self.client.sdk_configuration.get_server_details()[0])
147
+
143
148
  async def request(
144
149
  self,
145
150
  messages: list[ModelMessage],
@@ -491,6 +496,8 @@ class MistralModel(Model):
491
496
  content.append(MistralImageURLChunk(image_url=image_url, type='image_url'))
492
497
  else:
493
498
  raise RuntimeError('Only image binary content is supported for Mistral.')
499
+ elif isinstance(item, DocumentUrl):
500
+ raise RuntimeError('DocumentUrl is not supported in Mistral.')
494
501
  else: # pragma: no cover
495
502
  raise RuntimeError(f'Unsupported content type: {type(item)}')
496
503
  return MistralUserMessage(content=content)
@@ -18,6 +18,7 @@ from .._utils import guard_tool_call_id as _guard_tool_call_id
18
18
  from ..messages import (
19
19
  AudioUrl,
20
20
  BinaryContent,
21
+ DocumentUrl,
21
22
  ImageUrl,
22
23
  ModelMessage,
23
24
  ModelRequest,
@@ -187,6 +188,10 @@ class OpenAIModel(Model):
187
188
  self.system_prompt_role = system_prompt_role
188
189
  self._system = system
189
190
 
191
+ @property
192
+ def base_url(self) -> str:
193
+ return str(self.client.base_url)
194
+
190
195
  async def request(
191
196
  self,
192
197
  messages: list[ModelMessage],
@@ -414,7 +419,8 @@ class OpenAIModel(Model):
414
419
  image_url = ImageURL(url=f'data:{item.media_type};base64,{base64_encoded}')
415
420
  content.append(ChatCompletionContentPartImageParam(image_url=image_url, type='image_url'))
416
421
  elif item.is_audio:
417
- audio = InputAudio(data=base64_encoded, format=item.audio_format)
422
+ assert item.format in ('wav', 'mp3')
423
+ audio = InputAudio(data=base64_encoded, format=item.format)
418
424
  content.append(ChatCompletionContentPartInputAudioParam(input_audio=audio, type='input_audio'))
419
425
  else: # pragma: no cover
420
426
  raise RuntimeError(f'Unsupported binary content type: {item.media_type}')
@@ -425,6 +431,25 @@ class OpenAIModel(Model):
425
431
  base64_encoded = base64.b64encode(response.content).decode('utf-8')
426
432
  audio = InputAudio(data=base64_encoded, format=response.headers.get('content-type'))
427
433
  content.append(ChatCompletionContentPartInputAudioParam(input_audio=audio, type='input_audio'))
434
+ elif isinstance(item, DocumentUrl): # pragma: no cover
435
+ raise NotImplementedError('DocumentUrl is not supported for OpenAI')
436
+ # The following implementation should have worked, but it seems we have the following error:
437
+ # pydantic_ai.exceptions.ModelHTTPError: status_code: 400, model_name: gpt-4o, body:
438
+ # {
439
+ # 'message': "Unknown parameter: 'messages[1].content[1].file.data'.",
440
+ # 'type': 'invalid_request_error',
441
+ # 'param': 'messages[1].content[1].file.data',
442
+ # 'code': 'unknown_parameter'
443
+ # }
444
+ #
445
+ # client = cached_async_http_client()
446
+ # response = await client.get(item.url)
447
+ # response.raise_for_status()
448
+ # base64_encoded = base64.b64encode(response.content).decode('utf-8')
449
+ # media_type = response.headers.get('content-type').split(';')[0]
450
+ # file_data = f'data:{media_type};base64,{base64_encoded}'
451
+ # file = File(file={'file_data': file_data, 'file_name': item.url, 'file_id': item.url}, type='file')
452
+ # content.append(file)
428
453
  else:
429
454
  assert_never(item)
430
455
  return chat.ChatCompletionUserMessageParam(role='user', content=content)
@@ -1,10 +1,10 @@
1
1
  from __future__ import annotations as _annotations
2
2
 
3
3
  import functools
4
- from collections.abc import AsyncGenerator
4
+ from collections.abc import AsyncGenerator, Mapping
5
5
  from datetime import datetime, timedelta
6
6
  from pathlib import Path
7
- from typing import Literal
7
+ from typing import Literal, overload
8
8
 
9
9
  import anyio.to_thread
10
10
  import httpx
@@ -52,19 +52,45 @@ class GoogleVertexProvider(Provider[httpx.AsyncClient]):
52
52
  def client(self) -> httpx.AsyncClient:
53
53
  return self._client
54
54
 
55
+ @overload
55
56
  def __init__(
56
57
  self,
58
+ *,
57
59
  service_account_file: Path | str | None = None,
58
60
  project_id: str | None = None,
59
61
  region: VertexAiRegion = 'us-central1',
60
62
  model_publisher: str = 'google',
61
63
  http_client: httpx.AsyncClient | None = None,
64
+ ) -> None: ...
65
+
66
+ @overload
67
+ def __init__(
68
+ self,
69
+ *,
70
+ service_account_info: Mapping[str, str] | None = None,
71
+ project_id: str | None = None,
72
+ region: VertexAiRegion = 'us-central1',
73
+ model_publisher: str = 'google',
74
+ http_client: httpx.AsyncClient | None = None,
75
+ ) -> None: ...
76
+
77
+ def __init__(
78
+ self,
79
+ *,
80
+ service_account_file: Path | str | None = None,
81
+ service_account_info: Mapping[str, str] | None = None,
82
+ project_id: str | None = None,
83
+ region: VertexAiRegion = 'us-central1',
84
+ model_publisher: str = 'google',
85
+ http_client: httpx.AsyncClient | None = None,
62
86
  ) -> None:
63
87
  """Create a new Vertex AI provider.
64
88
 
65
89
  Args:
66
90
  service_account_file: Path to a service account file.
67
- If not provided, the default environment credentials will be used.
91
+ If not provided, the service_account_info or default environment credentials will be used.
92
+ service_account_info: The loaded service_account_file contents.
93
+ If not provided, the service_account_file or default environment credentials will be used.
68
94
  project_id: The project ID to use, if not provided it will be taken from the credentials.
69
95
  region: The region to make requests to.
70
96
  model_publisher: The model publisher to use, I couldn't find a good list of available publishers,
@@ -73,13 +99,17 @@ class GoogleVertexProvider(Provider[httpx.AsyncClient]):
73
99
  Please create an issue or PR if you know how to use other publishers.
74
100
  http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
75
101
  """
102
+ if service_account_file and service_account_info:
103
+ raise ValueError('Only one of `service_account_file` or `service_account_info` can be provided.')
104
+
76
105
  self._client = http_client or cached_async_http_client()
77
106
  self.service_account_file = service_account_file
107
+ self.service_account_info = service_account_info
78
108
  self.project_id = project_id
79
109
  self.region = region
80
110
  self.model_publisher = model_publisher
81
111
 
82
- self._client.auth = _VertexAIAuth(service_account_file, project_id, region)
112
+ self._client.auth = _VertexAIAuth(service_account_file, service_account_info, project_id, region)
83
113
  self._client.base_url = self.base_url
84
114
 
85
115
 
@@ -91,10 +121,12 @@ class _VertexAIAuth(httpx.Auth):
91
121
  def __init__(
92
122
  self,
93
123
  service_account_file: Path | str | None = None,
124
+ service_account_info: Mapping[str, str] | None = None,
94
125
  project_id: str | None = None,
95
126
  region: VertexAiRegion = 'us-central1',
96
127
  ) -> None:
97
128
  self.service_account_file = service_account_file
129
+ self.service_account_info = service_account_info
98
130
  self.project_id = project_id
99
131
  self.region = region
100
132
 
@@ -119,6 +151,11 @@ class _VertexAIAuth(httpx.Auth):
119
151
  assert creds.project_id is None or isinstance(creds.project_id, str) # type: ignore[reportUnknownMemberType]
120
152
  creds_project_id: str | None = creds.project_id
121
153
  creds_source = 'service account file'
154
+ elif self.service_account_info is not None:
155
+ creds = await _creds_from_info(self.service_account_info)
156
+ assert creds.project_id is None or isinstance(creds.project_id, str) # type: ignore[reportUnknownMemberType]
157
+ creds_project_id: str | None = creds.project_id
158
+ creds_source = 'service account info'
122
159
  else:
123
160
  creds, creds_project_id = await _async_google_auth()
124
161
  creds_source = '`google.auth.default()`'
@@ -154,6 +191,14 @@ async def _creds_from_file(service_account_file: str | Path) -> ServiceAccountCr
154
191
  return await anyio.to_thread.run_sync(service_account_credentials_from_file, str(service_account_file))
155
192
 
156
193
 
194
+ async def _creds_from_info(service_account_info: Mapping[str, str]) -> ServiceAccountCredentials:
195
+ service_account_credentials_from_string = functools.partial(
196
+ ServiceAccountCredentials.from_service_account_info, # type: ignore[reportUnknownMemberType]
197
+ scopes=['https://www.googleapis.com/auth/cloud-platform'],
198
+ )
199
+ return await anyio.to_thread.run_sync(service_account_credentials_from_string, service_account_info)
200
+
201
+
157
202
  VertexAiRegion = Literal[
158
203
  'asia-east1',
159
204
  'asia-east2',
@@ -64,7 +64,7 @@ class Usage:
64
64
  }
65
65
  for key, value in (self.details or {}).items():
66
66
  result[f'gen_ai.usage.details.{key}'] = value
67
- return {k: v for k, v in result.items() if v is not None}
67
+ return {k: v for k, v in result.items() if v}
68
68
 
69
69
 
70
70
  @dataclass
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "pydantic-ai-slim"
7
- version = "0.0.36"
7
+ version = "0.0.38"
8
8
  description = "Agent Framework / shim to use Pydantic with LLMs, slim package"
9
9
  authors = [{ name = "Samuel Colvin", email = "samuel@pydantic.dev" }]
10
10
  license = "MIT"
@@ -36,7 +36,7 @@ dependencies = [
36
36
  "griffe>=1.3.2",
37
37
  "httpx>=0.27",
38
38
  "pydantic>=2.10",
39
- "pydantic-graph==0.0.36",
39
+ "pydantic-graph==0.0.38",
40
40
  "exceptiongroup; python_version < '3.11'",
41
41
  "opentelemetry-api>=1.28.0",
42
42
  "typing-inspection>=0.4.0",