mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(arize): _set_usage_outputs handles raw OpenAI Pydantic CompletionUsage (#26506)
* [Feat] Day-0 support for GPT-5.5 and GPT-5.5 Pro (#26449) * feat(openai): day-0 support for GPT-5.5 and GPT-5.5 Pro Add pricing + capability entries for the new GPT-5.5 family launched by OpenAI on 2026-04-24: - gpt-5.5 / gpt-5.5-2026-04-23 (chat): $5/$30/$0.50 per 1M input/output/cached input - gpt-5.5-pro / gpt-5.5-pro-2026-04-23 (responses-only): $60/$360/$6 per 1M input/output/cached input Other fees (long-context >272k, flex, batches, priority, cache discounts) follow the same ratios as GPT-5.4, with context window retained at 1.05M input / 128K output. No transformation / classifier code changes are required: OpenAIGPT5Config.is_model_gpt_5_4_plus_model() already matches 5.5+ via numeric version parsing, and model registration is driven from the JSON. The existing responses-API bridge for tools + reasoning_effort (litellm/main.py:970) already covers gpt-5.5-pro. Tests: - GPT5_MODELS regression list now covers gpt-5.5-pro and dated variants - New test_generic_cost_per_token_gpt55_pro cost-calc test - Updated test_generic_cost_per_token_gpt55 for long-context fields * fix(openai): mirror reasoning_effort flags onto gpt-5.5 dated variants gpt-5.5-2026-04-23 and gpt-5.5-pro-2026-04-23 were missing the supports_none_reasoning_effort, supports_xhigh_reasoning_effort, and supports_minimal_reasoning_effort flags that their non-dated counterparts define. Reasoning-effort routing in OpenAIGPT5Config is fully capability-driven from these JSON flags — since an absent flag is treated as False for opt-in levels (xhigh), users pinning to a dated snapshot would silently lose xhigh support and diverge from the base alias on logprobs + flexible temperature handling. Copy the flags onto both dated variants so every dated snapshot inherits the base model's reasoning-effort capability profile. Adds a parametrized regression test that asserts supports_{none,minimal,xhigh}_reasoning_effort parity between each dated variant and its non-dated counterpart, preventing future drift when new snapshots are added. * [Feat] Add azure/gpt-5.5 + azure/gpt-5.5-pro entries (+ dated variants) (#26361) * feat(azure): add azure/gpt-5.5 + azure/gpt-5.5-pro entries (+ dated variants) Azure variants of OpenAI's GPT-5.5 family. Microsoft has not yet shipped GPT-5.5 on Azure OpenAI (latest GA on the Foundry models page is GPT-5.4 as of 2026-04-24), but adding the entries day-0 mirrors the established precedent for azure/gpt-5.4* (which were in the cost map before the Azure rollout) so cost tracking and capability flags work the moment customers deploy. Schema follows the existing azure/gpt-5.4* shape: - Same base/long-context pricing as openai/gpt-5.5*: $5/$30 chat, $60/$360 pro per 1M, with priority tier 2x base - Azure variants drop the flex/batches keys (Azure has no flex tier) but keep priority pricing, matching gpt-5.4* precedent - mode=chat for the thinking model, mode=responses for pro reasoning_effort capability flags mirror the OpenAI variants exactly since Azure proxies the same API contract: minimal rejection on both chat and pro, low/none rejection on pro. Once #26456 (which sets supports_low_reasoning_effort + minimal=false on openai/gpt-5.5*) lands, OpenAI and Azure flag profiles align. Tests pin entry presence + pricing for all four Azure variants and verify the live-API-derived reasoning_effort flags. * test: register supports_low_reasoning_effort in cost-map JSON schema azure/gpt-5.5-pro and azure/gpt-5.5-pro-2026-04-23 added in this branch carry supports_low_reasoning_effort=false. The strict 'additionalProperties: false' schema in test_aaamodel_prices_and_context_window_json_is_valid rejected the new key. Register it alongside the other supports_*_reasoning_effort entries. Note: the runtime side of this flag (code that reads it) lands in #26456. Until that PR merges the flag is inert for both Azure and OpenAI pro entries, but having the schema accept it lets cost-map tests pass on either merge order. * fix(arize/langfuse_otel): handle Pydantic usage objects without `.get` `_set_usage_outputs` called `usage.get(...)` and `usage.get('output_tokens_details', {}).get('reasoning_tokens')`. These crash with `AttributeError: 'CompletionUsage' object has no attribute 'get'` when `usage` (or the nested token-details object) is a raw OpenAI Pydantic model rather than a dict / litellm `Usage` wrapper. Reproduces on the langfuse_otel + arize Responses API logging paths. Fixes #13672. Changes: - Add `_safe_get(obj, key, default)` that prefers dict-style `.get` when available and otherwise falls back to `getattr`. Works uniformly for dicts, litellm's `Usage`, and plain Pydantic models like `openai.types.completion_usage.CompletionUsage` / `CompletionTokensDetails` / `OutputTokensDetails`. - Use `_safe_get` for total / completion / prompt / output tokens. - Look for reasoning tokens in `completion_tokens_details` (Chat Completions API) before falling back to `output_tokens_details` (Responses API). Previously reasoning tokens from the Chat Completions API were silently dropped. Tests: - `test_set_usage_outputs_pydantic_completion_usage` — covers the chat completions path with raw `CompletionUsage` + `CompletionTokensDetails`. - `test_set_usage_outputs_pydantic_response_api_usage` — covers the Responses API path with a Pydantic usage object lacking `.get`. Both tests fail on main before this commit and pass after. --------- Co-authored-by: yuneng-jiang <yuneng@berri.ai> Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: alvinttang <alvin@pm.me> Co-authored-by: Krrish Dholakia <krrish+github@berri.ai>
This commit is contained in:
parent
f63a6f1b26
commit
98a9005c76
5 changed files with 450 additions and 4 deletions
|
|
@ -220,23 +220,57 @@ def _set_structured_outputs(span: "Span", response_obj, msg_attrs, span_attrs):
|
|||
safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_ROLE}", message_role)
|
||||
|
||||
|
||||
def _safe_get(obj, key, default=None):
|
||||
"""Read ``key`` from a dict-like or Pydantic-model-like object.
|
||||
|
||||
The arize/langfuse_otel logger receives ``usage`` objects from many sources:
|
||||
plain dicts, litellm ``Usage`` (which exposes ``.get``), and raw OpenAI
|
||||
Pydantic models (e.g. ``openai.types.completion_usage.CompletionUsage`` and
|
||||
nested ``CompletionTokensDetails`` / ``OutputTokensDetails``) which do NOT
|
||||
expose ``.get``. Calling ``.get`` on the latter raised ``AttributeError`` —
|
||||
see https://github.com/BerriAI/litellm/issues/13672.
|
||||
"""
|
||||
if obj is None:
|
||||
return default
|
||||
getter = getattr(obj, "get", None)
|
||||
if callable(getter):
|
||||
try:
|
||||
return getter(key, default)
|
||||
except TypeError:
|
||||
# Some objects expose `.get` with a different signature
|
||||
pass
|
||||
return getattr(obj, key, default)
|
||||
|
||||
|
||||
def _set_usage_outputs(span: "Span", response_obj, span_attrs):
|
||||
usage = response_obj and response_obj.get("usage")
|
||||
if not usage:
|
||||
return
|
||||
|
||||
safe_set_attribute(
|
||||
span, span_attrs.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens")
|
||||
span, span_attrs.LLM_TOKEN_COUNT_TOTAL, _safe_get(usage, "total_tokens")
|
||||
)
|
||||
completion_tokens = _safe_get(usage, "completion_tokens") or _safe_get(
|
||||
usage, "output_tokens"
|
||||
)
|
||||
completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens")
|
||||
if completion_tokens:
|
||||
safe_set_attribute(
|
||||
span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens
|
||||
)
|
||||
prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens")
|
||||
prompt_tokens = _safe_get(usage, "prompt_tokens") or _safe_get(
|
||||
usage, "input_tokens"
|
||||
)
|
||||
if prompt_tokens:
|
||||
safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_PROMPT, prompt_tokens)
|
||||
reasoning_tokens = usage.get("output_tokens_details", {}).get("reasoning_tokens")
|
||||
|
||||
# Reasoning tokens live in `completion_tokens_details` for Chat Completions
|
||||
# API (Usage) and in `output_tokens_details` for Responses API
|
||||
# (ResponseAPIUsage). Both nested objects may be plain Pydantic models
|
||||
# without `.get`.
|
||||
token_details = _safe_get(usage, "completion_tokens_details") or _safe_get(
|
||||
usage, "output_tokens_details"
|
||||
)
|
||||
reasoning_tokens = _safe_get(token_details, "reasoning_tokens")
|
||||
if reasoning_tokens:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
|
|
|
|||
|
|
@ -4645,6 +4645,169 @@
|
|||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"azure/gpt-5.5": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
|
||||
"cache_read_input_token_cost_priority": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1e-05,
|
||||
"input_cost_per_token_priority": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 2e-05,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.5e-05,
|
||||
"output_cost_per_token_priority": 6e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 9e-05,
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_service_tier": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false
|
||||
},
|
||||
"azure/gpt-5.5-2026-04-23": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
|
||||
"cache_read_input_token_cost_priority": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1e-05,
|
||||
"input_cost_per_token_priority": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 2e-05,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.5e-05,
|
||||
"output_cost_per_token_priority": 6e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 9e-05,
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_service_tier": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"azure/gpt-5.5-pro": {
|
||||
"cache_read_input_token_cost": 6e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
|
||||
"input_cost_per_token": 6e-05,
|
||||
"input_cost_per_token_above_272k_tokens": 0.00012,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 0.00036,
|
||||
"output_cost_per_token_above_272k_tokens": 0.00054,
|
||||
"supported_endpoints": [
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_none_reasoning_effort": false,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_low_reasoning_effort": false
|
||||
},
|
||||
"azure/gpt-5.5-pro-2026-04-23": {
|
||||
"cache_read_input_token_cost": 6e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
|
||||
"input_cost_per_token": 6e-05,
|
||||
"input_cost_per_token_above_272k_tokens": 0.00012,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 0.00036,
|
||||
"output_cost_per_token_above_272k_tokens": 0.00054,
|
||||
"supported_endpoints": [
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"azure/gpt-5.4-mini": {
|
||||
"cache_read_input_token_cost": 7.5e-08,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
|
|
|
|||
|
|
@ -4659,6 +4659,169 @@
|
|||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"azure/gpt-5.5": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
|
||||
"cache_read_input_token_cost_priority": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1e-05,
|
||||
"input_cost_per_token_priority": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 2e-05,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.5e-05,
|
||||
"output_cost_per_token_priority": 6e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 9e-05,
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_service_tier": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false
|
||||
},
|
||||
"azure/gpt-5.5-2026-04-23": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
|
||||
"cache_read_input_token_cost_priority": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1e-05,
|
||||
"input_cost_per_token_priority": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 2e-05,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.5e-05,
|
||||
"output_cost_per_token_priority": 6e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 9e-05,
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_service_tier": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"azure/gpt-5.5-pro": {
|
||||
"cache_read_input_token_cost": 6e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
|
||||
"input_cost_per_token": 6e-05,
|
||||
"input_cost_per_token_above_272k_tokens": 0.00012,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 0.00036,
|
||||
"output_cost_per_token_above_272k_tokens": 0.00054,
|
||||
"supported_endpoints": [
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_none_reasoning_effort": false,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_low_reasoning_effort": false
|
||||
},
|
||||
"azure/gpt-5.5-pro-2026-04-23": {
|
||||
"cache_read_input_token_cost": 6e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1.2e-05,
|
||||
"input_cost_per_token": 6e-05,
|
||||
"input_cost_per_token_above_272k_tokens": 0.00012,
|
||||
"litellm_provider": "azure",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"output_cost_per_token": 0.00036,
|
||||
"output_cost_per_token_above_272k_tokens": 0.00054,
|
||||
"supported_endpoints": [
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": false,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true
|
||||
},
|
||||
"azure/gpt-5.4-mini": {
|
||||
"cache_read_input_token_cost": 7.5e-08,
|
||||
"input_cost_per_token": 7.5e-07,
|
||||
|
|
|
|||
|
|
@ -273,6 +273,91 @@ def test_arize_set_attributes_responses_api():
|
|||
)
|
||||
|
||||
|
||||
def test_set_usage_outputs_pydantic_completion_usage():
|
||||
"""
|
||||
Regression test for https://github.com/BerriAI/litellm/issues/13672
|
||||
|
||||
`_set_usage_outputs` previously called `usage.get(...)` which crashes when
|
||||
`usage` is a plain Pydantic model (e.g. openai.types.completion_usage.CompletionUsage)
|
||||
that does not implement dict-style `.get()`. Same crash for nested
|
||||
`output_tokens_details` / `completion_tokens_details`.
|
||||
|
||||
The function must:
|
||||
1. Read total/prompt/completion tokens from a Pydantic usage without `.get`.
|
||||
2. Read reasoning_tokens from `completion_tokens_details` (chat completions API)
|
||||
OR `output_tokens_details` (responses API), even when those nested objects
|
||||
are Pydantic models without `.get`.
|
||||
3. Not raise AttributeError; not call span.record_exception.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from openai.types.completion_usage import (
|
||||
CompletionTokensDetails,
|
||||
CompletionUsage,
|
||||
)
|
||||
|
||||
from litellm.integrations.arize._utils import _set_usage_outputs
|
||||
|
||||
span = MagicMock()
|
||||
|
||||
# Plain OpenAI Pydantic model — has no `.get()`
|
||||
usage = CompletionUsage(
|
||||
completion_tokens=60,
|
||||
prompt_tokens=40,
|
||||
total_tokens=100,
|
||||
completion_tokens_details=CompletionTokensDetails(reasoning_tokens=25),
|
||||
)
|
||||
assert not hasattr(usage, "get"), "precondition: CompletionUsage must lack .get"
|
||||
|
||||
response_obj = {"usage": usage}
|
||||
|
||||
# Must not raise
|
||||
_set_usage_outputs(span, response_obj, SpanAttributes)
|
||||
|
||||
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 100)
|
||||
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 40)
|
||||
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 60)
|
||||
# reasoning_tokens for chat completions live in completion_tokens_details
|
||||
span.set_attribute.assert_any_call(
|
||||
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 25
|
||||
)
|
||||
|
||||
|
||||
def test_set_usage_outputs_pydantic_response_api_usage():
|
||||
"""
|
||||
Same crash also affects Responses API with `output_tokens_details` as a
|
||||
Pydantic model that lacks `.get()`. Verifies the responses-API path.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.integrations.arize._utils import _set_usage_outputs
|
||||
from litellm.types.llms.openai import OutputTokensDetails
|
||||
|
||||
# Build an object that mimics openai ResponsesAPI usage but lacks `.get`
|
||||
# (uses a plain class — not BaseLiteLLMOpenAIResponseObject)
|
||||
class PlainResponsesUsage:
|
||||
def __init__(self):
|
||||
self.total_tokens = 370
|
||||
self.input_tokens = 120
|
||||
self.output_tokens = 250
|
||||
self.output_tokens_details = OutputTokensDetails(reasoning_tokens=180)
|
||||
|
||||
usage = PlainResponsesUsage()
|
||||
assert not hasattr(usage, "get")
|
||||
|
||||
span = MagicMock()
|
||||
response_obj = {"usage": usage}
|
||||
|
||||
_set_usage_outputs(span, response_obj, SpanAttributes)
|
||||
|
||||
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 370)
|
||||
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 120)
|
||||
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 250)
|
||||
span.set_attribute.assert_any_call(
|
||||
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 180
|
||||
)
|
||||
|
||||
|
||||
class TestArizeLogger(CustomLogger):
|
||||
"""
|
||||
Custom logger implementation to capture standard_callback_dynamic_params.
|
||||
|
|
|
|||
|
|
@ -769,6 +769,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
|||
"uses_embed_content": {"type": "boolean"},
|
||||
"supports_reasoning": {"type": "boolean"},
|
||||
"supports_minimal_reasoning_effort": {"type": "boolean"},
|
||||
"supports_low_reasoning_effort": {"type": "boolean"},
|
||||
"supports_none_reasoning_effort": {"type": "boolean"},
|
||||
"supports_xhigh_reasoning_effort": {"type": "boolean"},
|
||||
"supports_max_reasoning_effort": {"type": "boolean"},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue