Merge 5ec0c9286a into 34be7ffceb

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******"}
         """
+
+        # 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,
             )
 
-        if os.getenv("OTEL_EXPORTER") == "in_memory":
             return cls(exporter=InMemorySpanExporter())
+
         return cls(
-            exporter=os.getenv("OTEL_EXPORTER", "console"),
+            exporter=exporter,
-            endpoint=os.getenv("OTEL_ENDPOINT"),
+            endpoint=endpoint,
-            headers=os.getenv(
+            headers=headers,
-                "OTEL_HEADERS"
+            debug=str(debug).lower(),
-            ),  # example: OTEL_HEADERS=x-honeycomb-team=B85YgLm96***"
         )
 
 
@@ -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 self.config.debug == "true":
-
-        if _debug_otel == "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())

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
+import sys
-from datetime import datetime
-import traceback
 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