diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 86f80047e52..3f7416f908a 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -143,7 +143,7 @@ def langfuse_trace_id( class _UsageObject(Protocol): """Token-count surface the Langfuse logger reads off a response usage payload.""" - def get(self, key: Literal["cache_creation_input_tokens", "cache_read_input_tokens"], /) -> int | None: ... + def get(self, key: str, /) -> int | None: ... def _chat_usage(raw_usage: _UsageObject | None) -> _UsageObject | None: @@ -167,7 +167,44 @@ def _chat_usage(raw_usage: _UsageObject | None) -> _UsageObject | None: return raw_usage -def _extract_cache_read_input_tokens(usage_obj) -> int: +def _usage_token_value(usage_obj: object, standard_name: str, decision_name: str | None = None) -> int: + """ + Read a token count from a usage object, tolerating dict-like and Pydantic usage models. + + Decisions responses (DecisionsUsage/OpenAIDecisionUsage) are Pydantic models + without a dict ``.get`` and name their counts ``input_tokens``/``output_tokens`` + (and ``cached_tokens``/``cache_write_tokens``). Fall back to those names when the + standard OpenAI-style attribute is absent. + + Args: + usage_obj: Usage object from an LLM/decisions response. + standard_name: OpenAI-style attribute name (e.g. ``prompt_tokens``). + decision_name: Optional decisions-response field to fall back to. + + Returns: + int: The token count, defaulting to 0. + """ + if isinstance(usage_obj, dict): + usage_proto: Final[_UsageObject] = ( + cast( # cast-ok: plain dict at runtime, narrowed by isinstance; _UsageObject only describes .get + "_UsageObject", usage_obj + ) + ) + # A raw dict here means _chat_usage could not validate it as a chat or + # responses usage (e.g. a malformed /v1/responses payload whose token + # counts fail validation) and passed it through unchanged. Read only the + # standard chat field so a failed validation logs as zero, matching + # getattr(..., None) or 0; do not fall back to decisions-style names, + # which would resurrect a value from a payload that failed validation. + dict_value: Final = usage_proto.get(standard_name) + return dict_value if isinstance(dict_value, int) else 0 + if (attr_value := getattr(usage_obj, standard_name, None)) is None and decision_name is not None: + decision_attr_value: Final = getattr(usage_obj, decision_name, None) + return decision_attr_value if isinstance(decision_attr_value, int) else 0 + return attr_value if isinstance(attr_value, int) else 0 + + +def _extract_cache_read_input_tokens(usage_obj: object) -> int: """ Extract cache_read_input_tokens from usage object. @@ -183,7 +220,7 @@ def _extract_cache_read_input_tokens(usage_obj) -> int: Returns: int: Number of cached tokens read, defaults to 0 """ - cache_read_input_tokens = usage_obj.get("cache_read_input_tokens") or 0 + cache_read_input_tokens = _usage_token_value(usage_obj, "cache_read_input_tokens", "cached_tokens") # Check prompt_tokens_details.cached_tokens (used by Gemini and other providers) if hasattr(usage_obj, "prompt_tokens_details"): @@ -866,11 +903,13 @@ class LangFuseLogger: if _usage_obj: # Safely get usage values, defaulting None to 0 for Langfuse compatibility. # Some providers may return null for token counts. - prompt_tokens: Final = getattr(_usage_obj, "prompt_tokens", None) or 0 - completion_tokens: Final = getattr(_usage_obj, "completion_tokens", None) or 0 - total_tokens: Final = getattr(_usage_obj, "total_tokens", None) or 0 + prompt_tokens: Final = _usage_token_value(_usage_obj, "prompt_tokens", "input_tokens") + completion_tokens: Final = _usage_token_value(_usage_obj, "completion_tokens", "output_tokens") + total_tokens: Final = _usage_token_value(_usage_obj, "total_tokens") - cache_creation_input_tokens: Final = _usage_obj.get("cache_creation_input_tokens") or 0 + cache_creation_input_tokens: Final = _usage_token_value( + _usage_obj, "cache_creation_input_tokens", "cache_write_tokens" + ) cache_read_input_tokens: Final = _extract_cache_read_input_tokens(_usage_obj) usage = { diff --git a/tests/unit/integrations/langfuse/test_langfuse_decisions_usage.py b/tests/unit/integrations/langfuse/test_langfuse_decisions_usage.py new file mode 100644 index 00000000000..c44ec58866d --- /dev/null +++ b/tests/unit/integrations/langfuse/test_langfuse_decisions_usage.py @@ -0,0 +1,74 @@ +# -*- coding: utf-8 -*- +""" +Regression tests for decisions-route responses flowing into the legacy +Langfuse logger (issue #45584). + +DecisionsUsage / OpenAIDecisionUsage are Pydantic models without a dict-style +``.get`` and name their token counts input_tokens/output_tokens (and +cached_tokens/cache_write_tokens). The Langfuse logger used to crash on those +(AttributeError: 'DecisionsUsage' object has no attribute 'get') and drop the +whole trace; token counts were also misread as zero. +""" + +from litellm.integrations.langfuse.langfuse import ( + _extract_cache_read_input_tokens, + _usage_token_value, +) +from litellm.types.decisions import DecisionsUsage + + +class _PlainUsage: + prompt_tokens = 5 + completion_tokens = 7 + total_tokens = 12 + cache_creation_input_tokens = 2 + cache_read_input_tokens = 3 + + +class _EmptyUsage: + pass + + +def test_usage_token_value_reads_standard_attributes() -> None: + usage = _PlainUsage() + assert _usage_token_value(usage, "prompt_tokens") == 5 + assert _usage_token_value(usage, "completion_tokens") == 7 + assert _usage_token_value(usage, "total_tokens") == 12 + assert _usage_token_value(usage, "cache_creation_input_tokens") == 2 + + +def test_usage_token_value_reads_dict_payloads() -> None: + usage = {"prompt_tokens": 5, "completion_tokens": 7, "cache_creation_input_tokens": 2} + assert _usage_token_value(usage, "prompt_tokens") == 5 + assert _usage_token_value(usage, "completion_tokens") == 7 + assert _usage_token_value(usage, "cache_creation_input_tokens") == 2 + + +def test_usage_token_value_defaults_missing_counts_to_zero() -> None: + assert _usage_token_value(_EmptyUsage(), "prompt_tokens") == 0 + assert _usage_token_value({}, "prompt_tokens") == 0 + + +def test_usage_token_value_falls_back_to_decisions_field_names() -> None: + usage = DecisionsUsage(input_tokens=5, output_tokens=7, cached_tokens=3, cache_write_tokens=2) + assert _usage_token_value(usage, "prompt_tokens", "input_tokens") == 5 + assert _usage_token_value(usage, "completion_tokens", "output_tokens") == 7 + # DecisionsUsage has no total_tokens; must default to 0 rather than crash. + assert _usage_token_value(usage, "total_tokens") == 0 + assert _usage_token_value(usage, "cache_creation_input_tokens", "cache_write_tokens") == 2 + + +def test_extract_cache_read_accepts_pydantic_decisions_usage() -> None: + usage = DecisionsUsage(input_tokens=5, output_tokens=7, cached_tokens=3, cache_write_tokens=2) + assert _extract_cache_read_input_tokens(usage) == 3 + + +def test_extract_cache_read_keeps_dict_semantics() -> None: + usage = {"cache_read_input_tokens": 4, "prompt_tokens": 10} + assert _extract_cache_read_input_tokens(usage) == 4 + assert _extract_cache_read_input_tokens({}) == 0 + + +def test_extract_cache_read_keeps_plain_attribute_semantics() -> None: + usage = _PlainUsage() + assert _extract_cache_read_input_tokens(usage) == 3