diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 9921d47aff..4697245f57 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -12,20 +12,23 @@ else: def set_attributes(span: Span, kwargs, response_obj): - from openinference.semconv.trace import ( + from litellm.integrations._types.open_inference import ( MessageAttributes, OpenInferenceSpanKindValues, SpanAttributes, ) try: - litellm_params = kwargs.get("litellm_params", {}) or {} - + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object" + ) + ############################################# ############ LLM CALL METADATA ############## ############################################# - metadata = litellm_params.get("metadata", {}) or {} - span.set_attribute(SpanAttributes.METADATA, str(metadata)) + + if standard_logging_payload and (metadata := standard_logging_payload["metadata"]): + span.set_attribute(SpanAttributes.METADATA, json.dumps(metadata)) ############################################# ########## LLM Request Attributes ########### @@ -62,9 +65,6 @@ def set_attributes(span: Span, kwargs, response_obj): msg.get("content", ""), ) - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) if standard_logging_payload and (model_params := standard_logging_payload["model_parameters"]): # The Generative AI Provider: Azure, OpenAI, etc. span.set_attribute( diff --git a/tests/local_testing/test_arize_ai.py b/tests/local_testing/test_arize_ai.py index 9d2f27ca30..b68981ffd5 100644 --- a/tests/local_testing/test_arize_ai.py +++ b/tests/local_testing/test_arize_ai.py @@ -1,11 +1,14 @@ import asyncio +import json import logging +from litellm import Choices import pytest from dotenv import load_dotenv import litellm from litellm._logging import verbose_logger, verbose_proxy_logger +from litellm.integrations._types.open_inference import SpanAttributes from litellm.integrations.arize.arize import ArizeConfig, ArizeLogger load_dotenv() @@ -58,3 +61,40 @@ def test_get_arize_config_with_endpoints(mock_env_vars, monkeypatch): config = ArizeLogger.get_arize_config() assert config.endpoint == "grpc://test.endpoint" assert config.protocol == "otlp_grpc" + + +def test_arize_set_attributes(): + """ + Test setting attributes for Arize + """ + from unittest.mock import MagicMock + from litellm.types.utils import ModelResponse + + span = MagicMock() + kwargs = { + "role": "user", + "content": "simple arize test", + "model": "gpt-4o", + "messages": [{"role": "user", "content": "basic arize test"}], + "standard_logging_object": {"model_parameters": {"user": "test_user"}, "metadata": {"key": "value", "key2": None}}, + } + response_obj = ModelResponse(usage={"total_tokens": 100, "completion_tokens": 60, "prompt_tokens": 40}, + choices=[Choices(message={"role": "assistant", "content": "response content"})]) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + assert span.set_attribute.call_count == 14 + span.set_attribute.assert_any_call(SpanAttributes.METADATA, json.dumps({"key": "value", "key2": None})) + span.set_attribute.assert_any_call(SpanAttributes.LLM_MODEL_NAME, "gpt-4o") + span.set_attribute.assert_any_call(SpanAttributes.OPENINFERENCE_SPAN_KIND, "LLM") + span.set_attribute.assert_any_call(SpanAttributes.INPUT_VALUE, "basic arize test") + span.set_attribute.assert_any_call("llm.input_messages.0.message.role", "user") + span.set_attribute.assert_any_call("llm.input_messages.0.message.content", "basic arize test") + span.set_attribute.assert_any_call(SpanAttributes.LLM_INVOCATION_PARAMETERS, '{"user": "test_user"}') + span.set_attribute.assert_any_call(SpanAttributes.USER_ID, "test_user") + span.set_attribute.assert_any_call(SpanAttributes.OUTPUT_VALUE, "response content") + span.set_attribute.assert_any_call("llm.output_messages.0.message.role", "assistant") + span.set_attribute.assert_any_call("llm.output_messages.0.message.content", "response content") + span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 100) + span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 60) + span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 40) diff --git a/tests/logging_callback_tests/test_arize_logging.py b/tests/logging_callback_tests/test_arize_logging.py new file mode 100644 index 0000000000..27f067d728 --- /dev/null +++ b/tests/logging_callback_tests/test_arize_logging.py @@ -0,0 +1,42 @@ +import os +import sys +import time +from unittest.mock import Mock, patch + +from litellm.main import completion +import opentelemetry.exporter.otlp.proto.grpc.trace_exporter + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system-path + +import litellm + + +def test_arize_callback(): + litellm.callbacks = ["arize"] + os.environ["ARIZE_SPACE_KEY"] = "test_space_key" + os.environ["ARIZE_API_KEY"] = "test_api_key" + os.environ["ARIZE_ENDPOINT"] = "https://otlp.arize.com/v1" + + # Set the batch span processor to quickly flush after a span has been added + # This is to ensure that the span is exported before the test ends + os.environ["OTEL_BSP_MAX_QUEUE_SIZE"] = "1" + os.environ["OTEL_BSP_MAX_EXPORT_BATCH_SIZE"] = "1" + os.environ["OTEL_BSP_SCHEDULE_DELAY_MILLIS"] = "1" + os.environ["OTEL_BSP_EXPORT_TIMEOUT_MILLIS"] = "5" + + with patch.object( + opentelemetry.exporter.otlp.proto.grpc.trace_exporter.OTLPSpanExporter, + 'export', + new=Mock() + ) as patched_export: + completion( + model="openai/test-model", + messages=[{"role": "user", "content": "arize test content"}], + stream=False, + mock_response="hello there!", + ) + + time.sleep(1) # Wait for the batch span processor to flush + assert patched_export.called