diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 4f1f3b1c25b..95237136b7b 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -529,14 +529,16 @@ def _set_request_attributes( litellm_params.get("custom_llm_provider", "Unknown"), ) - for source_key, request_attribute in ( - ("max_tokens", "llm.request.max_tokens"), - ("temperature", "llm.request.temperature"), - ("top_p", "llm.request.top_p"), + for source_keys, request_attribute in ( + (("max_tokens", "max_completion_tokens"), "llm.request.max_tokens"), + (("temperature",), "llm.request.temperature"), + (("top_p",), "llm.request.top_p"), ): - request_value = optional_params.get(source_key) - if request_value is not None: - safe_set_attribute(span, request_attribute, request_value) + for source_key in source_keys: + request_value = optional_params.get(source_key) + if request_value is not None: + safe_set_attribute(span, request_attribute, request_value) + break safe_set_attribute(span, "llm.is_streaming", str(optional_params.get("stream", False))) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 0630adba202..998bf48cd86 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -2469,14 +2469,16 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): value=provider_name, ) - for source_key, request_attribute in ( - ("max_tokens", SpanAttributes.LLM_REQUEST_MAX_TOKENS.value), - ("temperature", SpanAttributes.LLM_REQUEST_TEMPERATURE.value), - ("top_p", SpanAttributes.LLM_REQUEST_TOP_P.value), + for source_keys, request_attribute in ( + (("max_tokens", "max_completion_tokens"), SpanAttributes.LLM_REQUEST_MAX_TOKENS.value), + (("temperature",), SpanAttributes.LLM_REQUEST_TEMPERATURE.value), + (("top_p",), SpanAttributes.LLM_REQUEST_TOP_P.value), ): - request_value = optional_params.get(source_key) - if request_value is not None: - self.safe_set_attribute(span=span, key=request_attribute, value=request_value) + for source_key in source_keys: + request_value = optional_params.get(source_key) + if request_value is not None: + self.safe_set_attribute(span=span, key=request_attribute, value=request_value) + break if self._gen_ai_semconv_latest_experimental: # Semconv emits gen_ai.request.stream (only when streaming) via diff --git a/tests/unit/integrations/arize/test_arize_utils.py b/tests/unit/integrations/arize/test_arize_utils.py index 3ee723d439b..478a8321ede 100644 --- a/tests/unit/integrations/arize/test_arize_utils.py +++ b/tests/unit/integrations/arize/test_arize_utils.py @@ -219,6 +219,50 @@ def test_arize_records_zero_valued_sampling_params(): span.set_attribute.assert_any_call("llm.request.max_tokens", 0) +@pytest.mark.parametrize( + "optional_params, expected", + [ + ({"max_completion_tokens": 256}, 256), + ({"max_tokens": 128}, 128), + ({"max_tokens": 128, "max_completion_tokens": 256}, 128), + ({"max_tokens": 0, "max_completion_tokens": 256}, 0), + ], +) +def test_arize_records_max_tokens_from_either_request_field(optional_params, expected): + """Callers send max_completion_tokens instead of max_tokens, so the span has to + read both. max_tokens wins when present, including when it is 0, which matches + how LLMRequestParams.from_model_parameters resolves the two. + """ + from unittest.mock import MagicMock + + from litellm.types.utils import ModelResponse + + span = MagicMock() + + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Basic Request Content"}], + "standard_logging_object": { + "model_parameters": {}, + "metadata": {}, + "call_type": "completion", + }, + "optional_params": optional_params, + "litellm_params": {"custom_llm_provider": "openai"}, + } + + response_obj = ModelResponse( + usage={"total_tokens": 100, "completion_tokens": 60, "prompt_tokens": 40}, + choices=[Choices(message={"role": "assistant", "content": "Basic Response Content"})], + model="gpt-4o", + id="chatcmpl-ID", + ) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + span.set_attribute.assert_any_call("llm.request.max_tokens", expected) + + def test_arize_set_attributes_responses_api(): """ Test setting attributes for Responses API with mixed output (reasoning + message). diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 86c7b783431..16cc77ae815 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -442,6 +442,49 @@ class TestOpenTelemetryZeroValuedRequestParams(unittest.TestCase): assert recorded.get(SpanAttributes.LLM_REQUEST_TOP_P.value) == 0 assert recorded.get(SpanAttributes.LLM_REQUEST_MAX_TOKENS.value) == 0 + @parameterized.expand( + [ + ({"max_completion_tokens": 256}, 256), + ({"max_tokens": 128}, 128), + ({"max_tokens": 128, "max_completion_tokens": 256}, 128), + ({"max_tokens": 0, "max_completion_tokens": 256}, 0), + ] + ) + def test_max_tokens_is_read_from_either_request_field(self, optional_params, expected): + """Callers send max_completion_tokens instead of max_tokens, so the span has to + read both. max_tokens wins when present, including when it is 0, which matches + how LLMRequestParams.from_model_parameters resolves the two. + """ + from litellm.proxy._types import SpanAttributes + + otel = OpenTelemetry() + mock_span = MagicMock() + + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "optional_params": optional_params, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": { + "id": "test-id", + "call_type": "completion", + "metadata": {}, + }, + } + + response_obj = { + "id": "test-response-id", + "model": "gpt-4", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + } + + otel.set_attributes(span=mock_span, kwargs=kwargs, response_obj=response_obj) + + recorded = dict(call[0] for call in mock_span.set_attribute.call_args_list if len(call[0]) == 2) + + assert recorded.get(SpanAttributes.LLM_REQUEST_MAX_TOKENS.value) == expected + class TestOpenTelemetryProviderInitialization(unittest.TestCase): """Test suite for verifying provider initialization respects existing providers"""