add usage

Author: ajac-zero

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -4,7 +4,7 @@
 import itertools
 import json
 import warnings
-from collections.abc import AsyncIterable, AsyncIterator, Sequence
+from collections.abc import AsyncIterable, AsyncIterator, Iterable, Sequence
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field, replace
 from datetime import datetime
@@ -62,8 +62,9 @@
         ChatCompletionContentPartInputAudioParam,
         ChatCompletionContentPartParam,
         ChatCompletionContentPartTextParam,
+        chat_completion,
+        chat_completion_chunk,
     )
-    from openai.types.chat.chat_completion_chunk import Choice
     from openai.types.chat.chat_completion_content_part_image_param import ImageURL
     from openai.types.chat.chat_completion_content_part_input_audio_param import InputAudio
     from openai.types.chat.chat_completion_content_part_param import File, FileFile
@@ -543,28 +544,7 @@ def _process_provider_details(self, response: chat.ChatCompletion) -> dict[str,
 
         This method may be overridden by subclasses of `OpenAIChatModel` to apply custom mappings.
         """
-        choice = response.choices[0]
-        provider_details: dict[str, Any] = {}
-
-        # Add logprobs to vendor_details if available
-        if choice.logprobs is not None and choice.logprobs.content:
-            # Convert logprobs to a serializable format
-            provider_details['logprobs'] = [
-                {
-                    'token': lp.token,
-                    'bytes': lp.bytes,
-                    'logprob': lp.logprob,
-                    'top_logprobs': [
-                        {'token': tlp.token, 'bytes': tlp.bytes, 'logprob': tlp.logprob} for tlp in lp.top_logprobs
-                    ],
-                }
-                for lp in choice.logprobs.content
-            ]
-
-        raw_finish_reason = choice.finish_reason
-        provider_details['finish_reason'] = raw_finish_reason
-
-        return provider_details
+        return _map_provider_details(response.choices[0])
 
     def _process_response(self, response: chat.ChatCompletion | str) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
@@ -618,7 +598,7 @@ def _process_response(self, response: chat.ChatCompletion | str) -> ModelRespons
 
         return ModelResponse(
             parts=items,
-            usage=_map_usage(response, self._provider.name, self._provider.base_url, self._model_name),
+            usage=self._map_usage(response),
             model_name=response.model,
             timestamp=timestamp,
             provider_details=self._process_provider_details(response),
@@ -680,6 +660,9 @@ def _streamed_response_cls(self) -> type[OpenAIStreamedResponse]:
         """
         return OpenAIStreamedResponse
 
+    def _map_usage(self, response: chat.ChatCompletion) -> usage.RequestUsage:
+        return _map_usage(response, self._provider.name, self._provider.base_url, self._model_name)
+
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]:
         return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
@@ -1767,7 +1750,7 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             for event in self._map_part_delta(choice):
                 yield event
 
-    def _validate_response(self):
+    def _validate_response(self) -> AsyncIterable[ChatCompletionChunk]:
         """Hook that validates incoming chunks.
 
         This method may be overridden by subclasses of `OpenAIStreamedResponse` to apply custom chunk validations.
@@ -1776,7 +1759,7 @@ def _validate_response(self):
         """
         return self._response
 
-    def _map_part_delta(self, choice: Choice):
+    def _map_part_delta(self, choice: chat_completion_chunk.Choice) -> Iterable[ModelResponseStreamEvent]:
         """Hook that determines the sequence of mappings that will be called to produce events.
 
         This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the mapping.
@@ -1785,7 +1768,7 @@ def _map_part_delta(self, choice: Choice):
             self._map_thinking_delta(choice), self._map_text_delta(choice), self._map_tool_call_delta(choice)
         )
 
-    def _map_thinking_delta(self, choice: Choice):
+    def _map_thinking_delta(self, choice: chat_completion_chunk.Choice) -> Iterable[ModelResponseStreamEvent]:
         """Hook that maps thinking delta content to events.
 
         This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the mapping.
@@ -1811,7 +1794,7 @@ def _map_thinking_delta(self, choice: Choice):
                 provider_name=self.provider_name,
             )
 
-    def _map_text_delta(self, choice: Choice):
+    def _map_text_delta(self, choice: chat_completion_chunk.Choice) -> Iterable[ModelResponseStreamEvent]:
         """Hook that maps text delta content to events.
 
         This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the mapping.
@@ -1831,7 +1814,7 @@ def _map_text_delta(self, choice: Choice):
                     maybe_event.part.provider_name = self.provider_name
                 yield maybe_event
 
-    def _map_tool_call_delta(self, choice: Choice):
+    def _map_tool_call_delta(self, choice: chat_completion_chunk.Choice) -> Iterable[ModelResponseStreamEvent]:
         """Hook that maps tool call delta content to events.
 
         This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the mapping.
@@ -1851,11 +1834,9 @@ def _map_provider_details(self, chunk: ChatCompletionChunk) -> dict[str, str] |
 
         This method may be overridden by subclasses of `OpenAIStreamResponse` to customize the provider details.
         """
-        choice = chunk.choices[0]
-        if raw_finish_reason := choice.finish_reason:
-            return {'finish_reason': raw_finish_reason}
+        return _map_provider_details(chunk.choices[0])
 
-    def _map_usage(self, response: ChatCompletionChunk):
+    def _map_usage(self, response: ChatCompletionChunk) -> usage.RequestUsage:
         return _map_usage(response, self._provider_name, self._provider_url, self._model_name)
 
     def _map_finish_reason(
@@ -2177,7 +2158,7 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     UserWarning,
                 )
 
-    def _map_usage(self, response: responses.Response):
+    def _map_usage(self, response: responses.Response) -> usage.RequestUsage:
         return _map_usage(response, self._provider_name, self._provider_url, self._model_name)
 
     @property
@@ -2237,6 +2218,32 @@ def _map_usage(
     )
 
 
+def _map_provider_details(
+    choice: chat_completion_chunk.Choice | chat_completion.Choice,
+) -> dict[str, Any]:
+    provider_details: dict[str, Any] = {}
+
+    # Add logprobs to vendor_details if available
+    if choice.logprobs is not None and choice.logprobs.content:
+        # Convert logprobs to a serializable format
+        provider_details['logprobs'] = [
+            {
+                'token': lp.token,
+                'bytes': lp.bytes,
+                'logprob': lp.logprob,
+                'top_logprobs': [
+                    {'token': tlp.token, 'bytes': tlp.bytes, 'logprob': tlp.logprob} for tlp in lp.top_logprobs
+                ],
+            }
+            for lp in choice.logprobs.content
+        ]
+
+    if raw_finish_reason := choice.finish_reason:
+        provider_details['finish_reason'] = raw_finish_reason
+
+    return provider_details
+
+
 def _split_combined_tool_call_id(combined_id: str) -> tuple[str, str | None]:
     # When reasoning, the Responses API requires the `ResponseFunctionToolCall` to be returned with both the `call_id` and `id` fields.
     # Before our `ToolCallPart` gained the `id` field alongside `tool_call_id` field, we combined the two fields into a single string stored on `tool_call_id`.

pydantic_ai_slim/pydantic_ai/models/openrouter.py

@@ -24,7 +24,7 @@
 
 try:
     from openai import APIError
-    from openai.types import chat
+    from openai.types import chat, completion_usage
     from openai.types.chat import chat_completion, chat_completion_chunk
 
     from .openai import OpenAIChatModel, OpenAIChatModelSettings, OpenAIStreamedResponse
@@ -220,6 +220,12 @@ class WebPlugin(TypedDict, total=False):
 OpenRouterPlugin = WebPlugin
 
 
+class OpenRouterUsageConfig(TypedDict, total=False):
+    """Configuration for OpenRouter usage."""
+
+    include: bool
+
+
 class OpenRouterModelSettings(ModelSettings, total=False):
     """Settings used for an OpenRouter model request."""
 
@@ -254,6 +260,16 @@ class OpenRouterModelSettings(ModelSettings, total=False):
     """
 
     openrouter_plugins: list[OpenRouterPlugin]
+    """To enable plugins in the request.
+
+    Plugins are tools that can be used to extend the functionality of the model. [See more](https://openrouter.ai/docs/features/web-search)
+    """
+
+    openrouter_usage: OpenRouterUsageConfig
+    """To control the usage of the model.
+
+    The usage config object consolidates settings for enabling detailed usage information. [See more](https://openrouter.ai/docs/use-cases/usage-accounting)
+    """
 
 
 class OpenRouterError(BaseModel):
@@ -357,6 +373,30 @@ class OpenRouterChoice(chat_completion.Choice):
     """A wrapped chat completion message with OpenRouter specific attributes."""
 
 
+class OpenRouterCostDetails(BaseModel):
+    """OpenRouter specific cost details."""
+
+    upstream_inference_cost: int | None = None
+
+
+class OpenRouterCompletionTokenDetails(completion_usage.CompletionTokensDetails):
+    """Wraps OpenAI completion token details with OpenRouter specific attributes."""
+
+    image_tokens: int | None = None
+
+
+class OpenRouterUsage(completion_usage.CompletionUsage):
+    """Wraps OpenAI completion usage with OpenRouter specific attributes."""
+
+    cost: float | None = None
+
+    cost_details: OpenRouterCostDetails | None = None
+
+    is_byok: bool | None = None
+
+    completion_tokens_details: OpenRouterCompletionTokenDetails | None = None  # type: ignore[reportIncompatibleVariableOverride]
+
+
 class OpenRouterChatCompletion(chat.ChatCompletion):
     """Wraps OpenAI chat completion with OpenRouter specific attributes."""
 
@@ -369,6 +409,9 @@ class OpenRouterChatCompletion(chat.ChatCompletion):
     error: OpenRouterError | None = None
     """OpenRouter specific error attribute."""
 
+    usage: OpenRouterUsage | None = None  # type: ignore[reportIncompatibleVariableOverride]
+    """OpenRouter specific usage attribute."""
+
 
 def _openrouter_settings_to_openai_settings(model_settings: OpenRouterModelSettings) -> OpenAIChatModelSettings:
     """Transforms a 'OpenRouterModelSettings' object into an 'OpenAIChatModelSettings' object.
@@ -389,6 +432,8 @@ def _openrouter_settings_to_openai_settings(model_settings: OpenRouterModelSetti
         extra_body['preset'] = preset
     if transforms := model_settings.pop('openrouter_transforms', None):
         extra_body['transforms'] = transforms
+    if usage := model_settings.pop('openrouter_usage', None):
+        extra_body['usage'] = usage
 
     model_settings['extra_body'] = extra_body
 
@@ -401,30 +446,40 @@ def _map_usage(
     provider_url: str,
     model: str,
 ) -> RequestUsage:
+    assert isinstance(response, OpenRouterChatCompletion) or isinstance(response, OpenRouterChatCompletionChunk)
+    builder = RequestUsage()
+
     response_usage = response.usage
     if response_usage is None:
-        return RequestUsage()
-
-    usage_data = response_usage.model_dump(exclude_none=True)
-    details = {
-        k: v
-        for k, v in usage_data.items()
-        if k not in {'prompt_tokens', 'completion_tokens', 'input_tokens', 'output_tokens', 'total_tokens'}
-        if isinstance(v, int)
-    }
-    response_data = dict(model=model, usage=usage_data)
-
-    if response_usage.completion_tokens_details is not None:  # pragma: lax no cover
-        details.update(response_usage.completion_tokens_details.model_dump(exclude_none=True))
-
-    return RequestUsage.extract(
-        response_data,
-        provider=provider,
-        provider_url=provider_url,
-        provider_fallback='openai',
-        api_flavor='chat',
-        details=details,
-    )
+        return builder
+
+    builder.input_tokens = response_usage.prompt_tokens
+    builder.output_tokens = response_usage.completion_tokens
+
+    if prompt_token_details := response_usage.prompt_tokens_details:
+        if cached_tokens := prompt_token_details.cached_tokens:
+            builder.cache_read_tokens = cached_tokens
+
+        if audio_tokens := prompt_token_details.audio_tokens:  # pragma: lax no cover
+            builder.input_audio_tokens = audio_tokens
+
+        if video_tokens := prompt_token_details.video_tokens:  # pragma: lax no cover
+            builder.details['input_video_tokens'] = video_tokens
+
+    if completion_token_details := response_usage.completion_tokens_details:
+        if reasoning_tokens := completion_token_details.reasoning_tokens:
+            builder.details['reasoning_tokens'] = reasoning_tokens
+
+        if image_tokens := completion_token_details.image_tokens:  # pragma: lax no cover
+            builder.details['output_image_tokens'] = image_tokens
+
+    if (is_byok := response_usage.is_byok) is not None:
+        builder.details['is_byok'] = is_byok
+
+    if cost := response_usage.cost:
+        builder.details['cost'] = int(cost * 1000000)  # convert to microcost
+
+    return builder
 
 
 class OpenRouterModel(OpenAIChatModel):
@@ -524,6 +579,10 @@ def _map_model_response(self, message: ModelResponse) -> chat.ChatCompletionMess
     def _streamed_response_cls(self):
         return OpenRouterStreamedResponse
 
+    @override
+    def _map_usage(self, response: chat.ChatCompletion) -> RequestUsage:
+        return _map_usage(response, self._provider.name, self._provider.base_url, self._model_name)
+
     @override
     def _map_finish_reason(  # type: ignore[reportIncompatibleMethodOverride]
         self, key: Literal['stop', 'length', 'tool_calls', 'content_filter', 'error']
@@ -566,6 +625,9 @@ class OpenRouterChatCompletionChunk(chat.ChatCompletionChunk):
     choices: list[OpenRouterChunkChoice]  # type: ignore[reportIncompatibleVariableOverride]
     """A list of chat completion chunk choices modified with OpenRouter specific attributes."""
 
+    usage: OpenRouterUsage | None = None  # type: ignore[reportIncompatibleVariableOverride]
+    """Usage statistics for the completion request."""
+
 
 @dataclass
 class OpenRouterStreamedResponse(OpenAIStreamedResponse):