pydantic-ai-slim 0.7.4__tar.gz → 0.7.6__tar.gz

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.7.4
+Version: 0.7.6
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Project-URL: Homepage, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
 Project-URL: Source, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
@@ -35,7 +35,7 @@ Requires-Dist: genai-prices>=0.0.22
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.7.4
+Requires-Dist: pydantic-graph==0.7.6
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -57,9 +57,9 @@ Requires-Dist: cohere>=5.16.0; (platform_system != 'Emscripten') and extra == 'c
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.7.4; extra == 'evals'
+Requires-Dist: pydantic-evals==0.7.6; extra == 'evals'
 Provides-Extra: google
-Requires-Dist: google-genai>=1.28.0; extra == 'google'
+Requires-Dist: google-genai>=1.31.0; extra == 'google'
 Provides-Extra: groq
 Requires-Dist: groq>=0.25.0; extra == 'groq'
 Provides-Extra: huggingface

pydantic_ai_slim-0.7.6/pydantic_ai/_otel_messages.py

@@ -0,0 +1,67 @@
+"""Type definitions of OpenTelemetry GenAI spec message parts.
+
+Based on https://github.com/lmolkova/semantic-conventions/blob/eccd1f806e426a32c98271c3ce77585492d26de2/docs/gen-ai/non-normative/models.ipynb
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import JsonValue
+from typing_extensions import NotRequired, TypeAlias, TypedDict
+
+
+class TextPart(TypedDict):
+    type: Literal['text']
+    content: NotRequired[str]
+
+
+class ToolCallPart(TypedDict):
+    type: Literal['tool_call']
+    id: str
+    name: str
+    arguments: NotRequired[JsonValue]
+
+
+class ToolCallResponsePart(TypedDict):
+    type: Literal['tool_call_response']
+    id: str
+    name: str
+    result: NotRequired[JsonValue]
+
+
+class MediaUrlPart(TypedDict):
+    type: Literal['image-url', 'audio-url', 'video-url', 'document-url']
+    url: NotRequired[str]
+
+
+class BinaryDataPart(TypedDict):
+    type: Literal['binary']
+    media_type: str
+    content: NotRequired[str]
+
+
+class ThinkingPart(TypedDict):
+    type: Literal['thinking']
+    content: NotRequired[str]
+
+
+MessagePart: TypeAlias = 'TextPart | ToolCallPart | ToolCallResponsePart | MediaUrlPart | BinaryDataPart | ThinkingPart'
+
+
+Role = Literal['system', 'user', 'assistant']
+
+
+class ChatMessage(TypedDict):
+    role: Role
+    parts: list[MessagePart]
+
+
+InputMessages: TypeAlias = list[ChatMessage]
+
+
+class OutputMessage(ChatMessage):
+    finish_reason: NotRequired[str]
+
+
+OutputMessages: TypeAlias = list[OutputMessage]

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/agent/__init__.py

@@ -677,16 +677,23 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
     def _run_span_end_attributes(
         self, state: _agent_graph.GraphAgentState, usage: _usage.RunUsage, settings: InstrumentationSettings
     ):
+        if settings.version == 1:
+            attr_name = 'all_messages_events'
+            value = [
+                InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(state.message_history)
+            ]
+        else:
+            attr_name = 'pydantic_ai.all_messages'
+            value = settings.messages_to_otel_messages(state.message_history)
+
         return {
             **usage.opentelemetry_attributes(),
-            'all_messages_events': json.dumps(
-                [InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(state.message_history)]
-            ),
+            attr_name: json.dumps(value),
             'logfire.json_schema': json.dumps(
                 {
                     'type': 'object',
                     'properties': {
-                        'all_messages_events': {'type': 'array'},
+                        attr_name: {'type': 'array'},
                         'final_result': {'type': 'object'},
                     },
                 }

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/builtin_tools.py

@@ -29,6 +29,7 @@ class WebSearchTool(AbstractBuiltinTool):
     * Anthropic
     * OpenAI
     * Groq
+    * Google
     """
 
     search_context_size: Literal['low', 'medium', 'high'] = 'medium'

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/durable_exec/temporal/_model.py

@@ -55,6 +55,10 @@ class TemporalStreamedResponse(StreamedResponse):
     def model_name(self) -> str:
         return self.response.model_name or ''  # pragma: no cover
 
+    @property
+    def provider_name(self) -> str:
+        return self.response.provider_name or ''  # pragma: no cover
+
     @property
     def timestamp(self) -> datetime:
         return self.response.timestamp  # pragma: no cover

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/messages.py

@@ -10,10 +10,11 @@ from typing import TYPE_CHECKING, Annotated, Any, Literal, Union, cast, overload
 
 import pydantic
 import pydantic_core
+from genai_prices import calc_price, types as genai_types
 from opentelemetry._events import Event  # pyright: ignore[reportPrivateImportUsage]
 from typing_extensions import TypeAlias, deprecated
 
-from . import _utils
+from . import _otel_messages, _utils
 from ._utils import (
     generate_tool_call_id as _generate_tool_call_id,
     now_utc as _now_utc,
@@ -82,6 +83,9 @@ class SystemPromptPart:
             body={'role': 'system', **({'content': self.content} if settings.include_content else {})},
         )
 
+    def otel_message_parts(self, settings: InstrumentationSettings) -> list[_otel_messages.MessagePart]:
+        return [_otel_messages.TextPart(type='text', **{'content': self.content} if settings.include_content else {})]
+
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
@@ -504,25 +508,41 @@ class UserPromptPart:
     """Part type identifier, this is available on all parts as a discriminator."""
 
     def otel_event(self, settings: InstrumentationSettings) -> Event:
-        content: str | list[dict[str, Any] | str] | dict[str, Any]
-        if isinstance(self.content, str):
-            content = self.content if settings.include_content else {'kind': 'text'}
-        else:
-            content = []
-            for part in self.content:
-                if isinstance(part, str):
-                    content.append(part if settings.include_content else {'kind': 'text'})
-                elif isinstance(part, (ImageUrl, AudioUrl, DocumentUrl, VideoUrl)):
-                    content.append({'kind': part.kind, **({'url': part.url} if settings.include_content else {})})
-                elif isinstance(part, BinaryContent):
-                    converted_part = {'kind': part.kind, 'media_type': part.media_type}
-                    if settings.include_content and settings.include_binary_content:
-                        converted_part['binary_content'] = base64.b64encode(part.data).decode()
-                    content.append(converted_part)
-                else:
-                    content.append({'kind': part.kind})  # pragma: no cover
+        content = [{'kind': part.pop('type'), **part} for part in self.otel_message_parts(settings)]
+        for part in content:
+            if part['kind'] == 'binary' and 'content' in part:
+                part['binary_content'] = part.pop('content')
+        content = [
+            part['content'] if part == {'kind': 'text', 'content': part.get('content')} else part for part in content
+        ]
+        if content in ([{'kind': 'text'}], [self.content]):
+            content = content[0]
         return Event('gen_ai.user.message', body={'content': content, 'role': 'user'})
 
+    def otel_message_parts(self, settings: InstrumentationSettings) -> list[_otel_messages.MessagePart]:
+        parts: list[_otel_messages.MessagePart] = []
+        content: Sequence[UserContent] = [self.content] if isinstance(self.content, str) else self.content
+        for part in content:
+            if isinstance(part, str):
+                parts.append(
+                    _otel_messages.TextPart(type='text', **({'content': part} if settings.include_content else {}))
+                )
+            elif isinstance(part, (ImageUrl, AudioUrl, DocumentUrl, VideoUrl)):
+                parts.append(
+                    _otel_messages.MediaUrlPart(
+                        type=part.kind,
+                        **{'url': part.url} if settings.include_content else {},
+                    )
+                )
+            elif isinstance(part, BinaryContent):
+                converted_part = _otel_messages.BinaryDataPart(type='binary', media_type=part.media_type)
+                if settings.include_content and settings.include_binary_content:
+                    converted_part['content'] = base64.b64encode(part.data).decode()
+                parts.append(converted_part)
+            else:
+                parts.append({'type': part.kind})  # pragma: no cover
+        return parts
+
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
@@ -576,6 +596,18 @@ class BaseToolReturnPart:
             },
         )
 
+    def otel_message_parts(self, settings: InstrumentationSettings) -> list[_otel_messages.MessagePart]:
+        from .models.instrumented import InstrumentedModel
+
+        return [
+            _otel_messages.ToolCallResponsePart(
+                type='tool_call_response',
+                id=self.tool_call_id,
+                name=self.tool_name,
+                **({'result': InstrumentedModel.serialize_any(self.content)} if settings.include_content else {}),
+            )
+        ]
+
     def has_content(self) -> bool:
         """Return `True` if the tool return has content."""
         return self.content is not None  # pragma: no cover
@@ -669,6 +701,19 @@ class RetryPromptPart:
                 },
             )
 
+    def otel_message_parts(self, settings: InstrumentationSettings) -> list[_otel_messages.MessagePart]:
+        if self.tool_name is None:
+            return [_otel_messages.TextPart(type='text', content=self.model_response())]
+        else:
+            return [
+                _otel_messages.ToolCallResponsePart(
+                    type='tool_call_response',
+                    id=self.tool_call_id,
+                    name=self.tool_name,
+                    **({'result': self.model_response()} if settings.include_content else {}),
+                )
+            ]
+
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
@@ -848,6 +893,9 @@ class ModelResponse:
     kind: Literal['response'] = 'response'
     """Message type identifier, this is available on all parts as a discriminator."""
 
+    provider_name: str | None = None
+    """The name of the LLM provider that generated the response."""
+
     provider_details: dict[str, Any] | None = field(default=None)
     """Additional provider-specific details in a serializable format.
 
@@ -858,6 +906,19 @@ class ModelResponse:
     provider_request_id: str | None = None
     """request ID as specified by the model provider. This can be used to track the specific request to the model."""
 
+    def price(self) -> genai_types.PriceCalculation:
+        """Calculate the price of the usage.
+
+        Uses [`genai-prices`](https://github.com/pydantic/genai-prices).
+        """
+        assert self.model_name, 'Model name is required to calculate price'
+        return calc_price(
+            self.usage,
+            self.model_name,
+            provider_id=self.provider_name,
+            genai_request_timestamp=self.timestamp,
+        )
+
     def otel_events(self, settings: InstrumentationSettings) -> list[Event]:
         """Return OpenTelemetry events for the response."""
         result: list[Event] = []
@@ -894,6 +955,36 @@ class ModelResponse:
 
         return result
 
+    def otel_message_parts(self, settings: InstrumentationSettings) -> list[_otel_messages.MessagePart]:
+        parts: list[_otel_messages.MessagePart] = []
+        for part in self.parts:
+            if isinstance(part, TextPart):
+                parts.append(
+                    _otel_messages.TextPart(
+                        type='text',
+                        **({'content': part.content} if settings.include_content else {}),
+                    )
+                )
+            elif isinstance(part, ThinkingPart):
+                parts.append(
+                    _otel_messages.ThinkingPart(
+                        type='thinking',
+                        **({'content': part.content} if settings.include_content else {}),
+                    )
+                )
+            elif isinstance(part, ToolCallPart):
+                call_part = _otel_messages.ToolCallPart(type='tool_call', id=part.tool_call_id, name=part.tool_name)
+                if settings.include_content and part.args is not None:
+                    from .models.instrumented import InstrumentedModel
+
+                    if isinstance(part.args, str):
+                        call_part['arguments'] = part.args
+                    else:
+                        call_part['arguments'] = {k: InstrumentedModel.serialize_any(v) for k, v in part.args.items()}
+
+                parts.append(call_part)
+        return parts
+
     @property
     @deprecated('`vendor_details` is deprecated, use `provider_details` instead')
     def vendor_details(self) -> dict[str, Any] | None:

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/models/__init__.py

@@ -111,6 +111,15 @@ KnownModelName = TypeAliasType(
         'bedrock:mistral.mixtral-8x7b-instruct-v0:1',
         'bedrock:mistral.mistral-large-2402-v1:0',
         'bedrock:mistral.mistral-large-2407-v1:0',
+        'cerebras:gpt-oss-120b',
+        'cerebras:llama3.1-8b',
+        'cerebras:llama-3.3-70b',
+        'cerebras:llama-4-scout-17b-16e-instruct',
+        'cerebras:llama-4-maverick-17b-128e-instruct',
+        'cerebras:qwen-3-235b-a22b-instruct-2507',
+        'cerebras:qwen-3-32b',
+        'cerebras:qwen-3-coder-480b',
+        'cerebras:qwen-3-235b-a22b-thinking-2507',
         'claude-3-5-haiku-20241022',
         'claude-3-5-haiku-latest',
         'claude-3-5-sonnet-20240620',
@@ -598,6 +607,7 @@ class StreamedResponse(ABC):
             model_name=self.model_name,
             timestamp=self.timestamp,
             usage=self.usage(),
+            provider_name=self.provider_name,
         )
 
     def usage(self) -> RequestUsage:
@@ -610,6 +620,12 @@ class StreamedResponse(ABC):
         """Get the model name of the response."""
         raise NotImplementedError()
 
+    @property
+    @abstractmethod
+    def provider_name(self) -> str | None:
+        """Get the provider name."""
+        raise NotImplementedError()
+
     @property
     @abstractmethod
     def timestamp(self) -> datetime:
@@ -688,21 +704,23 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
 
         return CohereModel(model_name, provider=provider)
     elif provider in (
-        'openai',
-        'deepseek',
         'azure',
-        'openrouter',
-        'vercel',
+        'deepseek',
+        'cerebras',
+        'fireworks',
+        'github',
         'grok',
+        'heroku',
         'moonshotai',
-        'fireworks',
+        'openai',
+        'openai-chat',
+        'openrouter',
         'together',
-        'heroku',
-        'github',
+        'vercel',
     ):
-        from .openai import OpenAIModel
+        from .openai import OpenAIChatModel
 
-        return OpenAIModel(model_name, provider=provider)
+        return OpenAIChatModel(model_name, provider=provider)
     elif provider == 'openai-responses':
         from .openai import OpenAIResponsesModel
 

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/models/anthropic.py

@@ -37,12 +37,13 @@ from ..messages import (
 )
 from ..profiles import ModelProfileSpec
 from ..providers import Provider, infer_provider
+from ..providers.anthropic import AsyncAnthropicClient
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, download_item, get_user_agent
 
 try:
-    from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropic, AsyncStream
+    from anthropic import NOT_GIVEN, APIStatusError, AsyncStream
     from anthropic.types.beta import (
         BetaBase64PDFBlockParam,
         BetaBase64PDFSourceParam,
@@ -134,16 +135,16 @@ class AnthropicModel(Model):
     Apart from `__init__`, all methods are private or match those of the base class.
     """
 
-    client: AsyncAnthropic = field(repr=False)
+    client: AsyncAnthropicClient = field(repr=False)
 
     _model_name: AnthropicModelName = field(repr=False)
-    _provider: Provider[AsyncAnthropic] = field(repr=False)
+    _provider: Provider[AsyncAnthropicClient] = field(repr=False)
 
     def __init__(
         self,
         model_name: AnthropicModelName,
         *,
-        provider: Literal['anthropic'] | Provider[AsyncAnthropic] = 'anthropic',
+        provider: Literal['anthropic'] | Provider[AsyncAnthropicClient] = 'anthropic',
         profile: ModelProfileSpec | None = None,
         settings: ModelSettings | None = None,
     ):
@@ -153,7 +154,7 @@ class AnthropicModel(Model):
             model_name: The name of the Anthropic model to use. List of model names available
                 [here](https://docs.anthropic.com/en/docs/about-claude/models).
             provider: The provider to use for the Anthropic API. Can be either the string 'anthropic' or an
-                instance of `Provider[AsyncAnthropic]`. If not provided, the other parameters will be used.
+                instance of `Provider[AsyncAnthropicClient]`. If not provided, the other parameters will be used.
             profile: The model profile to use. Defaults to a profile picked by the provider based on the model name.
             settings: Default model settings for this model instance.
         """
@@ -326,7 +327,11 @@ class AnthropicModel(Model):
                 )
 
         return ModelResponse(
-            items, usage=_map_usage(response), model_name=response.model, provider_request_id=response.id
+            items,
+            usage=_map_usage(response),
+            model_name=response.model,
+            provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -344,6 +349,7 @@ class AnthropicModel(Model):
             _model_name=self._model_name,
             _response=peekable_response,
             _timestamp=timestamp,
+            _provider_name=self._provider.name,
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[BetaToolParam]:
@@ -469,9 +475,9 @@ class AnthropicModel(Model):
                     anthropic_messages.append(BetaMessageParam(role='assistant', content=assistant_content_params))
             else:
                 assert_never(m)
-        system_prompt = '\n\n'.join(system_prompt_parts)
         if instructions := self._get_instructions(messages):
-            system_prompt = f'{instructions}\n\n{system_prompt}'
+            system_prompt_parts.insert(0, instructions)
+        system_prompt = '\n\n'.join(system_prompt_parts)
         return system_prompt, anthropic_messages
 
     @staticmethod
@@ -574,6 +580,7 @@ class AnthropicStreamedResponse(StreamedResponse):
     _model_name: AnthropicModelName
     _response: AsyncIterable[BetaRawMessageStreamEvent]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         current_block: BetaContentBlock | None = None
@@ -655,6 +662,11 @@ class AnthropicStreamedResponse(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_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/models/bedrock.py

@@ -240,10 +240,7 @@ class BedrockConverseModel(Model):
 
     @staticmethod
     def _map_tool_definition(f: ToolDefinition) -> ToolTypeDef:
-        tool_spec: ToolSpecificationTypeDef = {
-            'name': f.name,
-            'inputSchema': {'json': f.parameters_json_schema},
-        }
+        tool_spec: ToolSpecificationTypeDef = {'name': f.name, 'inputSchema': {'json': f.parameters_json_schema}}
 
         if f.description:  # pragma: no branch
             tool_spec['description'] = f.description
@@ -275,6 +272,7 @@ class BedrockConverseModel(Model):
             model_request_parameters=model_request_parameters,
             _model_name=self.model_name,
             _event_stream=response,
+            _provider_name=self._provider.name,
         )
 
     async def _process_response(self, response: ConverseResponseTypeDef) -> ModelResponse:
@@ -304,7 +302,9 @@ class BedrockConverseModel(Model):
             output_tokens=response['usage']['outputTokens'],
         )
         vendor_id = response.get('ResponseMetadata', {}).get('RequestId', None)
-        return ModelResponse(items, usage=u, model_name=self.model_name, provider_request_id=vendor_id)
+        return ModelResponse(
+            items, usage=u, model_name=self.model_name, provider_request_id=vendor_id, provider_name=self._provider.name
+        )
 
     @overload
     async def _messages_create(
@@ -423,7 +423,7 @@ class BedrockConverseModel(Model):
         for message in messages:
             if isinstance(message, ModelRequest):
                 for part in message.parts:
-                    if isinstance(part, SystemPromptPart):
+                    if isinstance(part, SystemPromptPart) and part.content:
                         system_prompt.append({'text': part.content})
                     elif isinstance(part, UserPromptPart):
                         bedrock_messages.extend(await self._map_user_prompt(part, document_count))
@@ -594,6 +594,7 @@ class BedrockStreamedResponse(StreamedResponse):
 
     _model_name: BedrockModelName
     _event_stream: EventStream[ConverseStreamOutputTypeDef]
+    _provider_name: str
     _timestamp: datetime = field(default_factory=_utils.now_utc)
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
@@ -660,15 +661,20 @@ class BedrockStreamedResponse(StreamedResponse):
                     if maybe_event:  # pragma: no branch
                         yield maybe_event
 
-    @property
-    def timestamp(self) -> datetime:
-        return self._timestamp
-
     @property
     def model_name(self) -> str:
         """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:
+        return self._timestamp
+
     def _map_usage(self, metadata: ConverseStreamMetadataEventTypeDef) -> usage.RequestUsage:
         return usage.RequestUsage(
             input_tokens=metadata['usage']['inputTokens'],

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/models/cohere.py

@@ -205,7 +205,9 @@ class CohereModel(Model):
                         tool_call_id=c.id or _generate_tool_call_id(),
                     )
                 )
-        return ModelResponse(parts=parts, usage=_map_usage(response), model_name=self._model_name)
+        return ModelResponse(
+            parts=parts, usage=_map_usage(response), model_name=self._model_name, provider_name=self._provider.name
+        )
 
     def _map_messages(self, messages: list[ModelMessage]) -> list[ChatMessageV2]:
         """Just maps a `pydantic_ai.Message` to a `cohere.ChatMessageV2`."""

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/models/function.py

@@ -304,6 +304,11 @@ class FunctionStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> None:
+        """Get the provider name."""
+        return None
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

{pydantic_ai_slim-0.7.4 → pydantic_ai_slim-0.7.6}/pydantic_ai/models/gemini.py

@@ -305,6 +305,7 @@ class GeminiModel(Model):
             _model_name=self._model_name,
             _content=content,
             _stream=aiter_bytes,
+            _provider_name=self._provider.name,
         )
 
     async def _message_to_gemini_content(
@@ -425,6 +426,7 @@ class GeminiStreamedResponse(StreamedResponse):
     _model_name: GeminiModelName
     _content: bytearray
     _stream: AsyncIterator[bytes]
+    _provider_name: str
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
@@ -495,6 +497,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."""
@@ -883,7 +890,7 @@ def _metadata_as_usage(response: _GeminiResponse) -> usage.RequestUsage:
 
     return usage.RequestUsage(
         input_tokens=metadata.get('prompt_token_count', 0),
-        output_tokens=metadata.get('candidates_token_count', 0),
+        output_tokens=metadata.get('candidates_token_count', 0) + thoughts_token_count,
         cache_read_tokens=cached_content_token_count,
         input_audio_tokens=input_audio_tokens,
         output_audio_tokens=output_audio_tokens,