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

use inmemory exporter for testing

Browse source
This commit is contained in:
yujonglee 2024-06-04 09:04:19 +09:00
parent 61c8e9d2ec
commit 0d8a7d5cf0
2 changed files with 13 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

9
litellm/integrations/opentelemetry.py
View file

@ -1,4 +1,4 @@
from dataclasses import dataclass, field from dataclasses import dataclass
from typing import Optional from typing import Optional
import os import os
@ -18,6 +18,8 @@ from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
OTLPSpanExporter as OTLPSpanExporterGRPC, OTLPSpanExporter as OTLPSpanExporterGRPC,
) )
from opentelemetry.sdk.trace.export import ( from opentelemetry.sdk.trace.export import (
SpanExporter,
SimpleSpanProcessor,
BatchSpanProcessor, BatchSpanProcessor,
ConsoleSpanExporter, ConsoleSpanExporter,
) )
@ -28,7 +30,7 @@ LITELLM_RESOURCE = {"service.name": "litellm"}
@dataclass @dataclass
class OpenTelemetryConfig: class OpenTelemetryConfig:
exporter: str = field(default="console") exporter: str | SpanExporter = "console"
endpoint: Optional[str] = None endpoint: Optional[str] = None
bearer_token: Optional[str] = None bearer_token: Optional[str] = None
@ -106,6 +108,9 @@ class OpenTelemetry(CustomLogger):
return TraceContextTextMapPropagator().extract(carrier=carrier) return TraceContextTextMapPropagator().extract(carrier=carrier)
def _get_span_processor(self): def _get_span_processor(self):
if isinstance(self.config.exporter, SpanExporter):
return SimpleSpanProcessor(self.config.exporter)
if self.config.exporter == "console": if self.config.exporter == "console":
return BatchSpanProcessor(ConsoleSpanExporter()) return BatchSpanProcessor(ConsoleSpanExporter())
elif self.config.exporter == "otlp_http": elif self.config.exporter == "otlp_http":

10
litellm/tests/test_proxy_server.py
View file

@ -604,11 +604,12 @@ def test_load_router_config(mock_cache, fake_env_vars):
# test_load_router_config() # test_load_router_config()
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
@mock_patch_acompletion() @mock_patch_acompletion()
def test_otel_with_proxy_server(mock_acompletion, client_no_auth, capsys): def test_otel_with_proxy_server(mock_acompletion, client_no_auth):
litellm.callbacks = [OpenTelemetry(OpenTelemetryConfig(exporter="console"))] exporter = InMemorySpanExporter()
litellm.callbacks = [OpenTelemetry(OpenTelemetryConfig(exporter=exporter))]
data = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]} data = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]}
@ -626,4 +627,5 @@ def test_otel_with_proxy_server(mock_acompletion, client_no_auth, capsys):
assert response.status_code == 200 assert response.status_code == 200
assert response.json() == example_completion_result assert response.json() == example_completion_result
print(capsys.readouterr()) spans = exporter.get_finished_spans()
assert len(spans) == 0