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

pydantic_ai/models/groq.py

@@ -5,10 +5,13 @@ from collections.abc import AsyncIterable, AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, Literal, Union, cast, overload
+from typing import Any, Literal, cast, overload
 
+from pydantic import BaseModel, Json, ValidationError
 from typing_extensions import assert_never
 
+from pydantic_ai._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition
+
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._run_context import RunContext
 from .._thinking_part import split_content_into_text_and_thinking
@@ -48,7 +51,7 @@ from . import (
 )
 
 try:
-    from groq import NOT_GIVEN, APIStatusError, AsyncGroq, AsyncStream
+    from groq import NOT_GIVEN, APIError, APIStatusError, AsyncGroq, AsyncStream
     from groq.types import chat
     from groq.types.chat.chat_completion_content_part_image_param import ImageURL
 except ImportError as _import_error:
@@ -88,7 +91,7 @@ PreviewGroqModelNames = Literal[
 ]
 """Preview Groq models from <https://console.groq.com/docs/models#preview-models>."""
 
-GroqModelName = Union[str, ProductionGroqModelNames, PreviewGroqModelNames]
+GroqModelName = str | ProductionGroqModelNames | PreviewGroqModelNames
 """Possible Groq model names.
 
 Since Groq supports a variety of models and the list changes frequencly, we explicitly list the named models as of 2025-03-31
@@ -169,9 +172,24 @@ class GroqModel(Model):
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         check_allow_model_requests()
-        response = await self._completions_create(
-            messages, False, cast(GroqModelSettings, model_settings or {}), model_request_parameters
-        )
+        try:
+            response = await self._completions_create(
+                messages, False, cast(GroqModelSettings, model_settings or {}), model_request_parameters
+            )
+        except ModelHTTPError as e:
+            if isinstance(e.body, dict):  # pragma: no branch
+                # The Groq SDK tries to be helpful by raising an exception when generated tool arguments don't match the schema,
+                # but we'd rather handle it ourselves so we can tell the model to retry the tool call.
+                try:
+                    error = _GroqToolUseFailedError.model_validate(e.body)  # pyright: ignore[reportUnknownMemberType]
+                    tool_call_part = ToolCallPart(
+                        tool_name=error.error.failed_generation.name,
+                        args=error.error.failed_generation.arguments,
+                    )
+                    return ModelResponse(parts=[tool_call_part])
+                except ValidationError:
+                    pass
+            raise
         model_response = self._process_response(response)
         return model_response
 
@@ -228,6 +246,18 @@ class GroqModel(Model):
 
         groq_messages = self._map_messages(messages)
 
+        response_format: chat.completion_create_params.ResponseFormat | None = None
+        if model_request_parameters.output_mode == 'native':
+            output_object = model_request_parameters.output_object
+            assert output_object is not None
+            response_format = self._map_json_schema(output_object)
+        elif (
+            model_request_parameters.output_mode == 'prompted'
+            and not tools
+            and self.profile.supports_json_object_output
+        ):  # pragma: no branch
+            response_format = {'type': 'json_object'}
+
         try:
             extra_headers = model_settings.get('extra_headers', {})
             extra_headers.setdefault('User-Agent', get_user_agent())
@@ -240,6 +270,7 @@ class GroqModel(Model):
                 tool_choice=tool_choice or NOT_GIVEN,
                 stop=model_settings.get('stop_sequences', NOT_GIVEN),
                 stream=stream,
+                response_format=response_format or NOT_GIVEN,
                 max_tokens=model_settings.get('max_tokens', NOT_GIVEN),
                 temperature=model_settings.get('temperature', NOT_GIVEN),
                 top_p=model_settings.get('top_p', NOT_GIVEN),
@@ -285,11 +316,11 @@ class GroqModel(Model):
             for c in choice.message.tool_calls:
                 items.append(ToolCallPart(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id))
         return ModelResponse(
-            items,
+            parts=items,
             usage=_map_usage(response),
             model_name=response.model,
             timestamp=timestamp,
-            provider_request_id=response.id,
+            provider_response_id=response.id,
             provider_name=self._provider.name,
         )
 
@@ -347,7 +378,7 @@ class GroqModel(Model):
                     elif isinstance(item, ThinkingPart):
                         # Skip thinking parts when mapping to Groq messages
                         continue
-                    elif isinstance(item, (BuiltinToolCallPart, BuiltinToolReturnPart)):  # pragma: no cover
+                    elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
                         # This is currently never returned from groq
                         pass
                     else:
@@ -385,6 +416,19 @@ class GroqModel(Model):
             },
         }
 
+    def _map_json_schema(self, o: OutputObjectDefinition) -> chat.completion_create_params.ResponseFormat:
+        response_format_param: chat.completion_create_params.ResponseFormatResponseFormatJsonSchema = {
+            'type': 'json_schema',
+            'json_schema': {
+                'name': o.name or DEFAULT_OUTPUT_TOOL_NAME,
+                'schema': o.json_schema,
+                'strict': o.strict,
+            },
+        }
+        if o.description:  # pragma: no branch
+            response_format_param['json_schema']['description'] = o.description
+        return response_format_param
+
     @classmethod
     def _map_user_message(cls, message: ModelRequest) -> Iterable[chat.ChatCompletionMessageParam]:
         for part in message.parts:
@@ -449,36 +493,52 @@ class GroqStreamedResponse(StreamedResponse):
     _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
-        async for chunk in self._response:
-            self._usage += _map_usage(chunk)
-
-            try:
-                choice = chunk.choices[0]
-            except IndexError:
-                continue
-
-            # Handle the text part of the response
-            content = choice.delta.content
-            if content is not None:
-                maybe_event = self._parts_manager.handle_text_delta(
-                    vendor_part_id='content',
-                    content=content,
-                    thinking_tags=self._model_profile.thinking_tags,
-                    ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
-                )
-                if maybe_event is not None:  # pragma: no branch
-                    yield maybe_event
-
-            # Handle the tool calls
-            for dtc in choice.delta.tool_calls or []:
-                maybe_event = self._parts_manager.handle_tool_call_delta(
-                    vendor_part_id=dtc.index,
-                    tool_name=dtc.function and dtc.function.name,
-                    args=dtc.function and dtc.function.arguments,
-                    tool_call_id=dtc.id,
-                )
-                if maybe_event is not None:
-                    yield maybe_event
+        try:
+            async for chunk in self._response:
+                self._usage += _map_usage(chunk)
+
+                try:
+                    choice = chunk.choices[0]
+                except IndexError:
+                    continue
+
+                # Handle the text part of the response
+                content = choice.delta.content
+                if content is not None:
+                    maybe_event = self._parts_manager.handle_text_delta(
+                        vendor_part_id='content',
+                        content=content,
+                        thinking_tags=self._model_profile.thinking_tags,
+                        ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
+                    )
+                    if maybe_event is not None:  # pragma: no branch
+                        yield maybe_event
+
+                # Handle the tool calls
+                for dtc in choice.delta.tool_calls or []:
+                    maybe_event = self._parts_manager.handle_tool_call_delta(
+                        vendor_part_id=dtc.index,
+                        tool_name=dtc.function and dtc.function.name,
+                        args=dtc.function and dtc.function.arguments,
+                        tool_call_id=dtc.id,
+                    )
+                    if maybe_event is not None:
+                        yield maybe_event
+        except APIError as e:
+            if isinstance(e.body, dict):  # pragma: no branch
+                # The Groq SDK tries to be helpful by raising an exception when generated tool arguments don't match the schema,
+                # but we'd rather handle it ourselves so we can tell the model to retry the tool call
+                try:
+                    error = _GroqToolUseFailedInnerError.model_validate(e.body)  # pyright: ignore[reportUnknownMemberType]
+                    yield self._parts_manager.handle_tool_call_part(
+                        vendor_part_id='tool_use_failed',
+                        tool_name=error.failed_generation.name,
+                        args=error.failed_generation.arguments,
+                    )
+                    return
+                except ValidationError as e:  # pragma: no cover
+                    pass
+            raise  # pragma: no cover
 
     @property
     def model_name(self) -> GroqModelName:
@@ -510,3 +570,31 @@ def _map_usage(completion: chat.ChatCompletionChunk | chat.ChatCompletion) -> us
         input_tokens=response_usage.prompt_tokens,
         output_tokens=response_usage.completion_tokens,
     )
+
+
+class _GroqToolUseFailedGeneration(BaseModel):
+    name: str
+    arguments: dict[str, Any]
+
+
+class _GroqToolUseFailedInnerError(BaseModel):
+    message: str
+    type: Literal['invalid_request_error']
+    code: Literal['tool_use_failed']
+    failed_generation: Json[_GroqToolUseFailedGeneration]
+
+
+class _GroqToolUseFailedError(BaseModel):
+    # The Groq SDK tries to be helpful by raising an exception when generated tool arguments don't match the schema,
+    # but we'd rather handle it ourselves so we can tell the model to retry the tool call.
+    # Example payload from `exception.body`:
+    # {
+    #     'error': {
+    #         'message': "Tool call validation failed: tool call validation failed: parameters for tool get_something_by_name did not match schema: errors: [missing properties: 'name', additionalProperties 'foo' not allowed]",
+    #         'type': 'invalid_request_error',
+    #         'code': 'tool_use_failed',
+    #         'failed_generation': '{"name": "get_something_by_name", "arguments": {\n  "foo": "bar"\n}}',
+    #     }
+    # }
+
+    error: _GroqToolUseFailedInnerError

pydantic_ai/models/huggingface.py

@@ -5,7 +5,7 @@ from collections.abc import AsyncIterable, AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import Any, Literal, Union, cast, overload
+from typing import Any, Literal, cast, overload
 
 from typing_extensions import assert_never
 
@@ -88,7 +88,7 @@ LatestHuggingFaceModelNames = Literal[
 """Latest Hugging Face models."""
 
 
-HuggingFaceModelName = Union[str, LatestHuggingFaceModelNames]
+HuggingFaceModelName = str | LatestHuggingFaceModelNames
 """Possible Hugging Face model names.
 
 You can browse available models [here](https://huggingface.co/models?pipeline_tag=text-generation&inference_provider=all&sort=trending).
@@ -267,11 +267,11 @@ class HuggingFaceModel(Model):
             for c in tool_calls:
                 items.append(ToolCallPart(c.function.name, c.function.arguments, tool_call_id=c.id))
         return ModelResponse(
-            items,
+            parts=items,
             usage=_map_usage(response),
             model_name=response.model,
             timestamp=timestamp,
-            provider_request_id=response.id,
+            provider_response_id=response.id,
             provider_name=self._provider.name,
         )
 
@@ -320,7 +320,7 @@ class HuggingFaceModel(Model):
                         # please open an issue. The below code is the code to send thinking to the provider.
                         # texts.append(f'<think>\n{item.content}\n</think>')
                         pass
-                    elif isinstance(item, (BuiltinToolCallPart, BuiltinToolReturnPart)):  # pragma: no cover
+                    elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
                         # This is currently never returned from huggingface
                         pass
                     else:

pydantic_ai/models/instrumented.py

@@ -2,10 +2,11 @@ from __future__ import annotations
 
 import itertools
 import json
-from collections.abc import AsyncIterator, Iterator, Mapping
+import warnings
+from collections.abc import AsyncIterator, Callable, Iterator, Mapping
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
-from typing import Any, Callable, Literal, cast
+from typing import Any, Literal, cast
 from urllib.parse import urlparse
 
 from opentelemetry._events import (
@@ -93,36 +94,41 @@ class InstrumentationSettings:
     def __init__(
         self,
         *,
-        event_mode: Literal['attributes', 'logs'] = 'attributes',
         tracer_provider: TracerProvider | None = None,
         meter_provider: MeterProvider | None = None,
-        event_logger_provider: EventLoggerProvider | None = None,
         include_binary_content: bool = True,
         include_content: bool = True,
-        version: Literal[1, 2] = 1,
+        version: Literal[1, 2] = 2,
+        event_mode: Literal['attributes', 'logs'] = 'attributes',
+        event_logger_provider: EventLoggerProvider | None = None,
     ):
         """Create instrumentation options.
 
         Args:
-            event_mode: The mode for emitting events. If `'attributes'`, events are attached to the span as attributes.
-                If `'logs'`, events are emitted as OpenTelemetry log-based events.
             tracer_provider: The OpenTelemetry tracer provider to use.
                 If not provided, the global tracer provider is used.
                 Calling `logfire.configure()` sets the global tracer provider, so most users don't need this.
             meter_provider: The OpenTelemetry meter provider to use.
                 If not provided, the global meter provider is used.
                 Calling `logfire.configure()` sets the global meter provider, so most users don't need this.
-            event_logger_provider: The OpenTelemetry event logger provider to use.
-                If not provided, the global event logger provider is used.
-                Calling `logfire.configure()` sets the global event logger provider, so most users don't need this.
-                This is only used if `event_mode='logs'`.
             include_binary_content: Whether to include binary content in the instrumentation events.
             include_content: Whether to include prompts, completions, and tool call arguments and responses
                 in the instrumentation events.
-            version: Version of the data format.
-                Version 1 is based on the legacy event-based OpenTelemetry GenAI spec.
-                Version 2 stores messages in the attributes `gen_ai.input.messages` and `gen_ai.output.messages`.
-                Version 2 is still WIP and experimental, but will become the default in Pydantic AI v1.
+            version: Version of the data format. This is unrelated to the Pydantic AI package version.
+                Version 1 is based on the legacy event-based OpenTelemetry GenAI spec
+                    and will be removed in a future release.
+                    The parameters `event_mode` and `event_logger_provider` are only relevant for version 1.
+                Version 2 uses the newer OpenTelemetry GenAI spec and stores messages in the following attributes:
+                    - `gen_ai.system_instructions` for instructions passed to the agent.
+                    - `gen_ai.input.messages` and `gen_ai.output.messages` on model request spans.
+                    - `pydantic_ai.all_messages` on agent run spans.
+            event_mode: The mode for emitting events in version 1.
+                If `'attributes'`, events are attached to the span as attributes.
+                If `'logs'`, events are emitted as OpenTelemetry log-based events.
+            event_logger_provider: The OpenTelemetry event logger provider to use.
+                If not provided, the global event logger provider is used.
+                Calling `logfire.configure()` sets the global event logger provider, so most users don't need this.
+                This is only used if `event_mode='logs'` and `version=1`.
         """
         from pydantic_ai import __version__
 
@@ -136,6 +142,14 @@ class InstrumentationSettings:
         self.event_mode = event_mode
         self.include_binary_content = include_binary_content
         self.include_content = include_content
+
+        if event_mode == 'logs' and version != 1:
+            warnings.warn(
+                'event_mode is only relevant for version=1 which is deprecated and will be removed in a future release.',
+                stacklevel=2,
+            )
+            version = 1
+
         self.version = version
 
         # As specified in the OpenTelemetry GenAI metrics spec:
@@ -236,27 +250,36 @@ class InstrumentationSettings:
             if response.provider_details and 'finish_reason' in response.provider_details:
                 output_message['finish_reason'] = response.provider_details['finish_reason']
             instructions = InstrumentedModel._get_instructions(input_messages)  # pyright: ignore [reportPrivateUsage]
+            system_instructions_attributes = self.system_instructions_attributes(instructions)
             attributes = {
                 'gen_ai.input.messages': json.dumps(self.messages_to_otel_messages(input_messages)),
                 'gen_ai.output.messages': json.dumps([output_message]),
+                **system_instructions_attributes,
                 'logfire.json_schema': json.dumps(
                     {
                         'type': 'object',
                         'properties': {
                             'gen_ai.input.messages': {'type': 'array'},
                             'gen_ai.output.messages': {'type': 'array'},
-                            **({'gen_ai.system_instructions': {'type': 'array'}} if instructions else {}),
+                            **(
+                                {'gen_ai.system_instructions': {'type': 'array'}}
+                                if system_instructions_attributes
+                                else {}
+                            ),
                             'model_request_parameters': {'type': 'object'},
                         },
                     }
                 ),
             }
-            if instructions is not None:
-                attributes['gen_ai.system_instructions'] = json.dumps(
-                    [_otel_messages.TextPart(type='text', content=instructions)]
-                )
             span.set_attributes(attributes)
 
+    def system_instructions_attributes(self, instructions: str | None) -> dict[str, str]:
+        if instructions and self.include_content:
+            return {
+                'gen_ai.system_instructions': json.dumps([_otel_messages.TextPart(type='text', content=instructions)]),
+            }
+        return {}
+
     def _emit_events(self, span: Span, events: list[Event]) -> None:
         if self.event_mode == 'logs':
             for event in events:
@@ -357,7 +380,7 @@ class InstrumentedModel(WrapperModel):
 
         if model_settings:
             for key in MODEL_SETTING_ATTRIBUTES:
-                if isinstance(value := model_settings.get(key), (float, int)):
+                if isinstance(value := model_settings.get(key), float | int):
                     attributes[f'gen_ai.request.{key}'] = value
 
         record_metrics: Callable[[], None] | None = None
@@ -397,10 +420,15 @@ class InstrumentedModel(WrapperModel):
                         return
 
                     self.instrumentation_settings.handle_messages(messages, response, system, span)
+                    try:
+                        cost_attributes = {'operation.cost': float(response.cost().total_price)}
+                    except LookupError:
+                        cost_attributes = {}
                     span.set_attributes(
                         {
                             **response.usage.opentelemetry_attributes(),
                             'gen_ai.response.model': response_model,
+                            **cost_attributes,
                         }
                     )
                     span.update_name(f'{operation} {request_model}')

pydantic_ai/models/mcp_sampling.py

@@ -2,7 +2,7 @@ from __future__ import annotations as _annotations
 
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
-from dataclasses import dataclass
+from dataclasses import KW_ONLY, dataclass
 from typing import TYPE_CHECKING, Any, cast
 
 from .. import _mcp, exceptions
@@ -36,6 +36,8 @@ class MCPSamplingModel(Model):
     session: ServerSession
     """The MCP server session to use for sampling."""
 
+    _: KW_ONLY
+
     default_max_tokens: int = 16_384
     """Default max tokens to use if not set in [`ModelSettings`][pydantic_ai.settings.ModelSettings.max_tokens].
 

pydantic_ai/models/mistral.py

@@ -5,7 +5,7 @@ from collections.abc import AsyncIterable, AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, Literal, Union, cast
+from typing import Any, Literal, cast
 
 import pydantic_core
 from httpx import Timeout
@@ -79,7 +79,7 @@ try:
     from mistralai.models.usermessage import UserMessage as MistralUserMessage
     from mistralai.types.basemodel import Unset as MistralUnset
     from mistralai.utils.eventstreaming import EventStreamAsync as MistralEventStreamAsync
-except ImportError as e:  # pragma: no cover
+except ImportError as e:
     raise ImportError(
         'Please install `mistral` to use the Mistral model, '
         'you can use the `mistral` optional group — `pip install "pydantic-ai-slim[mistral]"`'
@@ -90,7 +90,7 @@ LatestMistralModelNames = Literal[
 ]
 """Latest  Mistral models."""
 
-MistralModelName = Union[str, LatestMistralModelNames]
+MistralModelName = str | LatestMistralModelNames
 """Possible Mistral model names.
 
 Since Mistral supports a variety of date-stamped models, we explicitly list the most popular models but
@@ -117,7 +117,7 @@ class MistralModel(Model):
     """
 
     client: Mistral = field(repr=False)
-    json_mode_schema_prompt: str = """Answer in JSON Object, respect the format:\n```\n{schema}\n```\n"""
+    json_mode_schema_prompt: str
 
     _model_name: MistralModelName = field(repr=False)
     _provider: Provider[Mistral] = field(repr=False)
@@ -348,11 +348,11 @@ class MistralModel(Model):
                 parts.append(tool)
 
         return ModelResponse(
-            parts,
+            parts=parts,
             usage=_map_usage(response),
             model_name=response.model,
             timestamp=timestamp,
-            provider_request_id=response.id,
+            provider_response_id=response.id,
             provider_name=self._provider.name,
         )
 
@@ -515,7 +515,7 @@ class MistralModel(Model):
                         pass
                     elif isinstance(part, ToolCallPart):
                         tool_calls.append(self._map_tool_call(part))
-                    elif isinstance(part, (BuiltinToolCallPart, BuiltinToolReturnPart)):  # pragma: no cover
+                    elif isinstance(part, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
                         # This is currently never returned from mistral
                         pass
                     else:
@@ -576,7 +576,7 @@ class MistralModel(Model):
         return MistralUserMessage(content=content)
 
 
-MistralToolCallId = Union[str, None]
+MistralToolCallId = str | None
 
 
 @dataclass