diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index f4fe40738b..1a743e5c09 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -49,28 +49,37 @@ class OpenTelemetryConfig: exporter: Union[str, SpanExporter] = "console" endpoint: Optional[str] = None headers: Optional[str] = None + debug: Optional[str] = None @classmethod def from_env(cls): """ - OTEL_HEADERS=x-honeycomb-team=B85YgLm9**** OTEL_EXPORTER="otlp_http" OTEL_ENDPOINT="https://api.honeycomb.io/v1/traces" + OTEL_HEADERS=x-honeycomb-team=B85YgLm9**** + DEBUG_OTEL="true" OTEL_HEADERS gets sent as headers = {"x-honeycomb-team": "B85YgLm96******"} """ - from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( - InMemorySpanExporter, - ) - if os.getenv("OTEL_EXPORTER") == "in_memory": + # Declare LiteLLM variables + exporter = os.getenv("OTEL_EXPORTER", "console") + endpoint = os.getenv("OTEL_ENDPOINT") + headers = os.getenv("OTEL_HEADERS") + debug = os.getenv("DEBUG_OTEL") + + if exporter == "in_memory": + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + return cls(exporter=InMemorySpanExporter()) + return cls( - exporter=os.getenv("OTEL_EXPORTER", "console"), - endpoint=os.getenv("OTEL_ENDPOINT"), - headers=os.getenv( - "OTEL_HEADERS" - ), # example: OTEL_HEADERS=x-honeycomb-team=B85YgLm96***" + exporter=exporter, + endpoint=endpoint, + headers=headers, + debug=str(debug).lower(), ) @@ -82,29 +91,20 @@ class OpenTelemetry(CustomLogger): **kwargs, ): from opentelemetry import trace - from opentelemetry.sdk.resources import Resource - from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import SpanKind if config is None: config = OpenTelemetryConfig.from_env() self.config = config + self.callback_name = callback_name self.OTEL_EXPORTER = self.config.exporter self.OTEL_ENDPOINT = self.config.endpoint self.OTEL_HEADERS = self.config.headers - provider = TracerProvider(resource=Resource(attributes=LITELLM_RESOURCE)) - provider.add_span_processor(self._get_span_processor()) self.callback_name = callback_name - - trace.set_tracer_provider(provider) - self.tracer = trace.get_tracer(LITELLM_TRACER_NAME) - self.span_kind = SpanKind - _debug_otel = str(os.getenv("DEBUG_OTEL", "False")).lower() - - if _debug_otel == "true": + if self.config.debug == "true": # Set up logging import logging @@ -115,6 +115,16 @@ class OpenTelemetry(CustomLogger): otel_exporter_logger = logging.getLogger("opentelemetry.sdk.trace.export") otel_exporter_logger.setLevel(logging.DEBUG) + # Don't override a tracer provider already set by the user + if trace.get_tracer_provider() is None: + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + + provider = TracerProvider(resource=Resource(attributes=LITELLM_RESOURCE)) + provider.add_span_processor(self._get_span_processor()) + trace.set_tracer_provider(provider) + self.tracer = trace.get_tracer(LITELLM_TRACER_NAME) + # init CustomLogger params super().__init__(**kwargs) self._init_otel_logger_on_litellm_proxy() @@ -816,12 +826,6 @@ class OpenTelemetry(CustomLogger): return TraceContextTextMapPropagator().extract(carrier=carrier), None def _get_span_processor(self, dynamic_headers: Optional[dict] = None): - from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( - OTLPSpanExporter as OTLPSpanExporterGRPC, - ) - from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( - OTLPSpanExporter as OTLPSpanExporterHTTP, - ) from opentelemetry.sdk.trace.export import ( BatchSpanProcessor, ConsoleSpanExporter, @@ -843,22 +847,26 @@ class OpenTelemetry(CustomLogger): self.OTEL_EXPORTER, "export" ): # Check if it has the export method that SpanExporter requires verbose_logger.debug( - "OpenTelemetry: intiializing SpanExporter. Value of OTEL_EXPORTER: %s", + "OpenTelemetry: initializing SpanExporter. Value of OTEL_EXPORTER: %s", self.OTEL_EXPORTER, ) return SimpleSpanProcessor(cast(SpanExporter, self.OTEL_EXPORTER)) if self.OTEL_EXPORTER == "console": verbose_logger.debug( - "OpenTelemetry: intiializing console exporter. Value of OTEL_EXPORTER: %s", + "OpenTelemetry: initializing console exporter. Value of OTEL_EXPORTER: %s", self.OTEL_EXPORTER, ) return BatchSpanProcessor(ConsoleSpanExporter()) elif self.OTEL_EXPORTER == "otlp_http": verbose_logger.debug( - "OpenTelemetry: intiializing http exporter. Value of OTEL_EXPORTER: %s", + "OpenTelemetry: initializing http exporter. Value of OTEL_EXPORTER: %s", self.OTEL_EXPORTER, ) + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter as OTLPSpanExporterHTTP, + ) + return BatchSpanProcessor( OTLPSpanExporterHTTP( endpoint=self.OTEL_ENDPOINT, headers=_split_otel_headers @@ -866,9 +874,13 @@ class OpenTelemetry(CustomLogger): ) elif self.OTEL_EXPORTER == "otlp_grpc": verbose_logger.debug( - "OpenTelemetry: intiializing grpc exporter. Value of OTEL_EXPORTER: %s", + "OpenTelemetry: initializing grpc exporter. Value of OTEL_EXPORTER: %s", self.OTEL_EXPORTER, ) + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter as OTLPSpanExporterGRPC, + ) + return BatchSpanProcessor( OTLPSpanExporterGRPC( endpoint=self.OTEL_ENDPOINT, headers=_split_otel_headers @@ -876,7 +888,7 @@ class OpenTelemetry(CustomLogger): ) else: verbose_logger.debug( - "OpenTelemetry: intiializing console exporter. Value of OTEL_EXPORTER: %s", + "OpenTelemetry: initializing console exporter. Value of OTEL_EXPORTER: %s", self.OTEL_EXPORTER, ) return BatchSpanProcessor(ConsoleSpanExporter()) diff --git a/tests/logging_callback_tests/test_opentelemetry_unit_tests.py b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py index b0d09562c5..b914bef914 100644 --- a/tests/logging_callback_tests/test_opentelemetry_unit_tests.py +++ b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py @@ -3,12 +3,14 @@ # What is this? ## Unit test for presidio pii masking -import sys, os, asyncio, time, random -from datetime import datetime -import traceback +import sys from dotenv import load_dotenv -load_dotenv() +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + import os import asyncio @@ -17,7 +19,8 @@ sys.path.insert( ) # Adds the parent directory to the system path import pytest import litellm -from unittest.mock import patch, MagicMock, AsyncMock + +from unittest.mock import MagicMock, patch from base_test import BaseLoggingCallbackTest from litellm.types.utils import ModelResponse @@ -37,12 +40,31 @@ class TestOpentelemetryUnitTests(BaseLoggingCallbackTest): f"{SpanAttributes.LLM_COMPLETIONS}.1.function_call.name": "get_news", } + @patch("opentelemetry.trace") + def test_sets_tracer_provider_when_none_exists(self, mock_trace): + mock_trace.get_tracer_provider.return_value = None + + OpenTelemetry(config=OpenTelemetryConfig()) + + mock_trace.set_tracer_provider.assert_called_once() + + @patch("opentelemetry.trace") + def test_does_not_override_existing_tracer_provider(self, mock_trace): + existing_tracer_provider = MagicMock() + mock_trace.get_tracer_provider.return_value = existing_tracer_provider + + OpenTelemetry(config=OpenTelemetryConfig()) + + mock_trace.set_tracer_provider.assert_not_called() + @pytest.mark.asyncio async def test_opentelemetry_integration(self): """ Unit test to confirm the parent otel span is ended """ + load_dotenv() + parent_otel_span = MagicMock() litellm.callbacks = ["otel"] @@ -56,3 +78,61 @@ class TestOpentelemetryUnitTests(BaseLoggingCallbackTest): await asyncio.sleep(1) parent_otel_span.end.assert_called_once() + + +class TestOpenTelemetryConfigUnitTests: + + @pytest.mark.parametrize( + "name, env_vars, expected", + [ + ( + "default", + {}, + OpenTelemetryConfig(exporter="console"), + ), + ( + "OTEL_ENDPOINT -> endpoint", + { + "OTEL_EXPORTER": "otlp_http", + "OTEL_ENDPOINT": "http://localhost:4318/v1/traces" + }, + OpenTelemetryConfig(exporter="otlp_http", endpoint="http://localhost:4318/v1/traces"), + ), + ( + "OTEL_EXPORTER=in_memory -> exporter=InMemorySpanExporter", + {"OTEL_EXPORTER": "in_memory"}, + OpenTelemetryConfig(exporter=InMemorySpanExporter), + ), + ( + "OTEL_HEADERS -> headers", + { + "OTEL_HEADERS": "Authorization=Bearer token123" + }, + OpenTelemetryConfig(exporter="console", headers="Authorization=Bearer token123"), + ), + ( + "DEBUG_OTEL=TrUe -> debug=true", + {"DEBUG_OTEL": "TrUe"}, + OpenTelemetryConfig(exporter="console", debug="true"), + ), + ], + ) + def test_env_variable_prioritization(self, name, monkeypatch, env_vars, expected): + # Clear all environment variables + for var in os.environ: + monkeypatch.delenv(var, raising=False) + # Set test-specific environment variables + for key, value in env_vars.items(): + monkeypatch.setenv(key, value) + + # Call the method under test + config = OpenTelemetryConfig.from_env() + + # Validate the results + if isinstance(expected.exporter, type): + assert isinstance(config.exporter, expected.exporter) + else: + assert config.exporter == expected.exporter + + assert config.endpoint == expected.endpoint + assert config.headers == expected.headers