pydantic-ai-slim 0.7.4__py3-none-any.whl → 0.7.6__py3-none-any.whl

pydantic_ai/models/google.py

@@ -395,6 +395,7 @@ class GoogleModel(Model):
         return _process_response_from_parts(
             parts,
             response.model_version or self._model_name,
+            self._provider.name,
             usage,
             vendor_id=vendor_id,
             vendor_details=vendor_details,
@@ -414,6 +415,7 @@ class GoogleModel(Model):
             _model_name=self._model_name,
             _response=peekable_response,
             _timestamp=first_chunk.create_time or _utils.now_utc(),
+            _provider_name=self._provider.name,
         )
 
     async def _map_messages(self, messages: list[ModelMessage]) -> tuple[ContentDict | None, list[ContentUnionDict]]:
@@ -523,6 +525,7 @@ class GeminiStreamedResponse(StreamedResponse):
     _model_name: GoogleModelName
     _response: AsyncIterator[GenerateContentResponse]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
@@ -531,7 +534,10 @@ class GeminiStreamedResponse(StreamedResponse):
             assert chunk.candidates is not None
             candidate = chunk.candidates[0]
             if candidate.content is None or candidate.content.parts is None:
-                if candidate.finish_reason == 'SAFETY':  # pragma: no cover
+                if candidate.finish_reason == 'STOP':  # pragma: no cover
+                    # Normal completion - skip this chunk
+                    continue
+                elif candidate.finish_reason == 'SAFETY':  # pragma: no cover
                     raise UnexpectedModelBehavior('Safety settings triggered', str(chunk))
                 else:  # pragma: no cover
                     raise UnexpectedModelBehavior('Content field missing from streaming Gemini response', str(chunk))
@@ -561,6 +567,11 @@ class GeminiStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""
@@ -596,6 +607,7 @@ def _content_model_response(m: ModelResponse) -> ContentDict:
 def _process_response_from_parts(
     parts: list[Part],
     model_name: GoogleModelName,
+    provider_name: str,
     usage: usage.RequestUsage,
     vendor_id: str | None,
     vendor_details: dict[str, Any] | None = None,
@@ -633,7 +645,12 @@ def _process_response_from_parts(
                 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
             )
     return ModelResponse(
-        parts=items, model_name=model_name, usage=usage, provider_request_id=vendor_id, provider_details=vendor_details
+        parts=items,
+        model_name=model_name,
+        usage=usage,
+        provider_request_id=vendor_id,
+        provider_details=vendor_details,
+        provider_name=provider_name,
     )
 
 
@@ -661,7 +678,7 @@ def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
     if cached_content_token_count := metadata.cached_content_token_count:
         details['cached_content_tokens'] = cached_content_token_count
 
-    if thoughts_token_count := metadata.thoughts_token_count:
+    if thoughts_token_count := (metadata.thoughts_token_count or 0):
         details['thoughts_tokens'] = thoughts_token_count
 
     if tool_use_prompt_token_count := metadata.tool_use_prompt_token_count:
@@ -694,7 +711,7 @@ def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
 
     return usage.RequestUsage(
         input_tokens=metadata.prompt_token_count or 0,
-        output_tokens=metadata.candidates_token_count or 0,
+        output_tokens=(metadata.candidates_token_count or 0) + thoughts_token_count,
         cache_read_tokens=cached_content_token_count or 0,
         input_audio_tokens=input_audio_tokens,
         output_audio_tokens=output_audio_tokens,

pydantic_ai/models/groq.py

@@ -290,6 +290,7 @@ class GroqModel(Model):
             model_name=response.model,
             timestamp=timestamp,
             provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -309,6 +310,7 @@ class GroqModel(Model):
             _model_name=self._model_name,
             _model_profile=self.profile,
             _timestamp=number_to_datetime(first_chunk.created),
+            _provider_name=self._provider.name,
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]:
@@ -444,6 +446,7 @@ class GroqStreamedResponse(StreamedResponse):
     _model_profile: ModelProfile
     _response: AsyncIterable[chat.ChatCompletionChunk]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
@@ -482,6 +485,11 @@ class GroqStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai/models/huggingface.py

@@ -272,6 +272,7 @@ class HuggingFaceModel(Model):
             model_name=response.model,
             timestamp=timestamp,
             provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -291,6 +292,7 @@ class HuggingFaceModel(Model):
             _model_profile=self.profile,
             _response=peekable_response,
             _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
+            _provider_name=self._provider.name,
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ChatCompletionInputTool]:
@@ -437,6 +439,7 @@ class HuggingFaceStreamedResponse(StreamedResponse):
     _model_profile: ModelProfile
     _response: AsyncIterable[ChatCompletionStreamOutput]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
@@ -474,6 +477,11 @@ class HuggingFaceStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai/models/instrumented.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+import itertools
 import json
 from collections.abc import AsyncIterator, Iterator, Mapping
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
-from typing import Any, Callable, Literal
+from typing import Any, Callable, Literal, cast
 from urllib.parse import urlparse
 
 from opentelemetry._events import (
@@ -18,8 +19,14 @@ from opentelemetry.trace import Span, Tracer, TracerProvider, get_tracer_provide
 from opentelemetry.util.types import AttributeValue
 from pydantic import TypeAdapter
 
+from .. import _otel_messages
 from .._run_context import RunContext
-from ..messages import ModelMessage, ModelRequest, ModelResponse
+from ..messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    SystemPromptPart,
+)
 from ..settings import ModelSettings
 from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse
 from .wrapper import WrapperModel
@@ -80,6 +87,8 @@ class InstrumentationSettings:
     event_logger: EventLogger = field(repr=False)
     event_mode: Literal['attributes', 'logs'] = 'attributes'
     include_binary_content: bool = True
+    include_content: bool = True
+    version: Literal[1, 2] = 1
 
     def __init__(
         self,
@@ -90,6 +99,7 @@ class InstrumentationSettings:
         event_logger_provider: EventLoggerProvider | None = None,
         include_binary_content: bool = True,
         include_content: bool = True,
+        version: Literal[1, 2] = 1,
     ):
         """Create instrumentation options.
 
@@ -109,6 +119,10 @@ class InstrumentationSettings:
             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.
         """
         from pydantic_ai import __version__
 
@@ -122,6 +136,7 @@ class InstrumentationSettings:
         self.event_mode = event_mode
         self.include_binary_content = include_binary_content
         self.include_content = include_content
+        self.version = version
 
         # As specified in the OpenTelemetry GenAI metrics spec:
         # https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/#metric-gen_aiclienttokenusage
@@ -179,6 +194,90 @@ class InstrumentationSettings:
             event.body = InstrumentedModel.serialize_any(event.body)
         return events
 
+    def messages_to_otel_messages(self, messages: list[ModelMessage]) -> list[_otel_messages.ChatMessage]:
+        result: list[_otel_messages.ChatMessage] = []
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                for is_system, group in itertools.groupby(message.parts, key=lambda p: isinstance(p, SystemPromptPart)):
+                    message_parts: list[_otel_messages.MessagePart] = []
+                    for part in group:
+                        if hasattr(part, 'otel_message_parts'):
+                            message_parts.extend(part.otel_message_parts(self))
+                    result.append(
+                        _otel_messages.ChatMessage(role='system' if is_system else 'user', parts=message_parts)
+                    )
+            elif isinstance(message, ModelResponse):  # pragma: no branch
+                result.append(_otel_messages.ChatMessage(role='assistant', parts=message.otel_message_parts(self)))
+        return result
+
+    def handle_messages(self, input_messages: list[ModelMessage], response: ModelResponse, system: str, span: Span):
+        if self.version == 1:
+            events = self.messages_to_otel_events(input_messages)
+            for event in self.messages_to_otel_events([response]):
+                events.append(
+                    Event(
+                        'gen_ai.choice',
+                        body={
+                            'index': 0,
+                            'message': event.body,
+                        },
+                    )
+                )
+            for event in events:
+                event.attributes = {
+                    GEN_AI_SYSTEM_ATTRIBUTE: system,
+                    **(event.attributes or {}),
+                }
+            self._emit_events(span, events)
+        else:
+            output_messages = self.messages_to_otel_messages([response])
+            assert len(output_messages) == 1
+            output_message = cast(_otel_messages.OutputMessage, output_messages[0])
+            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]
+            attributes = {
+                'gen_ai.input.messages': json.dumps(self.messages_to_otel_messages(input_messages)),
+                'gen_ai.output.messages': json.dumps([output_message]),
+                '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 {}),
+                            '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 _emit_events(self, span: Span, events: list[Event]) -> None:
+        if self.event_mode == 'logs':
+            for event in events:
+                self.event_logger.emit(event)
+        else:
+            attr_name = 'events'
+            span.set_attributes(
+                {
+                    attr_name: json.dumps([InstrumentedModel.event_to_dict(event) for event in events]),
+                    'logfire.json_schema': json.dumps(
+                        {
+                            'type': 'object',
+                            'properties': {
+                                attr_name: {'type': 'array'},
+                                'model_request_parameters': {'type': 'object'},
+                            },
+                        }
+                    ),
+                }
+            )
+
 
 GEN_AI_SYSTEM_ATTRIBUTE = 'gen_ai.system'
 GEN_AI_REQUEST_MODEL_ATTRIBUTE = 'gen_ai.request.model'
@@ -269,7 +368,7 @@ class InstrumentedModel(WrapperModel):
                     # FallbackModel updates these span attributes.
                     attributes.update(getattr(span, 'attributes', {}))
                     request_model = attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE]
-                    system = attributes[GEN_AI_SYSTEM_ATTRIBUTE]
+                    system = cast(str, attributes[GEN_AI_SYSTEM_ATTRIBUTE])
 
                     response_model = response.model_name or request_model
 
@@ -297,18 +396,7 @@ class InstrumentedModel(WrapperModel):
                     if not span.is_recording():
                         return
 
-                    events = self.instrumentation_settings.messages_to_otel_events(messages)
-                    for event in self.instrumentation_settings.messages_to_otel_events([response]):
-                        events.append(
-                            Event(
-                                'gen_ai.choice',
-                                body={
-                                    # TODO finish_reason
-                                    'index': 0,
-                                    'message': event.body,
-                                },
-                            )
-                        )
+                    self.instrumentation_settings.handle_messages(messages, response, system, span)
                     span.set_attributes(
                         {
                             **response.usage.opentelemetry_attributes(),
@@ -316,12 +404,6 @@ class InstrumentedModel(WrapperModel):
                         }
                     )
                     span.update_name(f'{operation} {request_model}')
-                    for event in events:
-                        event.attributes = {
-                            GEN_AI_SYSTEM_ATTRIBUTE: system,
-                            **(event.attributes or {}),
-                        }
-                    self._emit_events(span, events)
 
                 yield finish
         finally:
@@ -330,27 +412,6 @@ class InstrumentedModel(WrapperModel):
                 # to prevent them from being redundantly recorded in the span itself by logfire.
                 record_metrics()
 
-    def _emit_events(self, span: Span, events: list[Event]) -> None:
-        if self.instrumentation_settings.event_mode == 'logs':
-            for event in events:
-                self.instrumentation_settings.event_logger.emit(event)
-        else:
-            attr_name = 'events'
-            span.set_attributes(
-                {
-                    attr_name: json.dumps([self.event_to_dict(event) for event in events]),
-                    'logfire.json_schema': json.dumps(
-                        {
-                            'type': 'object',
-                            'properties': {
-                                attr_name: {'type': 'array'},
-                                'model_request_parameters': {'type': 'object'},
-                            },
-                        }
-                    ),
-                }
-            )
-
     @staticmethod
     def model_attributes(model: Model):
         attributes: dict[str, AttributeValue] = {

pydantic_ai/models/mistral.py

@@ -353,6 +353,7 @@ class MistralModel(Model):
             model_name=response.model,
             timestamp=timestamp,
             provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -378,6 +379,7 @@ class MistralModel(Model):
             _response=peekable_response,
             _model_name=self._model_name,
             _timestamp=timestamp,
+            _provider_name=self._provider.name,
         )
 
     @staticmethod
@@ -584,6 +586,7 @@ class MistralStreamedResponse(StreamedResponse):
     _model_name: MistralModelName
     _response: AsyncIterable[MistralCompletionEvent]
     _timestamp: datetime
+    _provider_name: str
 
     _delta_content: str = field(default='', init=False)
 
@@ -631,6 +634,11 @@ class MistralStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""