pydantic-ai-slim 0.8.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

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 (75) hide show
  1. pydantic_ai/__init__.py +28 -2
  2. pydantic_ai/_a2a.py +1 -1
  3. pydantic_ai/_agent_graph.py +323 -156
  4. pydantic_ai/_function_schema.py +5 -5
  5. pydantic_ai/_griffe.py +2 -1
  6. pydantic_ai/_otel_messages.py +2 -2
  7. pydantic_ai/_output.py +31 -35
  8. pydantic_ai/_parts_manager.py +7 -5
  9. pydantic_ai/_run_context.py +3 -1
  10. pydantic_ai/_system_prompt.py +2 -2
  11. pydantic_ai/_tool_manager.py +32 -28
  12. pydantic_ai/_utils.py +14 -26
  13. pydantic_ai/ag_ui.py +82 -51
  14. pydantic_ai/agent/__init__.py +70 -9
  15. pydantic_ai/agent/abstract.py +35 -4
  16. pydantic_ai/agent/wrapper.py +6 -0
  17. pydantic_ai/builtin_tools.py +2 -2
  18. pydantic_ai/common_tools/duckduckgo.py +4 -2
  19. pydantic_ai/durable_exec/temporal/__init__.py +4 -2
  20. pydantic_ai/durable_exec/temporal/_agent.py +93 -11
  21. pydantic_ai/durable_exec/temporal/_function_toolset.py +53 -6
  22. pydantic_ai/durable_exec/temporal/_logfire.py +1 -1
  23. pydantic_ai/durable_exec/temporal/_mcp_server.py +2 -1
  24. pydantic_ai/durable_exec/temporal/_model.py +2 -2
  25. pydantic_ai/durable_exec/temporal/_run_context.py +2 -1
  26. pydantic_ai/durable_exec/temporal/_toolset.py +2 -1
  27. pydantic_ai/exceptions.py +45 -2
  28. pydantic_ai/format_prompt.py +2 -2
  29. pydantic_ai/mcp.py +15 -27
  30. pydantic_ai/messages.py +149 -42
  31. pydantic_ai/models/__init__.py +6 -4
  32. pydantic_ai/models/anthropic.py +9 -16
  33. pydantic_ai/models/bedrock.py +50 -56
  34. pydantic_ai/models/cohere.py +3 -3
  35. pydantic_ai/models/fallback.py +2 -2
  36. pydantic_ai/models/function.py +25 -23
  37. pydantic_ai/models/gemini.py +12 -13
  38. pydantic_ai/models/google.py +18 -4
  39. pydantic_ai/models/groq.py +126 -38
  40. pydantic_ai/models/huggingface.py +4 -4
  41. pydantic_ai/models/instrumented.py +35 -16
  42. pydantic_ai/models/mcp_sampling.py +3 -1
  43. pydantic_ai/models/mistral.py +6 -6
  44. pydantic_ai/models/openai.py +35 -40
  45. pydantic_ai/models/test.py +24 -4
  46. pydantic_ai/output.py +27 -32
  47. pydantic_ai/profiles/__init__.py +3 -3
  48. pydantic_ai/profiles/groq.py +1 -1
  49. pydantic_ai/profiles/openai.py +25 -4
  50. pydantic_ai/providers/__init__.py +4 -0
  51. pydantic_ai/providers/anthropic.py +2 -3
  52. pydantic_ai/providers/bedrock.py +3 -2
  53. pydantic_ai/providers/google_vertex.py +2 -1
  54. pydantic_ai/providers/groq.py +21 -2
  55. pydantic_ai/providers/litellm.py +134 -0
  56. pydantic_ai/result.py +144 -41
  57. pydantic_ai/retries.py +52 -31
  58. pydantic_ai/run.py +12 -5
  59. pydantic_ai/tools.py +127 -23
  60. pydantic_ai/toolsets/__init__.py +4 -1
  61. pydantic_ai/toolsets/_dynamic.py +4 -4
  62. pydantic_ai/toolsets/abstract.py +18 -2
  63. pydantic_ai/toolsets/approval_required.py +32 -0
  64. pydantic_ai/toolsets/combined.py +7 -12
  65. pydantic_ai/toolsets/{deferred.py → external.py} +11 -5
  66. pydantic_ai/toolsets/filtered.py +1 -1
  67. pydantic_ai/toolsets/function.py +58 -21
  68. pydantic_ai/toolsets/wrapper.py +2 -1
  69. pydantic_ai/usage.py +44 -8
  70. {pydantic_ai_slim-0.8.1.dist-info → pydantic_ai_slim-1.0.0.dist-info}/METADATA +8 -9
  71. pydantic_ai_slim-1.0.0.dist-info/RECORD +121 -0
  72. pydantic_ai_slim-0.8.1.dist-info/RECORD +0 -119
  73. {pydantic_ai_slim-0.8.1.dist-info → pydantic_ai_slim-1.0.0.dist-info}/WHEEL +0 -0
  74. {pydantic_ai_slim-0.8.1.dist-info → pydantic_ai_slim-1.0.0.dist-info}/entry_points.txt +0 -0
  75. {pydantic_ai_slim-0.8.1.dist-info → pydantic_ai_slim-1.0.0.dist-info}/licenses/LICENSE +0 -0
@@ -195,7 +195,10 @@ class TestModel(Model):
195
195
  # if there are tools, the first thing we want to do is call all of them
196
196
  if tool_calls and not any(isinstance(m, ModelResponse) for m in messages):
197
197
  return ModelResponse(
198
- parts=[ToolCallPart(name, self.gen_tool_args(args)) for name, args in tool_calls],
198
+ parts=[
199
+ ToolCallPart(name, self.gen_tool_args(args), tool_call_id=f'pyd_ai_tool_call_id__{name}')
200
+ for name, args in tool_calls
201
+ ],
199
202
  model_name=self._model_name,
200
203
  )
201
204
 
@@ -220,6 +223,7 @@ class TestModel(Model):
220
223
  output_wrapper.value
221
224
  if isinstance(output_wrapper, _WrappedToolOutput) and output_wrapper.value is not None
222
225
  else self.gen_tool_args(tool),
226
+ tool_call_id=f'pyd_ai_tool_call_id__{tool.name}',
223
227
  )
224
228
  for tool in output_tools
225
229
  if tool.name in new_retry_names
@@ -250,11 +254,27 @@ class TestModel(Model):
250
254
  output_tool = output_tools[self.seed % len(output_tools)]
251
255
  if custom_output_args is not None:
252
256
  return ModelResponse(
253
- parts=[ToolCallPart(output_tool.name, custom_output_args)], model_name=self._model_name
257
+ parts=[
258
+ ToolCallPart(
259
+ output_tool.name,
260
+ custom_output_args,
261
+ tool_call_id=f'pyd_ai_tool_call_id__{output_tool.name}',
262
+ )
263
+ ],
264
+ model_name=self._model_name,
254
265
  )
255
266
  else:
256
267
  response_args = self.gen_tool_args(output_tool)
257
- return ModelResponse(parts=[ToolCallPart(output_tool.name, response_args)], model_name=self._model_name)
268
+ return ModelResponse(
269
+ parts=[
270
+ ToolCallPart(
271
+ output_tool.name,
272
+ response_args,
273
+ tool_call_id=f'pyd_ai_tool_call_id__{output_tool.name}',
274
+ )
275
+ ],
276
+ model_name=self._model_name,
277
+ )
258
278
 
259
279
 
260
280
  @dataclass
@@ -293,7 +313,7 @@ class TestStreamedResponse(StreamedResponse):
293
313
  yield self._parts_manager.handle_tool_call_part(
294
314
  vendor_part_id=i, tool_name=part.tool_name, args=part.args, tool_call_id=part.tool_call_id
295
315
  )
296
- elif isinstance(part, (BuiltinToolCallPart, BuiltinToolReturnPart)): # pragma: no cover
316
+ elif isinstance(part, BuiltinToolCallPart | BuiltinToolReturnPart): # pragma: no cover
297
317
  # NOTE: These parts are not generated by TestModel, but we need to handle them for type checking
298
318
  assert False, f'Unexpected part type in TestModel: {type(part).__name__}'
299
319
  elif isinstance(part, ThinkingPart): # pragma: no cover
pydantic_ai/output.py CHANGED
@@ -1,17 +1,17 @@
1
1
  from __future__ import annotations
2
2
 
3
- from collections.abc import Awaitable, Sequence
3
+ from collections.abc import Awaitable, Callable, Sequence
4
4
  from dataclasses import dataclass
5
- from typing import Any, Callable, Generic, Literal, Union
5
+ from typing import Any, Generic, Literal
6
6
 
7
7
  from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
8
8
  from pydantic.json_schema import JsonSchemaValue
9
9
  from pydantic_core import core_schema
10
- from typing_extensions import TypeAliasType, TypeVar
10
+ from typing_extensions import TypeAliasType, TypeVar, deprecated
11
11
 
12
12
  from . import _utils
13
13
  from .messages import ToolCallPart
14
- from .tools import RunContext, ToolDefinition
14
+ from .tools import DeferredToolRequests, RunContext, ToolDefinition
15
15
 
16
16
  __all__ = (
17
17
  # classes
@@ -42,7 +42,7 @@ StructuredOutputMode = Literal['tool', 'native', 'prompted']
42
42
 
43
43
 
44
44
  OutputTypeOrFunction = TypeAliasType(
45
- 'OutputTypeOrFunction', Union[type[T_co], Callable[..., Union[Awaitable[T_co], T_co]]], type_params=(T_co,)
45
+ 'OutputTypeOrFunction', type[T_co] | Callable[..., Awaitable[T_co] | T_co], type_params=(T_co,)
46
46
  )
47
47
  """Definition of an output type or function.
48
48
 
@@ -54,10 +54,7 @@ See [output docs](../output.md) for more information.
54
54
 
55
55
  TextOutputFunc = TypeAliasType(
56
56
  'TextOutputFunc',
57
- Union[
58
- Callable[[RunContext, str], Union[Awaitable[T_co], T_co]],
59
- Callable[[str], Union[Awaitable[T_co], T_co]],
60
- ],
57
+ Callable[[RunContext, str], Awaitable[T_co] | T_co] | Callable[[str], Awaitable[T_co] | T_co],
61
58
  type_params=(T_co,),
62
59
  )
63
60
  """Definition of a function that will be called to process the model's plain text output. The function must take a single string argument.
@@ -135,10 +132,9 @@ class NativeOutput(Generic[OutputDataT]):
135
132
 
136
133
  Example:
137
134
  ```python {title="native_output.py" requires="tool_output.py"}
138
- from tool_output import Fruit, Vehicle
139
-
140
135
  from pydantic_ai import Agent, NativeOutput
141
136
 
137
+ from tool_output import Fruit, Vehicle
142
138
 
143
139
  agent = Agent(
144
140
  'openai:gpt-4o',
@@ -184,10 +180,11 @@ class PromptedOutput(Generic[OutputDataT]):
184
180
  Example:
185
181
  ```python {title="prompted_output.py" requires="tool_output.py"}
186
182
  from pydantic import BaseModel
187
- from tool_output import Vehicle
188
183
 
189
184
  from pydantic_ai import Agent, PromptedOutput
190
185
 
186
+ from tool_output import Vehicle
187
+
191
188
 
192
189
  class Device(BaseModel):
193
190
  name: str
@@ -286,18 +283,17 @@ def StructuredDict(
286
283
  ```python {title="structured_dict.py"}
287
284
  from pydantic_ai import Agent, StructuredDict
288
285
 
289
-
290
286
  schema = {
291
- "type": "object",
292
- "properties": {
293
- "name": {"type": "string"},
294
- "age": {"type": "integer"}
287
+ 'type': 'object',
288
+ 'properties': {
289
+ 'name': {'type': 'string'},
290
+ 'age': {'type': 'integer'}
295
291
  },
296
- "required": ["name", "age"]
292
+ 'required': ['name', 'age']
297
293
  }
298
294
 
299
295
  agent = Agent('openai:gpt-4o', output_type=StructuredDict(schema))
300
- result = agent.run_sync("Create a person")
296
+ result = agent.run_sync('Create a person')
301
297
  print(result.output)
302
298
  #> {'name': 'John Doe', 'age': 30}
303
299
  ```
@@ -333,16 +329,13 @@ def StructuredDict(
333
329
 
334
330
  _OutputSpecItem = TypeAliasType(
335
331
  '_OutputSpecItem',
336
- Union[OutputTypeOrFunction[T_co], ToolOutput[T_co], NativeOutput[T_co], PromptedOutput[T_co], TextOutput[T_co]],
332
+ OutputTypeOrFunction[T_co] | ToolOutput[T_co] | NativeOutput[T_co] | PromptedOutput[T_co] | TextOutput[T_co],
337
333
  type_params=(T_co,),
338
334
  )
339
335
 
340
336
  OutputSpec = TypeAliasType(
341
337
  'OutputSpec',
342
- Union[
343
- _OutputSpecItem[T_co],
344
- Sequence['OutputSpec[T_co]'],
345
- ],
338
+ _OutputSpecItem[T_co] | Sequence['OutputSpec[T_co]'],
346
339
  type_params=(T_co,),
347
340
  )
348
341
  """Specification of the agent's output data.
@@ -359,12 +352,14 @@ See [output docs](../output.md) for more information.
359
352
  """
360
353
 
361
354
 
362
- @dataclass
363
- class DeferredToolCalls:
364
- """Container for calls of deferred tools. This can be used as an agent's `output_type` and will be used as the output of the agent run if the model called any deferred tools.
365
-
366
- See [deferred toolset docs](../toolsets.md#deferred-toolset) for more information.
367
- """
355
+ @deprecated('`DeferredToolCalls` is deprecated, use `DeferredToolRequests` instead')
356
+ class DeferredToolCalls(DeferredToolRequests): # pragma: no cover
357
+ @property
358
+ @deprecated('`DeferredToolCalls.tool_calls` is deprecated, use `DeferredToolRequests.calls` instead')
359
+ def tool_calls(self) -> list[ToolCallPart]:
360
+ return self.calls
368
361
 
369
- tool_calls: list[ToolCallPart]
370
- tool_defs: dict[str, ToolDefinition]
362
+ @property
363
+ @deprecated('`DeferredToolCalls.tool_defs` is deprecated')
364
+ def tool_defs(self) -> dict[str, ToolDefinition]:
365
+ return {}
@@ -1,8 +1,8 @@
1
1
  from __future__ import annotations as _annotations
2
2
 
3
+ from collections.abc import Callable
3
4
  from dataclasses import dataclass, fields, replace
4
5
  from textwrap import dedent
5
- from typing import Callable, Union
6
6
 
7
7
  from typing_extensions import Self
8
8
 
@@ -18,7 +18,7 @@ __all__ = [
18
18
  ]
19
19
 
20
20
 
21
- @dataclass
21
+ @dataclass(kw_only=True)
22
22
  class ModelProfile:
23
23
  """Describes how requests to and responses from specific models or families of models need to be constructed and processed to get the best results, independent of the model and provider classes used."""
24
24
 
@@ -75,6 +75,6 @@ class ModelProfile:
75
75
  return replace(self, **non_default_attrs)
76
76
 
77
77
 
78
- ModelProfileSpec = Union[ModelProfile, Callable[[str], Union[ModelProfile, None]]]
78
+ ModelProfileSpec = ModelProfile | Callable[[str], ModelProfile | None]
79
79
 
80
80
  DEFAULT_PROFILE = ModelProfile()
@@ -5,7 +5,7 @@ from dataclasses import dataclass
5
5
  from . import ModelProfile
6
6
 
7
7
 
8
- @dataclass
8
+ @dataclass(kw_only=True)
9
9
  class GroqModelProfile(ModelProfile):
10
10
  """Profile for models used with GroqModel.
11
11
 
@@ -1,6 +1,7 @@
1
1
  from __future__ import annotations as _annotations
2
2
 
3
3
  import re
4
+ import warnings
4
5
  from collections.abc import Sequence
5
6
  from dataclasses import dataclass
6
7
  from typing import Any, Literal
@@ -11,7 +12,7 @@ from ._json_schema import JsonSchema, JsonSchemaTransformer
11
12
  OpenAISystemPromptRole = Literal['system', 'developer', 'user']
12
13
 
13
14
 
14
- @dataclass
15
+ @dataclass(kw_only=True)
15
16
  class OpenAIModelProfile(ModelProfile):
16
17
  """Profile for models used with `OpenAIChatModel`.
17
18
 
@@ -21,7 +22,6 @@ class OpenAIModelProfile(ModelProfile):
21
22
  openai_supports_strict_tool_definition: bool = True
22
23
  """This can be set by a provider or user if the OpenAI-"compatible" API doesn't support strict tool definitions."""
23
24
 
24
- # TODO(Marcelo): Deprecate this in favor of `openai_unsupported_model_settings`.
25
25
  openai_supports_sampling_settings: bool = True
26
26
  """Turn off to don't send sampling settings like `temperature` and `top_p` to models that don't support them, like OpenAI's o-series reasoning models."""
27
27
 
@@ -38,6 +38,14 @@ class OpenAIModelProfile(ModelProfile):
38
38
  openai_system_prompt_role: OpenAISystemPromptRole | None = None
39
39
  """The role to use for the system prompt message. If not provided, defaults to `'system'`."""
40
40
 
41
+ def __post_init__(self): # pragma: no cover
42
+ if not self.openai_supports_sampling_settings:
43
+ warnings.warn(
44
+ 'The `openai_supports_sampling_settings` has no effect, and it will be removed in future versions. '
45
+ 'Use `openai_unsupported_model_settings` instead.',
46
+ DeprecationWarning,
47
+ )
48
+
41
49
 
42
50
  def openai_model_profile(model_name: str) -> ModelProfile:
43
51
  """Get the model profile for an OpenAI model."""
@@ -46,6 +54,19 @@ def openai_model_profile(model_name: str) -> ModelProfile:
46
54
  # We leave it in here for all models because the `default_structured_output_mode` is `'tool'`, so `native` is only used
47
55
  # when the user specifically uses the `NativeOutput` marker, so an error from the API is acceptable.
48
56
 
57
+ if is_reasoning_model:
58
+ openai_unsupported_model_settings = (
59
+ 'temperature',
60
+ 'top_p',
61
+ 'presence_penalty',
62
+ 'frequency_penalty',
63
+ 'logit_bias',
64
+ 'logprobs',
65
+ 'top_logprobs',
66
+ )
67
+ else:
68
+ openai_unsupported_model_settings = ()
69
+
49
70
  # The o1-mini model doesn't support the `system` role, so we default to `user`.
50
71
  # See https://github.com/pydantic/pydantic-ai/issues/974 for more details.
51
72
  openai_system_prompt_role = 'user' if model_name.startswith('o1-mini') else None
@@ -54,7 +75,7 @@ def openai_model_profile(model_name: str) -> ModelProfile:
54
75
  json_schema_transformer=OpenAIJsonSchemaTransformer,
55
76
  supports_json_schema_output=True,
56
77
  supports_json_object_output=True,
57
- openai_supports_sampling_settings=not is_reasoning_model,
78
+ openai_unsupported_model_settings=openai_unsupported_model_settings,
58
79
  openai_system_prompt_role=openai_system_prompt_role,
59
80
  )
60
81
 
@@ -89,7 +110,7 @@ _STRICT_COMPATIBLE_STRING_FORMATS = [
89
110
  _sentinel = object()
90
111
 
91
112
 
92
- @dataclass
113
+ @dataclass(init=False)
93
114
  class OpenAIJsonSchemaTransformer(JsonSchemaTransformer):
94
115
  """Recursively handle the schema to make it compatible with OpenAI strict mode.
95
116
 
@@ -135,6 +135,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]: # noqa: C901
135
135
  from .github import GitHubProvider
136
136
 
137
137
  return GitHubProvider
138
+ elif provider == 'litellm':
139
+ from .litellm import LiteLLMProvider
140
+
141
+ return LiteLLMProvider
138
142
  else: # pragma: no cover
139
143
  raise ValueError(f'Unknown provider: {provider}')
140
144
 
@@ -1,10 +1,9 @@
1
1
  from __future__ import annotations as _annotations
2
2
 
3
3
  import os
4
- from typing import Union, overload
4
+ from typing import TypeAlias, overload
5
5
 
6
6
  import httpx
7
- from typing_extensions import TypeAlias
8
7
 
9
8
  from pydantic_ai.exceptions import UserError
10
9
  from pydantic_ai.models import cached_async_http_client
@@ -21,7 +20,7 @@ except ImportError as _import_error:
21
20
  ) from _import_error
22
21
 
23
22
 
24
- AsyncAnthropicClient: TypeAlias = Union[AsyncAnthropic, AsyncAnthropicBedrock]
23
+ AsyncAnthropicClient: TypeAlias = AsyncAnthropic | AsyncAnthropicBedrock
25
24
 
26
25
 
27
26
  class AnthropicProvider(Provider[AsyncAnthropicClient]):
@@ -2,8 +2,9 @@ from __future__ import annotations as _annotations
2
2
 
3
3
  import os
4
4
  import re
5
+ from collections.abc import Callable
5
6
  from dataclasses import dataclass
6
- from typing import Callable, Literal, overload
7
+ from typing import Literal, overload
7
8
 
8
9
  from pydantic_ai.exceptions import UserError
9
10
  from pydantic_ai.profiles import ModelProfile
@@ -27,7 +28,7 @@ except ImportError as _import_error:
27
28
  ) from _import_error
28
29
 
29
30
 
30
- @dataclass
31
+ @dataclass(kw_only=True)
31
32
  class BedrockModelProfile(ModelProfile):
32
33
  """Profile for models used with BedrockModel.
33
34
 
@@ -1,6 +1,7 @@
1
1
  from __future__ import annotations as _annotations
2
2
 
3
3
  import functools
4
+ from asyncio import Lock
4
5
  from collections.abc import AsyncGenerator, Mapping
5
6
  from pathlib import Path
6
7
  from typing import Literal, overload
@@ -118,7 +119,7 @@ class GoogleVertexProvider(Provider[httpx.AsyncClient]):
118
119
  class _VertexAIAuth(httpx.Auth):
119
120
  """Auth class for Vertex AI API."""
120
121
 
121
- _refresh_lock: anyio.Lock = anyio.Lock()
122
+ _refresh_lock: Lock = Lock()
122
123
 
123
124
  credentials: BaseCredentials | ServiceAccountCredentials | None
124
125
 
@@ -14,6 +14,7 @@ from pydantic_ai.profiles.groq import groq_model_profile
14
14
  from pydantic_ai.profiles.meta import meta_model_profile
15
15
  from pydantic_ai.profiles.mistral import mistral_model_profile
16
16
  from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
17
+ from pydantic_ai.profiles.openai import openai_model_profile
17
18
  from pydantic_ai.profiles.qwen import qwen_model_profile
18
19
  from pydantic_ai.providers import Provider
19
20
 
@@ -26,6 +27,23 @@ except ImportError as _import_error: # pragma: no cover
26
27
  ) from _import_error
27
28
 
28
29
 
30
+ def groq_moonshotai_model_profile(model_name: str) -> ModelProfile | None:
31
+ """Get the model profile for an MoonshotAI model used with the Groq provider."""
32
+ return ModelProfile(supports_json_object_output=True, supports_json_schema_output=True).update(
33
+ moonshotai_model_profile(model_name)
34
+ )
35
+
36
+
37
+ def meta_groq_model_profile(model_name: str) -> ModelProfile | None:
38
+ """Get the model profile for a Meta model used with the Groq provider."""
39
+ if model_name in {'llama-4-maverick-17b-128e-instruct', 'llama-4-scout-17b-16e-instruct'}:
40
+ return ModelProfile(supports_json_object_output=True, supports_json_schema_output=True).update(
41
+ meta_model_profile(model_name)
42
+ )
43
+ else:
44
+ return meta_model_profile(model_name)
45
+
46
+
29
47
  class GroqProvider(Provider[AsyncGroq]):
30
48
  """Provider for Groq API."""
31
49
 
@@ -44,13 +62,14 @@ class GroqProvider(Provider[AsyncGroq]):
44
62
  def model_profile(self, model_name: str) -> ModelProfile | None:
45
63
  prefix_to_profile = {
46
64
  'llama': meta_model_profile,
47
- 'meta-llama/': meta_model_profile,
65
+ 'meta-llama/': meta_groq_model_profile,
48
66
  'gemma': google_model_profile,
49
67
  'qwen': qwen_model_profile,
50
68
  'deepseek': deepseek_model_profile,
51
69
  'mistral': mistral_model_profile,
52
- 'moonshotai/': moonshotai_model_profile,
70
+ 'moonshotai/': groq_moonshotai_model_profile,
53
71
  'compound-': groq_model_profile,
72
+ 'openai/': openai_model_profile,
54
73
  }
55
74
 
56
75
  for prefix, profile_func in prefix_to_profile.items():
@@ -0,0 +1,134 @@
1
+ from __future__ import annotations as _annotations
2
+
3
+ from typing import overload
4
+
5
+ from httpx import AsyncClient as AsyncHTTPClient
6
+ from openai import AsyncOpenAI
7
+
8
+ from pydantic_ai.models import cached_async_http_client
9
+ from pydantic_ai.profiles import ModelProfile
10
+ from pydantic_ai.profiles.amazon import amazon_model_profile
11
+ from pydantic_ai.profiles.anthropic import anthropic_model_profile
12
+ from pydantic_ai.profiles.cohere import cohere_model_profile
13
+ from pydantic_ai.profiles.deepseek import deepseek_model_profile
14
+ from pydantic_ai.profiles.google import google_model_profile
15
+ from pydantic_ai.profiles.grok import grok_model_profile
16
+ from pydantic_ai.profiles.groq import groq_model_profile
17
+ from pydantic_ai.profiles.meta import meta_model_profile
18
+ from pydantic_ai.profiles.mistral import mistral_model_profile
19
+ from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
20
+ from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
21
+ from pydantic_ai.profiles.qwen import qwen_model_profile
22
+ from pydantic_ai.providers import Provider
23
+
24
+ try:
25
+ from openai import AsyncOpenAI
26
+ except ImportError as _import_error: # pragma: no cover
27
+ raise ImportError(
28
+ 'Please install the `openai` package to use the LiteLLM provider, '
29
+ 'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
30
+ ) from _import_error
31
+
32
+
33
+ class LiteLLMProvider(Provider[AsyncOpenAI]):
34
+ """Provider for LiteLLM API."""
35
+
36
+ @property
37
+ def name(self) -> str:
38
+ return 'litellm'
39
+
40
+ @property
41
+ def base_url(self) -> str:
42
+ return str(self.client.base_url)
43
+
44
+ @property
45
+ def client(self) -> AsyncOpenAI:
46
+ return self._client
47
+
48
+ def model_profile(self, model_name: str) -> ModelProfile | None:
49
+ # Map provider prefixes to their profile functions
50
+ provider_to_profile = {
51
+ 'anthropic': anthropic_model_profile,
52
+ 'openai': openai_model_profile,
53
+ 'google': google_model_profile,
54
+ 'mistralai': mistral_model_profile,
55
+ 'mistral': mistral_model_profile,
56
+ 'cohere': cohere_model_profile,
57
+ 'amazon': amazon_model_profile,
58
+ 'bedrock': amazon_model_profile,
59
+ 'meta-llama': meta_model_profile,
60
+ 'meta': meta_model_profile,
61
+ 'groq': groq_model_profile,
62
+ 'deepseek': deepseek_model_profile,
63
+ 'moonshotai': moonshotai_model_profile,
64
+ 'x-ai': grok_model_profile,
65
+ 'qwen': qwen_model_profile,
66
+ }
67
+
68
+ profile = None
69
+
70
+ # Check if model name contains a provider prefix (e.g., "anthropic/claude-3")
71
+ if '/' in model_name:
72
+ provider_prefix, model_suffix = model_name.split('/', 1)
73
+ if provider_prefix in provider_to_profile:
74
+ profile = provider_to_profile[provider_prefix](model_suffix)
75
+
76
+ # If no profile found, default to OpenAI profile
77
+ if profile is None:
78
+ profile = openai_model_profile(model_name)
79
+
80
+ # As LiteLLMProvider is used with OpenAIModel, which uses OpenAIJsonSchemaTransformer,
81
+ # we maintain that behavior
82
+ return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
83
+
84
+ @overload
85
+ def __init__(
86
+ self,
87
+ *,
88
+ api_key: str | None = None,
89
+ api_base: str | None = None,
90
+ ) -> None: ...
91
+
92
+ @overload
93
+ def __init__(
94
+ self,
95
+ *,
96
+ api_key: str | None = None,
97
+ api_base: str | None = None,
98
+ http_client: AsyncHTTPClient,
99
+ ) -> None: ...
100
+
101
+ @overload
102
+ def __init__(self, *, openai_client: AsyncOpenAI) -> None: ...
103
+
104
+ def __init__(
105
+ self,
106
+ *,
107
+ api_key: str | None = None,
108
+ api_base: str | None = None,
109
+ openai_client: AsyncOpenAI | None = None,
110
+ http_client: AsyncHTTPClient | None = None,
111
+ ) -> None:
112
+ """Initialize a LiteLLM provider.
113
+
114
+ Args:
115
+ api_key: API key for the model provider. If None, LiteLLM will try to get it from environment variables.
116
+ api_base: Base URL for the model provider. Use this for custom endpoints or self-hosted models.
117
+ openai_client: Pre-configured OpenAI client. If provided, other parameters are ignored.
118
+ http_client: Custom HTTP client to use.
119
+ """
120
+ if openai_client is not None:
121
+ self._client = openai_client
122
+ return
123
+
124
+ # Create OpenAI client that will be used with LiteLLM's completion function
125
+ # The actual API calls will be intercepted and routed through LiteLLM
126
+ if http_client is not None:
127
+ self._client = AsyncOpenAI(
128
+ base_url=api_base, api_key=api_key or 'litellm-placeholder', http_client=http_client
129
+ )
130
+ else:
131
+ http_client = cached_async_http_client(provider='litellm')
132
+ self._client = AsyncOpenAI(
133
+ base_url=api_base, api_key=api_key or 'litellm-placeholder', http_client=http_client
134
+ )