diff --git a/libs/partners/fireworks/langchain_fireworks/chat_models.py b/libs/partners/fireworks/langchain_fireworks/chat_models.py index ee94e9ed7a8..30567fa0e2a 100644 --- a/libs/partners/fireworks/langchain_fireworks/chat_models.py +++ b/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: diff --git a/libs/partners/fireworks/tests/unit_tests/test_chat_models.py b/libs/partners/fireworks/tests/unit_tests/test_chat_models.py index e586411fa4a..6f2df3b8314 100644 --- a/libs/partners/fireworks/tests/unit_tests/test_chat_models.py +++ b/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."""