phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-24 18:24:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge 5ec0c9286a into 34be7ffceb

Browse source
This commit is contained in:
Adrian Cole 2025-04-23 16:05:32 +02:00 • committed by GitHub
commit 511b2d84ba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 129 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

76
litellm/integrations/opentelemetry.py
View file

@ -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())

90
tests/logging_callback_tests/test_opentelemetry_unit_tests.py
View file

@ -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