openai[patch]: add usage metadata details

libs/core/langchain_core/messages/ai.py

@@ -2,7 +2,7 @@
 from typing import Any, Literal, Optional, Union
 
 from pydantic import model_validator
-from typing_extensions import Self, TypedDict
+from typing_extensions import NotRequired, Self, TypedDict
 
 from langchain_core.messages.base import (
     BaseMessage,
@@ -29,6 +29,66 @@
 from langchain_core.utils.json import parse_partial_json
 
 
+class InputTokenDetails(TypedDict, total=False):
+    """Breakdown of input token counts.
+
+    Does *not* need to sum to full input token count. Does *not* need to have all keys.
+
+    Example:
+
+        .. code-block:: python
+
+            {
+                "audio": 10,
+                "cache_creation": 200,
+                "cache_read": 100,
+            }
+
+    .. versionadded:: 0.3.9
+    """
+
+    audio: int
+    """Audio input tokens."""
+    cache_creation: int
+    """Input tokens that were cached and there was a cache miss.
+
+    Since there was a cache miss, the cache was created from these tokens.
+    """
+    cache_read: int
+    """Input tokens that were cached and there was a cache hit.
+
+    Since there was a cache hit, the tokens were read from the cache. More precisely,
+    the model state given these tokens was read from the cache.
+    """
+
+
+class OutputTokenDetails(TypedDict, total=False):
+    """Breakdown of output token counts.
+
+    Does *not* need to sum to full output token count. Does *not* need to have all keys.
+
+    Example:
+
+        .. code-block:: python
+
+            {
+                "audio": 10,
+                "reasoning": 200,
+            }
+
+    .. versionadded:: 0.3.9
+    """
+
+    audio: int
+    """Audio output tokens."""
+    reasoning: int
+    """Reasoning output tokens.
+
+    Tokens generated by the model in a chain of thought process (i.e. by OpenAI's o1
+    models) that are not returned as part of model output.
+    """
+
+
 class UsageMetadata(TypedDict):
     """Usage metadata for a message, such as token counts.
 
@@ -39,18 +99,41 @@ class UsageMetadata(TypedDict):
         .. code-block:: python
 
             {
-                "input_tokens": 10,
-                "output_tokens": 20,
-                "total_tokens": 30
+                "input_tokens": 350,
+                "output_tokens": 240,
+                "total_tokens": 590,
+                "input_token_details": {
+                    "audio": 10,
+                    "cache_creation": 200,
+                    "cache_read": 100,
+                },
+                "output_token_details": {
+                    "audio": 10,
+                    "reasoning": 200,
+                }
             }
+
+    .. versionchanged:: 0.3.9
+
+        Added ``input_token_details`` and ``output_token_details``.
     """
 
     input_tokens: int
-    """Count of input (or prompt) tokens."""
+    """Count of input (or prompt) tokens. Sum of all input token types."""
     output_tokens: int
-    """Count of output (or completion) tokens."""
+    """Count of output (or completion) tokens. Sum of all output token types."""
     total_tokens: int
-    """Total token count."""
+    """Total token count. Sum of input_tokens + output_tokens."""
+    input_token_details: NotRequired[InputTokenDetails]
+    """Breakdown of input token counts.
+
+    Does *not* need to sum to full input token count. Does *not* need to have all keys.
+    """
+    output_token_details: NotRequired[OutputTokenDetails]
+    """Breakdown of output token counts.
+
+    Does *not* need to sum to full output token count. Does *not* need to have all keys.
+    """
 
 
 class AIMessage(BaseMessage):

libs/core/pyproject.toml

@@ -45,7 +45,7 @@ python = ">=3.12.4"
 
 [tool.ruff.lint]
 select = [ "B", "C4", "E", "F", "I", "N", "PIE", "SIM", "T201", "UP", "W",]
-ignore = [ "UP007",]
+ignore = [ "UP007", 'W293']
 
 [tool.coverage.run]
 omit = [ "tests/*",]

libs/core/tests/unit_tests/prompts/__snapshots__/test_chat.ambr

@@ -677,6 +677,41 @@
         'title': 'HumanMessageChunk',
         'type': 'object',
       }),
+      'InputTokenDetails': dict({
+        'description': '''
+          Breakdown of input token counts.
+
+          Does *not* need to sum to full input token count. Does *not* need to have all keys.
+
+          Example:
+
+              .. code-block:: python
+
+                  {
+                      "audio": 10,
+                      "cache_creation": 200,
+                      "cache_read": 100,
+                  }
+
+          .. versionadded:: 0.3.9
+        ''',
+        'properties': dict({
+          'audio': dict({
+            'title': 'Audio',
+            'type': 'integer',
+          }),
+          'cache_creation': dict({
+            'title': 'Cache Creation',
+            'type': 'integer',
+          }),
+          'cache_read': dict({
+            'title': 'Cache Read',
+            'type': 'integer',
+          }),
+        }),
+        'title': 'InputTokenDetails',
+        'type': 'object',
+      }),
       'InvalidToolCall': dict({
         'description': '''
           Allowance for errors made by LLM.
@@ -743,6 +778,36 @@
         'title': 'InvalidToolCall',
         'type': 'object',
       }),
+      'OutputTokenDetails': dict({
+        'description': '''
+          Breakdown of output token counts.
+
+          Does *not* need to sum to full output token count. Does *not* need to have all keys.
+
+          Example:
+
+              .. code-block:: python
+
+                  {
+                      "audio": 10,
+                      "reasoning": 200,
+                  }
+
+          .. versionadded:: 0.3.9
+        ''',
+        'properties': dict({
+          'audio': dict({
+            'title': 'Audio',
+            'type': 'integer',
+          }),
+          'reasoning': dict({
+            'title': 'Reasoning',
+            'type': 'integer',
+          }),
+        }),
+        'title': 'OutputTokenDetails',
+        'type': 'object',
+      }),
       'SystemMessage': dict({
         'additionalProperties': True,
         'description': '''
@@ -1245,16 +1310,35 @@
               .. code-block:: python
 
                   {
-                      "input_tokens": 10,
-                      "output_tokens": 20,
-                      "total_tokens": 30
+                      "input_tokens": 350,
+                      "output_tokens": 240,
+                      "total_tokens": 590,
+                      "input_token_details": {
+                          "audio": 10,
+                          "cache_creation": 200,
+                          "cache_read": 100,
+                      },
+                      "output_token_details": {
+                          "audio": 10,
+                          "reasoning": 200,
+                      }
                   }
+
+          .. versionchanged:: 0.3.9
+
+              Added ``input_token_details`` and ``output_token_details``.
         ''',
         'properties': dict({
+          'input_token_details': dict({
+            '$ref': '#/$defs/InputTokenDetails',
+          }),
           'input_tokens': dict({
             'title': 'Input Tokens',
             'type': 'integer',
           }),
+          'output_token_details': dict({
+            '$ref': '#/$defs/OutputTokenDetails',
+          }),
           'output_tokens': dict({
             'title': 'Output Tokens',
             'type': 'integer',
@@ -2008,6 +2092,41 @@
         'title': 'HumanMessageChunk',
         'type': 'object',
       }),
+      'InputTokenDetails': dict({
+        'description': '''
+          Breakdown of input token counts.
+
+          Does *not* need to sum to full input token count. Does *not* need to have all keys.
+
+          Example:
+
+              .. code-block:: python
+
+                  {
+                      "audio": 10,
+                      "cache_creation": 200,
+                      "cache_read": 100,
+                  }
+
+          .. versionadded:: 0.3.9
+        ''',
+        'properties': dict({
+          'audio': dict({
+            'title': 'Audio',
+            'type': 'integer',
+          }),
+          'cache_creation': dict({
+            'title': 'Cache Creation',
+            'type': 'integer',
+          }),
+          'cache_read': dict({
+            'title': 'Cache Read',
+            'type': 'integer',
+          }),
+        }),
+        'title': 'InputTokenDetails',
+        'type': 'object',
+      }),
       'InvalidToolCall': dict({
         'description': '''
           Allowance for errors made by LLM.
@@ -2074,6 +2193,36 @@
         'title': 'InvalidToolCall',
         'type': 'object',
       }),
+      'OutputTokenDetails': dict({
+        'description': '''
+          Breakdown of output token counts.
+
+          Does *not* need to sum to full output token count. Does *not* need to have all keys.
+
+          Example:
+
+              .. code-block:: python
+
+                  {
+                      "audio": 10,
+                      "reasoning": 200,
+                  }
+
+          .. versionadded:: 0.3.9
+        ''',
+        'properties': dict({
+          'audio': dict({
+            'title': 'Audio',
+            'type': 'integer',
+          }),
+          'reasoning': dict({
+            'title': 'Reasoning',
+            'type': 'integer',
+          }),
+        }),
+        'title': 'OutputTokenDetails',
+        'type': 'object',
+      }),
       'SystemMessage': dict({
         'additionalProperties': True,
         'description': '''
@@ -2576,16 +2725,35 @@
               .. code-block:: python
 
                   {
-                      "input_tokens": 10,
-                      "output_tokens": 20,
-                      "total_tokens": 30
+                      "input_tokens": 350,
+                      "output_tokens": 240,
+                      "total_tokens": 590,
+                      "input_token_details": {
+                          "audio": 10,
+                          "cache_creation": 200,
+                          "cache_read": 100,
+                      },
+                      "output_token_details": {
+                          "audio": 10,
+                          "reasoning": 200,
+                      }
                   }
+
+          .. versionchanged:: 0.3.9
+
+              Added ``input_token_details`` and ``output_token_details``.
         ''',
         'properties': dict({
+          'input_token_details': dict({
+            '$ref': '#/$defs/InputTokenDetails',
+          }),
           'input_tokens': dict({
             'title': 'Input Tokens',
             'type': 'integer',
           }),
+          'output_token_details': dict({
+            '$ref': '#/$defs/OutputTokenDetails',
+          }),
           'output_tokens': dict({
             'title': 'Output Tokens',
             'type': 'integer',