fix(fireworks): report cached prompt token usage (#38751)

Closes #38648

`langchain-fireworks` now reports Fireworks cached prompt tokens in
`AIMessage.usage_metadata.input_token_details.cache_read` and no longer
crashes when combining nested token usage from batched `generate()`
calls.

---

Fireworks can return nested token usage details when prompt caching is
involved, including cached prompt token counts. Batched `generate()`
calls were crashing when those nested dictionaries were combined, and
regular chat results did not expose the cached-token breakdown in the
standard LangChain usage metadata shape.

This updates `ChatFireworks` so nested token usage is merged safely and
cached prompt tokens are reported as `input_token_details.cache_read`.
Users and downstream tracing systems can now distinguish cached
Fireworks input tokens from regular input tokens instead of treating the
full prompt as uncached input.

Thanks to @abcgco for the original report and recursive merge fix in
#38646, and to @abhi-0203 for independently identifying the same nested
`token_usage` failure in #38735. This PR builds on that work by using
the recursive merge approach and extending the fix to normalize cached
prompt tokens into standard usage metadata for tracing and cost
reporting.

Co-authored-by: Andrei Boldyrev <abcgco@gmail.com>

libs/partners/fireworks/langchain_fireworks/chat_models.py

@@ -11,6 +11,7 @@ from typing import (
     Any,
     Literal,
     NoReturn,
+    TypeAlias,
     cast,
 )
 
@@ -58,6 +59,7 @@ from langchain_core.messages import (
     ToolCall,
     ToolMessage,
     ToolMessageChunk,
+    UsageMetadata,
     is_data_content_block,
 )
 from langchain_core.messages.block_translators.openai import (
@@ -395,14 +397,72 @@ def _convert_message_to_dict(message: BaseMessage) -> dict:
     return message_dict
 
 
-def _usage_to_metadata(usage: Mapping[str, Any]) -> dict[str, int]:
-    input_tokens = usage.get("prompt_tokens", 0)
-    output_tokens = usage.get("completion_tokens", 0)
-    return {
+def _usage_to_metadata(usage: Mapping[str, Any]) -> UsageMetadata:
+    input_tokens = usage.get("prompt_tokens") or 0
+    output_tokens = usage.get("completion_tokens") or 0
+    usage_metadata: UsageMetadata = {
         "input_tokens": input_tokens,
         "output_tokens": output_tokens,
-        "total_tokens": usage.get("total_tokens", input_tokens + output_tokens),
+        "total_tokens": usage.get("total_tokens") or input_tokens + output_tokens,
     }
+    cached_tokens = (usage.get("prompt_tokens_details") or {}).get("cached_tokens")
+    if cached_tokens is not None:
+        usage_metadata["input_token_details"] = {"cache_read": cached_tokens}
+    return usage_metadata
+
+
+TokenUsageTree: TypeAlias = "int | dict[str, TokenUsageTree]"
+"""Raw provider token usage: a tree of `int` leaves and nested `dict` nodes
+(e.g. `prompt_tokens_details`).
+
+Modeled as a recursive alias so the merge helper's signature carries the shape
+rather than leaving it to `Any`.
+"""
+
+
+def _update_token_usage(
+    overall_token_usage: TokenUsageTree, new_usage: TokenUsageTree
+) -> TokenUsageTree:
+    """Recursively merge raw provider token usage across generations.
+
+    Token usage is a tree of `int` leaves (summed) and `dict` nodes such as
+    `prompt_tokens_details` (merged key-by-key, skipping `None` values).
+
+    A type mismatch between the accumulator and the incoming value (e.g. an
+    `int` on one side and a `dict` on the other) indicates malformed provider
+    data and is raised rather than silently coerced. An entirely unexpected
+    leaf type (neither `int` nor `dict`) is logged and passed through, so a
+    telemetry anomaly degrades gracefully instead of failing the response.
+    """
+    if isinstance(new_usage, int):
+        if not isinstance(overall_token_usage, int):
+            msg = (
+                "Got different types for token usage: "
+                f"{new_usage!r} ({type(new_usage).__name__}) and "
+                f"{overall_token_usage!r} ({type(overall_token_usage).__name__})"
+            )
+            raise ValueError(msg)
+        return overall_token_usage + new_usage
+    if isinstance(new_usage, dict):
+        if not isinstance(overall_token_usage, dict):
+            msg = (
+                "Got different types for token usage: "
+                f"{new_usage!r} ({type(new_usage).__name__}) and "
+                f"{overall_token_usage!r} ({type(overall_token_usage).__name__})"
+            )
+            raise ValueError(msg)
+        updated_token_usage = dict(overall_token_usage)
+        for key, value in new_usage.items():
+            if value is not None:
+                # Seed a first-seen key with an empty node of the same kind so a
+                # nested `dict` value merges rather than colliding with an `int`.
+                default: TokenUsageTree = {} if isinstance(value, dict) else 0
+                updated_token_usage[key] = _update_token_usage(
+                    overall_token_usage.get(key, default), value
+                )
+        return updated_token_usage
+    logger.warning("Unexpected type for token usage: %s", type(new_usage).__name__)
+    return new_usage
 
 
 def _convert_chunk_to_message_chunk(
@@ -423,7 +483,7 @@ def _convert_chunk_to_message_chunk(
         usage_metadata = _usage_to_metadata(usage) if usage else None
         return AIMessageChunk(
             content="",
-            usage_metadata=usage_metadata,  # type: ignore[arg-type]
+            usage_metadata=usage_metadata,
             response_metadata=response_metadata,
         )
     choice = choices[0]
@@ -458,7 +518,7 @@ def _convert_chunk_to_message_chunk(
             content=content,
             additional_kwargs=additional_kwargs,
             tool_call_chunks=tool_call_chunks,
-            usage_metadata=usage_metadata,  # type: ignore[arg-type]
+            usage_metadata=usage_metadata,
             response_metadata=response_metadata,
         )
     if role == "system" or default_class == SystemMessageChunk:
@@ -960,11 +1020,15 @@ class ChatFireworks(BaseChatModel):
             if output is None:
                 # Happens in streaming
                 continue
-            token_usage = output["token_usage"]
+            token_usage = output.get("token_usage")
             if token_usage is not None:
                 for k, v in token_usage.items():
+                    if v is None:
+                        continue
                     if k in overall_token_usage:
-                        overall_token_usage[k] += v
+                        overall_token_usage[k] = _update_token_usage(
+                            overall_token_usage[k], v
+                        )
                     else:
                         overall_token_usage[k] = v
             if system_fingerprint is None:
@@ -1064,11 +1128,7 @@ class ChatFireworks(BaseChatModel):
             message = _convert_dict_to_message(res["message"])
             if isinstance(message, AIMessage):
                 if token_usage:
-                    message.usage_metadata = {
-                        "input_tokens": token_usage.get("prompt_tokens", 0),
-                        "output_tokens": token_usage.get("completion_tokens", 0),
-                        "total_tokens": token_usage.get("total_tokens", 0),
-                    }
+                    message.usage_metadata = _usage_to_metadata(token_usage)
                     message.response_metadata["model_provider"] = "fireworks"
                     message.response_metadata["model_name"] = self.model_name
                 if service_tier:

libs/partners/fireworks/tests/unit_tests/test_chat_models.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import logging
 import os
 from typing import Any
 from unittest.mock import MagicMock
@@ -36,6 +37,7 @@ from langchain_fireworks.chat_models import (
     _convert_message_to_dict,
     _format_message_content,
     _sanitize_chat_completions_content,
+    _update_token_usage,
     _usage_to_metadata,
 )
 
@@ -1066,6 +1068,199 @@ class TestUsageToMetadata:
         result = _usage_to_metadata({})
         assert result == {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
 
+    def test_explicit_none_fields_coerced_to_zero(self) -> None:
+        """Provider may send explicit `None` values; coerce them to `0`.
+
+        Guards the `or`-based fallbacks against a `.get(key, default)` regression,
+        which would preserve `None` for a present-but-null key.
+        """
+        result = _usage_to_metadata(
+            {
+                "prompt_tokens": None,
+                "completion_tokens": None,
+                "total_tokens": None,
+            }
+        )
+        assert result == {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
+
+    def test_total_tokens_falls_back_to_sum_when_none(self) -> None:
+        """A null `total_tokens` falls back to `input + output`."""
+        result = _usage_to_metadata(
+            {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": None}
+        )
+        assert result == {"input_tokens": 7, "output_tokens": 3, "total_tokens": 10}
+
+    def test_cached_prompt_tokens_mapped_to_cache_read(self) -> None:
+        result = _usage_to_metadata(
+            {
+                "prompt_tokens": 10,
+                "completion_tokens": 5,
+                "total_tokens": 15,
+                "prompt_tokens_details": {"cached_tokens": 7},
+            }
+        )
+        assert result == {
+            "input_tokens": 10,
+            "output_tokens": 5,
+            "total_tokens": 15,
+            "input_token_details": {"cache_read": 7},
+        }
+
+    def test_cached_tokens_zero_preserved(self) -> None:
+        """A genuine `0` cache hit is reported, not dropped.
+
+        Guards the `is not None` check against a truthiness (`if cached_tokens:`)
+        regression that would silently omit `cache_read` for a real zero.
+        """
+        result = _usage_to_metadata(
+            {
+                "prompt_tokens": 10,
+                "completion_tokens": 5,
+                "total_tokens": 15,
+                "prompt_tokens_details": {"cached_tokens": 0},
+            }
+        )
+        assert result["input_token_details"] == {"cache_read": 0}
+
+    def test_prompt_tokens_details_without_cached_tokens_omits_detail(self) -> None:
+        """A details dict lacking (or nulling) `cached_tokens` adds no detail."""
+        assert "input_token_details" not in _usage_to_metadata(
+            {"prompt_tokens": 5, "prompt_tokens_details": {}}
+        )
+        assert "input_token_details" not in _usage_to_metadata(
+            {"prompt_tokens": 5, "prompt_tokens_details": {"cached_tokens": None}}
+        )
+
+
+class TestCombineLLMOutputs:
+    """Tests for combining raw provider token usage across generations."""
+
+    def test_combines_nested_token_usage(self) -> None:
+        model = _make_model()
+        result = model._combine_llm_outputs(
+            [
+                {
+                    "token_usage": {
+                        "prompt_tokens": 32,
+                        "completion_tokens": 51,
+                        "total_tokens": 83,
+                        "prompt_tokens_details": {"cached_tokens": 0},
+                    },
+                    "system_fingerprint": "fp-1",
+                },
+                {
+                    "token_usage": {
+                        "prompt_tokens": 44341,
+                        "completion_tokens": 10,
+                        "total_tokens": 44351,
+                        "prompt_tokens_details": {"cached_tokens": 41518},
+                    },
+                },
+            ]
+        )
+        assert result == {
+            "token_usage": {
+                "prompt_tokens": 44373,
+                "completion_tokens": 61,
+                "total_tokens": 44434,
+                "prompt_tokens_details": {"cached_tokens": 41518},
+            },
+            "model_name": MODEL_NAME,
+            "system_fingerprint": "fp-1",
+        }
+
+    def test_preserves_prior_nested_token_usage_keys(self) -> None:
+        model = _make_model()
+        result = model._combine_llm_outputs(
+            [
+                {
+                    "token_usage": {
+                        "prompt_tokens_details": {
+                            "audio_tokens": 4,
+                            "cached_tokens": 8,
+                        },
+                    },
+                },
+                {
+                    "token_usage": {
+                        "prompt_tokens_details": {
+                            "audio_tokens": 6,
+                        },
+                    },
+                },
+                {
+                    "token_usage": {
+                        "prompt_tokens_details": {
+                            "cached_tokens": None,
+                        },
+                    },
+                },
+            ]
+        )
+
+        assert result["token_usage"] == {
+            "prompt_tokens_details": {
+                "audio_tokens": 10,
+                "cached_tokens": 8,
+            },
+        }
+
+    def test_skips_none_token_usage_values(self) -> None:
+        model = _make_model()
+        result = model._combine_llm_outputs(
+            [
+                {"token_usage": {"prompt_tokens_details": None}},
+                {
+                    "token_usage": {
+                        "prompt_tokens_details": {"cached_tokens": 8},
+                    }
+                },
+            ]
+        )
+        assert result["token_usage"] == {"prompt_tokens_details": {"cached_tokens": 8}}
+
+    def test_skips_none_streaming_outputs(self) -> None:
+        """`None` entries (produced during streaming) are skipped, not dereferenced."""
+        model = _make_model()
+        result = model._combine_llm_outputs(
+            [None, {"token_usage": {"prompt_tokens": 5, "total_tokens": 5}}, None]
+        )
+        assert result["token_usage"] == {"prompt_tokens": 5, "total_tokens": 5}
+
+
+class TestUpdateTokenUsage:
+    """Tests for the recursive `_update_token_usage` merge helper.
+
+    The type-mismatch and unexpected-type branches are unreachable with today's
+    stable Fireworks payloads, so they are exercised directly here to lock in the
+    behavior: mismatches raise, while a wholly unexpected leaf type is logged and
+    passed through rather than failing the response.
+    """
+
+    def test_int_accumulator_with_dict_value_raises(self) -> None:
+        with pytest.raises(ValueError, match="Got different types for token usage"):
+            _update_token_usage(5, {"cached_tokens": 1})
+
+    def test_dict_accumulator_with_int_value_raises(self) -> None:
+        with pytest.raises(ValueError, match="Got different types for token usage"):
+            _update_token_usage({"cached_tokens": 1}, 5)
+
+    def test_unexpected_value_type_warns_and_passes_through(
+        self, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        with caplog.at_level(logging.WARNING):
+            result = _update_token_usage(0, 1.5)  # type: ignore[arg-type]
+        assert result == 1.5
+        assert "Unexpected type for token usage" in caplog.text
+
+    def test_first_seen_nested_dict_value_merges(self) -> None:
+        """A first-seen nested `dict` node seeds as a dict instead of raising."""
+        result = _update_token_usage(
+            {"details": {"a": 1}},
+            {"details": {"a": 2, "nested": {"b": 3}}},
+        )
+        assert result == {"details": {"a": 3, "nested": {"b": 3}}}
+
 
 class TestConvertChunkToMessageChunk:
     """Tests for `_convert_chunk_to_message_chunk` empty-choices handling."""
@@ -1108,6 +1303,56 @@ class TestConvertChunkToMessageChunk:
             "total_tokens": 3,
         }
 
+    def test_usage_chunk_maps_cached_prompt_tokens(self) -> None:
+        chunk = {
+            "choices": [],
+            "usage": {
+                "prompt_tokens": 10,
+                "completion_tokens": 2,
+                "total_tokens": 12,
+                "prompt_tokens_details": {"cached_tokens": 6},
+            },
+        }
+        result = _convert_chunk_to_message_chunk(chunk, AIMessageChunk)
+        assert isinstance(result, AIMessageChunk)
+        assert result.usage_metadata == {
+            "input_tokens": 10,
+            "output_tokens": 2,
+            "total_tokens": 12,
+            "input_token_details": {"cache_read": 6},
+        }
+
+
+class TestCreateChatResult:
+    """Tests for converting Fireworks responses into chat generations."""
+
+    def test_maps_cached_prompt_tokens_to_message_usage_metadata(self) -> None:
+        model = _make_model()
+        chat_result = model._create_chat_result(
+            {
+                "choices": [
+                    {
+                        "message": {"role": "assistant", "content": "ok"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 20,
+                    "completion_tokens": 3,
+                    "total_tokens": 23,
+                    "prompt_tokens_details": {"cached_tokens": 11},
+                },
+            }
+        )
+        message = chat_result.generations[0].message
+        assert isinstance(message, AIMessage)
+        assert message.usage_metadata == {
+            "input_tokens": 20,
+            "output_tokens": 3,
+            "total_tokens": 23,
+            "input_token_details": {"cache_read": 11},
+        }
+
 
 class TestExtraHeaders:
     """Tests for request-specific HTTP header plumbing."""