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.
- pydantic_ai/__init__.py +28 -2
- pydantic_ai/_a2a.py +1 -1
- pydantic_ai/_agent_graph.py +323 -156
- pydantic_ai/_function_schema.py +5 -5
- pydantic_ai/_griffe.py +2 -1
- pydantic_ai/_otel_messages.py +2 -2
- pydantic_ai/_output.py +31 -35
- pydantic_ai/_parts_manager.py +7 -5
- pydantic_ai/_run_context.py +3 -1
- pydantic_ai/_system_prompt.py +2 -2
- pydantic_ai/_tool_manager.py +32 -28
- pydantic_ai/_utils.py +14 -26
- pydantic_ai/ag_ui.py +82 -51
- pydantic_ai/agent/__init__.py +70 -9
- pydantic_ai/agent/abstract.py +35 -4
- pydantic_ai/agent/wrapper.py +6 -0
- pydantic_ai/builtin_tools.py +2 -2
- pydantic_ai/common_tools/duckduckgo.py +4 -2
- pydantic_ai/durable_exec/temporal/__init__.py +4 -2
- pydantic_ai/durable_exec/temporal/_agent.py +93 -11
- pydantic_ai/durable_exec/temporal/_function_toolset.py +53 -6
- pydantic_ai/durable_exec/temporal/_logfire.py +1 -1
- pydantic_ai/durable_exec/temporal/_mcp_server.py +2 -1
- pydantic_ai/durable_exec/temporal/_model.py +2 -2
- pydantic_ai/durable_exec/temporal/_run_context.py +2 -1
- pydantic_ai/durable_exec/temporal/_toolset.py +2 -1
- pydantic_ai/exceptions.py +45 -2
- pydantic_ai/format_prompt.py +2 -2
- pydantic_ai/mcp.py +15 -27
- pydantic_ai/messages.py +149 -42
- pydantic_ai/models/__init__.py +6 -4
- pydantic_ai/models/anthropic.py +9 -16
- pydantic_ai/models/bedrock.py +50 -56
- pydantic_ai/models/cohere.py +3 -3
- pydantic_ai/models/fallback.py +2 -2
- pydantic_ai/models/function.py +25 -23
- pydantic_ai/models/gemini.py +12 -13
- pydantic_ai/models/google.py +18 -4
- pydantic_ai/models/groq.py +126 -38
- pydantic_ai/models/huggingface.py +4 -4
- pydantic_ai/models/instrumented.py +35 -16
- pydantic_ai/models/mcp_sampling.py +3 -1
- pydantic_ai/models/mistral.py +6 -6
- pydantic_ai/models/openai.py +35 -40
- pydantic_ai/models/test.py +24 -4
- pydantic_ai/output.py +27 -32
- pydantic_ai/profiles/__init__.py +3 -3
- pydantic_ai/profiles/groq.py +1 -1
- pydantic_ai/profiles/openai.py +25 -4
- pydantic_ai/providers/__init__.py +4 -0
- pydantic_ai/providers/anthropic.py +2 -3
- pydantic_ai/providers/bedrock.py +3 -2
- pydantic_ai/providers/google_vertex.py +2 -1
- pydantic_ai/providers/groq.py +21 -2
- pydantic_ai/providers/litellm.py +134 -0
- pydantic_ai/result.py +144 -41
- pydantic_ai/retries.py +52 -31
- pydantic_ai/run.py +12 -5
- pydantic_ai/tools.py +127 -23
- pydantic_ai/toolsets/__init__.py +4 -1
- pydantic_ai/toolsets/_dynamic.py +4 -4
- pydantic_ai/toolsets/abstract.py +18 -2
- pydantic_ai/toolsets/approval_required.py +32 -0
- pydantic_ai/toolsets/combined.py +7 -12
- pydantic_ai/toolsets/{deferred.py → external.py} +11 -5
- pydantic_ai/toolsets/filtered.py +1 -1
- pydantic_ai/toolsets/function.py +58 -21
- pydantic_ai/toolsets/wrapper.py +2 -1
- pydantic_ai/usage.py +44 -8
- {pydantic_ai_slim-0.8.1.dist-info → pydantic_ai_slim-1.0.0.dist-info}/METADATA +8 -9
- pydantic_ai_slim-1.0.0.dist-info/RECORD +121 -0
- pydantic_ai_slim-0.8.1.dist-info/RECORD +0 -119
- {pydantic_ai_slim-0.8.1.dist-info → pydantic_ai_slim-1.0.0.dist-info}/WHEEL +0 -0
- {pydantic_ai_slim-0.8.1.dist-info → pydantic_ai_slim-1.0.0.dist-info}/entry_points.txt +0 -0
- {pydantic_ai_slim-0.8.1.dist-info → pydantic_ai_slim-1.0.0.dist-info}/licenses/LICENSE +0 -0
pydantic_ai/models/test.py
CHANGED
|
@@ -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=[
|
|
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=[
|
|
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(
|
|
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,
|
|
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,
|
|
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',
|
|
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
|
-
|
|
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
|
-
|
|
292
|
-
|
|
293
|
-
|
|
294
|
-
|
|
287
|
+
'type': 'object',
|
|
288
|
+
'properties': {
|
|
289
|
+
'name': {'type': 'string'},
|
|
290
|
+
'age': {'type': 'integer'}
|
|
295
291
|
},
|
|
296
|
-
|
|
292
|
+
'required': ['name', 'age']
|
|
297
293
|
}
|
|
298
294
|
|
|
299
295
|
agent = Agent('openai:gpt-4o', output_type=StructuredDict(schema))
|
|
300
|
-
result = agent.run_sync(
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
@
|
|
363
|
-
class DeferredToolCalls:
|
|
364
|
-
|
|
365
|
-
|
|
366
|
-
|
|
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
|
-
|
|
370
|
-
tool_defs
|
|
362
|
+
@property
|
|
363
|
+
@deprecated('`DeferredToolCalls.tool_defs` is deprecated')
|
|
364
|
+
def tool_defs(self) -> dict[str, ToolDefinition]:
|
|
365
|
+
return {}
|
pydantic_ai/profiles/__init__.py
CHANGED
|
@@ -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 =
|
|
78
|
+
ModelProfileSpec = ModelProfile | Callable[[str], ModelProfile | None]
|
|
79
79
|
|
|
80
80
|
DEFAULT_PROFILE = ModelProfile()
|
pydantic_ai/profiles/groq.py
CHANGED
pydantic_ai/profiles/openai.py
CHANGED
|
@@ -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
|
-
|
|
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
|
|
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 =
|
|
23
|
+
AsyncAnthropicClient: TypeAlias = AsyncAnthropic | AsyncAnthropicBedrock
|
|
25
24
|
|
|
26
25
|
|
|
27
26
|
class AnthropicProvider(Provider[AsyncAnthropicClient]):
|
pydantic_ai/providers/bedrock.py
CHANGED
|
@@ -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
|
|
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:
|
|
122
|
+
_refresh_lock: Lock = Lock()
|
|
122
123
|
|
|
123
124
|
credentials: BaseCredentials | ServiceAccountCredentials | None
|
|
124
125
|
|
pydantic_ai/providers/groq.py
CHANGED
|
@@ -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/':
|
|
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/':
|
|
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
|
+
)
|