pydantic-ai-slim 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

pydantic_ai/{models → profiles}/_json_schema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations as _annotations
+
 import re
 from abc import ABC, abstractmethod
 from copy import deepcopy
@@ -10,7 +12,7 @@ JsonSchema = dict[str, Any]
 
 
 @dataclass(init=False)
-class WalkJsonSchema(ABC):
+class JsonSchemaTransformer(ABC):
     """Walks a JSON schema, applying transformations to it at each level.
 
     Note: We may eventually want to rework tools to build the JSON schema from the type directly, using a subclass of
@@ -18,9 +20,18 @@ class WalkJsonSchema(ABC):
     """
 
     def __init__(
-        self, schema: JsonSchema, *, prefer_inlined_defs: bool = False, simplify_nullable_unions: bool = False
+        self,
+        schema: JsonSchema,
+        *,
+        strict: bool | None = None,
+        prefer_inlined_defs: bool = False,
+        simplify_nullable_unions: bool = False,
     ):
         self.schema = schema
+
+        self.strict = strict
+        self.is_strict_compatible = True  # Can be set to False by subclasses to set `strict` on `ToolDefinition` when set not set by user explicitly
+
         self.prefer_inlined_defs = prefer_inlined_defs
         self.simplify_nullable_unions = simplify_nullable_unions
 
@@ -164,3 +175,13 @@ class WalkJsonSchema(ABC):
                 return [cases[0]]
 
         return cases  # pragma: no cover
+
+
+class InlineDefsJsonSchemaTransformer(JsonSchemaTransformer):
+    """Transforms the JSON Schema to inline $defs."""
+
+    def __init__(self, schema: JsonSchema, *, strict: bool | None = None):
+        super().__init__(schema, strict=strict, prefer_inlined_defs=True)
+
+    def transform(self, schema: JsonSchema) -> JsonSchema:
+        return schema

pydantic_ai/profiles/amazon.py

@@ -0,0 +1,9 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+from ._json_schema import InlineDefsJsonSchemaTransformer
+
+
+def amazon_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for an Amazon model."""
+    return ModelProfile(json_schema_transformer=InlineDefsJsonSchemaTransformer)

pydantic_ai/profiles/anthropic.py

@@ -0,0 +1,8 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+
+
+def anthropic_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for an Anthropic model."""
+    return None

pydantic_ai/profiles/cohere.py

@@ -0,0 +1,8 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+
+
+def cohere_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a Cohere model."""
+    return None

pydantic_ai/profiles/deepseek.py

@@ -0,0 +1,8 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+
+
+def deepseek_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a DeepSeek model."""
+    return None

pydantic_ai/profiles/google.py

@@ -0,0 +1,100 @@
+from __future__ import annotations as _annotations
+
+import warnings
+
+from pydantic_ai.exceptions import UserError
+
+from . import ModelProfile
+from ._json_schema import JsonSchema, JsonSchemaTransformer
+
+
+def google_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a Google model."""
+    return ModelProfile(json_schema_transformer=GoogleJsonSchemaTransformer)
+
+
+class GoogleJsonSchemaTransformer(JsonSchemaTransformer):
+    """Transforms the JSON Schema from Pydantic to be suitable for Gemini.
+
+    Gemini which [supports](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations)
+    a subset of OpenAPI v3.0.3.
+
+    Specifically:
+    * gemini doesn't allow the `title` keyword to be set
+    * gemini doesn't allow `$defs` — we need to inline the definitions where possible
+    """
+
+    def __init__(self, schema: JsonSchema, *, strict: bool | None = None):
+        super().__init__(schema, strict=strict, prefer_inlined_defs=True, simplify_nullable_unions=True)
+
+    def transform(self, schema: JsonSchema) -> JsonSchema:
+        # Note: we need to remove `additionalProperties: False` since it is currently mishandled by Gemini
+        additional_properties = schema.pop(
+            'additionalProperties', None
+        )  # don't pop yet so it's included in the warning
+        if additional_properties:
+            original_schema = {**schema, 'additionalProperties': additional_properties}
+            warnings.warn(
+                '`additionalProperties` is not supported by Gemini; it will be removed from the tool JSON schema.'
+                f' Full schema: {self.schema}\n\n'
+                f'Source of additionalProperties within the full schema: {original_schema}\n\n'
+                'If this came from a field with a type like `dict[str, MyType]`, that field will always be empty.\n\n'
+                "If Google's APIs are updated to support this properly, please create an issue on the PydanticAI GitHub"
+                ' and we will fix this behavior.',
+                UserWarning,
+            )
+
+        schema.pop('title', None)
+        schema.pop('default', None)
+        schema.pop('$schema', None)
+        if (const := schema.pop('const', None)) is not None:  # pragma: no cover
+            # Gemini doesn't support const, but it does support enum with a single value
+            schema['enum'] = [const]
+        schema.pop('discriminator', None)
+        schema.pop('examples', None)
+
+        # TODO: Should we use the trick from pydantic_ai.models.openai._OpenAIJsonSchema
+        #   where we add notes about these properties to the field description?
+        schema.pop('exclusiveMaximum', None)
+        schema.pop('exclusiveMinimum', None)
+
+        # Gemini only supports string enums, so we need to convert any enum values to strings.
+        # Pydantic will take care of transforming the transformed string values to the correct type.
+        if enum := schema.get('enum'):
+            schema['type'] = 'string'
+            schema['enum'] = [str(val) for val in enum]
+
+        type_ = schema.get('type')
+        if 'oneOf' in schema and 'type' not in schema:  # pragma: no cover
+            # This gets hit when we have a discriminated union
+            # Gemini returns an API error in this case even though it says in its error message it shouldn't...
+            # Changing the oneOf to an anyOf prevents the API error and I think is functionally equivalent
+            schema['anyOf'] = schema.pop('oneOf')
+
+        if type_ == 'string' and (fmt := schema.pop('format', None)):
+            description = schema.get('description')
+            if description:
+                schema['description'] = f'{description} (format: {fmt})'
+            else:
+                schema['description'] = f'Format: {fmt}'
+
+        if '$ref' in schema:
+            raise UserError(f'Recursive `$ref`s in JSON Schema are not supported by Gemini: {schema["$ref"]}')
+
+        if 'prefixItems' in schema:
+            # prefixItems is not currently supported in Gemini, so we convert it to items for best compatibility
+            prefix_items = schema.pop('prefixItems')
+            items = schema.get('items')
+            unique_items = [items] if items is not None else []
+            for item in prefix_items:
+                if item not in unique_items:
+                    unique_items.append(item)
+            if len(unique_items) > 1:  # pragma: no cover
+                schema['items'] = {'anyOf': unique_items}
+            elif len(unique_items) == 1:  # pragma: no branch
+                schema['items'] = unique_items[0]
+            schema.setdefault('minItems', len(prefix_items))
+            if items is None:  # pragma: no branch
+                schema.setdefault('maxItems', len(prefix_items))
+
+        return schema

pydantic_ai/profiles/grok.py

@@ -0,0 +1,8 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+
+
+def grok_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a Grok model."""
+    return None

pydantic_ai/profiles/meta.py

@@ -0,0 +1,9 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+from ._json_schema import InlineDefsJsonSchemaTransformer
+
+
+def meta_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a Meta model."""
+    return ModelProfile(json_schema_transformer=InlineDefsJsonSchemaTransformer)

pydantic_ai/profiles/mistral.py

@@ -0,0 +1,8 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+
+
+def mistral_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a Mistral model."""
+    return None

pydantic_ai/profiles/openai.py

@@ -0,0 +1,144 @@
+from __future__ import annotations as _annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any
+
+from . import ModelProfile
+from ._json_schema import JsonSchema, JsonSchemaTransformer
+
+
+@dataclass
+class OpenAIModelProfile(ModelProfile):
+    """Profile for models used with OpenAIModel.
+
+    ALL FIELDS MUST BE `openai_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+    """
+
+    # This can be set by a provider or user if the OpenAI-"compatible" API doesn't support strict tool definitions
+    openai_supports_strict_tool_definition: bool = True
+
+
+def openai_model_profile(model_name: str) -> ModelProfile:
+    """Get the model profile for an OpenAI model."""
+    return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer)
+
+
+_STRICT_INCOMPATIBLE_KEYS = [
+    'minLength',
+    'maxLength',
+    'pattern',
+    'format',
+    'minimum',
+    'maximum',
+    'multipleOf',
+    'patternProperties',
+    'unevaluatedProperties',
+    'propertyNames',
+    'minProperties',
+    'maxProperties',
+    'unevaluatedItems',
+    'contains',
+    'minContains',
+    'maxContains',
+    'minItems',
+    'maxItems',
+    'uniqueItems',
+]
+
+_sentinel = object()
+
+
+@dataclass
+class OpenAIJsonSchemaTransformer(JsonSchemaTransformer):
+    """Recursively handle the schema to make it compatible with OpenAI strict mode.
+
+    See https://platform.openai.com/docs/guides/function-calling?api-mode=responses#strict-mode for more details,
+    but this basically just requires:
+    * `additionalProperties` must be set to false for each object in the parameters
+    * all fields in properties must be marked as required
+    """
+
+    def __init__(self, schema: JsonSchema, *, strict: bool | None = None):
+        super().__init__(schema, strict=strict)
+        self.root_ref = schema.get('$ref')
+
+    def walk(self) -> JsonSchema:
+        # Note: OpenAI does not support anyOf at the root in strict mode
+        # However, we don't need to check for it here because we ensure in pydantic_ai._utils.check_object_json_schema
+        # that the root schema either has type 'object' or is recursive.
+        result = super().walk()
+
+        # For recursive models, we need to tweak the schema to make it compatible with strict mode.
+        # Because the following should never change the semantics of the schema we apply it unconditionally.
+        if self.root_ref is not None:
+            result.pop('$ref', None)  # We replace references to the self.root_ref with just '#' in the transform method
+            root_key = re.sub(r'^#/\$defs/', '', self.root_ref)
+            result.update(self.defs.get(root_key) or {})
+
+        return result
+
+    def transform(self, schema: JsonSchema) -> JsonSchema:  # noqa C901
+        # Remove unnecessary keys
+        schema.pop('title', None)
+        schema.pop('default', None)
+        schema.pop('$schema', None)
+        schema.pop('discriminator', None)
+
+        if schema_ref := schema.get('$ref'):
+            if schema_ref == self.root_ref:
+                schema['$ref'] = '#'
+            if len(schema) > 1:
+                # OpenAI Strict mode doesn't support siblings to "$ref", but _does_ allow siblings to "anyOf".
+                # So if there is a "description" field or any other extra info, we move the "$ref" into an "anyOf":
+                schema['anyOf'] = [{'$ref': schema.pop('$ref')}]
+
+        # Track strict-incompatible keys
+        incompatible_values: dict[str, Any] = {}
+        for key in _STRICT_INCOMPATIBLE_KEYS:
+            value = schema.get(key, _sentinel)
+            if value is not _sentinel:
+                incompatible_values[key] = value
+        description = schema.get('description')
+        if incompatible_values:
+            if self.strict is True:
+                notes: list[str] = []
+                for key, value in incompatible_values.items():
+                    schema.pop(key)
+                    notes.append(f'{key}={value}')
+                notes_string = ', '.join(notes)
+                schema['description'] = notes_string if not description else f'{description} ({notes_string})'
+            elif self.strict is None:  # pragma: no branch
+                self.is_strict_compatible = False
+
+        schema_type = schema.get('type')
+        if 'oneOf' in schema:
+            # OpenAI does not support oneOf in strict mode
+            if self.strict is True:
+                schema['anyOf'] = schema.pop('oneOf')
+            else:
+                self.is_strict_compatible = False
+
+        if schema_type == 'object':
+            if self.strict is True:
+                # additional properties are disallowed
+                schema['additionalProperties'] = False
+
+                # all properties are required
+                if 'properties' not in schema:
+                    schema['properties'] = dict[str, Any]()
+                schema['required'] = list(schema['properties'].keys())
+
+            elif self.strict is None:
+                if (
+                    schema.get('additionalProperties') is not False
+                    or 'properties' not in schema
+                    or 'required' not in schema
+                ):
+                    self.is_strict_compatible = False
+                else:
+                    required = schema['required']
+                    for k in schema['properties'].keys():
+                        if k not in required:
+                            self.is_strict_compatible = False
+        return schema

pydantic_ai/profiles/qwen.py

@@ -0,0 +1,9 @@
+from __future__ import annotations as _annotations
+
+from . import ModelProfile
+from ._json_schema import InlineDefsJsonSchemaTransformer
+
+
+def qwen_model_profile(model_name: str) -> ModelProfile | None:
+    """Get the model profile for a Qwen model."""
+    return ModelProfile(json_schema_transformer=InlineDefsJsonSchemaTransformer)

pydantic_ai/providers/__init__.py

@@ -8,6 +8,8 @@ from __future__ import annotations as _annotations
 from abc import ABC, abstractmethod
 from typing import Any, Generic, TypeVar
 
+from pydantic_ai.profiles import ModelProfile
+
 InterfaceClient = TypeVar('InterfaceClient')
 
 
@@ -41,6 +43,10 @@ class Provider(ABC, Generic[InterfaceClient]):
         """The client for the provider."""
         raise NotImplementedError()
 
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        """The model profile for the named model, if available."""
+        return None  # pragma: no cover
+
 
 def infer_provider(provider: str) -> Provider[Any]:
     """Infer the provider from the provider name."""
@@ -89,5 +95,17 @@ def infer_provider(provider: str) -> Provider[Any]:
         from .cohere import CohereProvider
 
         return CohereProvider()
+    elif provider == 'grok':
+        from .grok import GrokProvider
+
+        return GrokProvider()
+    elif provider == 'fireworks':
+        from .fireworks import FireworksProvider
+
+        return FireworksProvider()
+    elif provider == 'together':
+        from .together import TogetherProvider
+
+        return TogetherProvider()
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')

pydantic_ai/providers/anthropic.py

@@ -7,6 +7,8 @@ import httpx
 
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.anthropic import anthropic_model_profile
 from pydantic_ai.providers import Provider
 
 try:
@@ -33,6 +35,9 @@ class AnthropicProvider(Provider[AsyncAnthropic]):
     def client(self) -> AsyncAnthropic:
         return self._client
 
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        return anthropic_model_profile(model_name)
+
     @overload
     def __init__(self, *, anthropic_client: AsyncAnthropic | None = None) -> None: ...
 

pydantic_ai/providers/azure.py

@@ -8,6 +8,13 @@ from openai import AsyncOpenAI
 
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.cohere import cohere_model_profile
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.grok import grok_model_profile
+from pydantic_ai.profiles.meta import meta_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
 from pydantic_ai.providers import Provider
 
 try:
@@ -38,6 +45,33 @@ class AzureProvider(Provider[AsyncOpenAI]):
     def client(self) -> AsyncOpenAI:
         return self._client
 
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        model_name = model_name.lower()
+
+        prefix_to_profile = {
+            'llama': meta_model_profile,
+            'meta-': meta_model_profile,
+            'deepseek': deepseek_model_profile,
+            'mistralai-': mistral_model_profile,
+            'mistral': mistral_model_profile,
+            'cohere-': cohere_model_profile,
+            'grok': grok_model_profile,
+        }
+
+        for prefix, profile_func in prefix_to_profile.items():
+            if model_name.startswith(prefix):
+                if prefix.endswith('-'):
+                    model_name = model_name[len(prefix) :]
+
+                profile = profile_func(model_name)
+
+                # As AzureProvider is always used with OpenAIModel, which used to unconditionally use OpenAIJsonSchemaTransformer,
+                # we need to maintain that behavior unless json_schema_transformer is set explicitly
+                return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
+
+        # OpenAI models are unprefixed
+        return openai_model_profile(model_name)
+
     @overload
     def __init__(self, *, openai_client: AsyncAzureOpenAI) -> None: ...
 

pydantic_ai/providers/bedrock.py

@@ -1,9 +1,18 @@
 from __future__ import annotations as _annotations
 
 import os
-from typing import overload
+import re
+from dataclasses import dataclass
+from typing import Callable, Literal, overload
 
 from pydantic_ai.exceptions import UserError
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.amazon import amazon_model_profile
+from pydantic_ai.profiles.anthropic import anthropic_model_profile
+from pydantic_ai.profiles.cohere import cohere_model_profile
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.meta import meta_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
 from pydantic_ai.providers import Provider
 
 try:
@@ -18,6 +27,17 @@ except ImportError as _import_error:
     ) from _import_error
 
 
+@dataclass
+class BedrockModelProfile(ModelProfile):
+    """Profile for models used with BedrockModel.
+
+    ALL FIELDS MUST BE `bedrock_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+    """
+
+    bedrock_supports_tool_choice: bool = True
+    bedrock_tool_result_format: Literal['text', 'json'] = 'text'
+
+
 class BedrockProvider(Provider[BaseClient]):
     """Provider for AWS Bedrock."""
 
@@ -33,6 +53,45 @@ class BedrockProvider(Provider[BaseClient]):
     def client(self) -> BaseClient:
         return self._client
 
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        provider_to_profile: dict[str, Callable[[str], ModelProfile | None]] = {
+            'anthropic': lambda model_name: BedrockModelProfile(bedrock_supports_tool_choice=False).update(
+                anthropic_model_profile(model_name)
+            ),
+            'mistral': lambda model_name: BedrockModelProfile(bedrock_tool_result_format='json').update(
+                mistral_model_profile(model_name)
+            ),
+            'cohere': cohere_model_profile,
+            'amazon': amazon_model_profile,
+            'meta': meta_model_profile,
+            'deepseek': deepseek_model_profile,
+        }
+
+        # Split the model name into parts
+        parts = model_name.split('.', 2)
+
+        # Handle regional prefixes (e.g. "us.")
+        if len(parts) > 2 and len(parts[0]) == 2:
+            parts = parts[1:]
+
+        if len(parts) < 2:
+            return None
+
+        provider = parts[0]
+        model_name_with_version = parts[1]
+
+        # Remove version suffix if it matches the format (e.g. "-v1:0" or "-v14")
+        version_match = re.match(r'(.+)-v\d+(?::\d+)?$', model_name_with_version)
+        if version_match:
+            model_name = version_match.group(1)
+        else:
+            model_name = model_name_with_version
+
+        if provider in provider_to_profile:
+            return provider_to_profile[provider](model_name)
+
+        return None
+
     @overload
     def __init__(self, *, bedrock_client: BaseClient) -> None: ...
 

pydantic_ai/providers/cohere.py

@@ -6,6 +6,8 @@ from httpx import AsyncClient as AsyncHTTPClient
 
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.cohere import cohere_model_profile
 from pydantic_ai.providers import Provider
 
 try:
@@ -33,6 +35,9 @@ class CohereProvider(Provider[AsyncClientV2]):
     def client(self) -> AsyncClientV2:
         return self._client
 
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        return cohere_model_profile(model_name)
+
     def __init__(
         self,
         *,

pydantic_ai/providers/deepseek.py

@@ -8,6 +8,9 @@ from openai import AsyncOpenAI
 
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
 from pydantic_ai.providers import Provider
 
 try:
@@ -34,6 +37,15 @@ class DeepSeekProvider(Provider[AsyncOpenAI]):
     def client(self) -> AsyncOpenAI:
         return self._client
 
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        profile = deepseek_model_profile(model_name)
+
+        # As DeepSeekProvider is always used with OpenAIModel, which used to unconditionally use OpenAIJsonSchemaTransformer,
+        # we need to maintain that behavior unless json_schema_transformer is set explicitly.
+        # This was not the case when using a DeepSeek model with another model class (e.g. BedrockConverseModel or GroqModel),
+        # so we won't do this in `deepseek_model_profile` unless we learn it's always needed.
+        return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
+
     @overload
     def __init__(self) -> None: ...