From c5e9e89288fdf064266473d2e9a8c63d3b9222e4 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Sun, 2 Jun 2024 19:49:34 +0900 Subject: [PATCH] remove mocks --- litellm/integrations/opentelemetry.py | 79 +++++++++++++------- litellm/tests/test_proxy_server.py | 101 +++++++++++++++----------- 2 files changed, 113 insertions(+), 67 deletions(-) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index e87590be4..7a325bc80 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1,6 +1,7 @@ +from dataclasses import dataclass, field +from typing import Optional import os -import litellm from litellm.integrations.custom_logger import CustomLogger from opentelemetry import trace @@ -16,19 +17,31 @@ from opentelemetry.sdk.trace.export import ( ConsoleSpanExporter, ) - LITELLM_TRACER_NAME = "litellm" LITELLM_RESOURCE = {"service.name": "litellm"} -MOCK_TRACE_PARENT = {"traceparent": "SOMETHING_FROM_PROXY_REQUEST"} -MOCK_SPAN_NAME = "TODO" + +@dataclass +class OpenTelemetryConfig: + exporter: str = field(default="console") + endpoint: Optional[str] = None + bearer_token: Optional[str] = None + + @classmethod + def from_env(cls): + return cls( + exporter=os.getenv("OTEL_EXPORTER", "console"), + endpoint=os.getenv("OTEL_ENDPOINT"), + bearer_token=os.getenv("OTEL_BEARER_TOKEN"), + ) class OpenTelemetry(CustomLogger): - def __init__(self): + def __init__(self, config=OpenTelemetryConfig.from_env()): + self.config = config provider = TracerProvider(resource=Resource(attributes=LITELLM_RESOURCE)) - provider.add_span_processor(self.get_span_processor()) - + provider.add_span_processor(self._get_span_processor()) + trace.set_tracer_provider(provider) self.tracer = trace.get_tracer(LITELLM_TRACER_NAME) @@ -46,42 +59,56 @@ class OpenTelemetry(CustomLogger): def _handle_sucess(self, kwargs, response_obj, start_time, end_time): span = self.tracer.start_span( - MOCK_SPAN_NAME, + name=self._get_span_name(kwargs), start_time=self._to_ns(start_time), - context=TraceContextTextMapPropagator().extract(carrier=MOCK_TRACE_PARENT), + context=self._get_span_context(kwargs), ) span.set_status(Status(StatusCode.OK)) - self._set_attributes(span, kwargs, response_obj) + self.set_attributes(span, kwargs, response_obj) span.end(end_time=self._to_ns(end_time)) def _handle_failure(self, kwargs, response_obj, start_time, end_time): span = self.tracer.start_span( - MOCK_SPAN_NAME, + name=self._get_span_name(kwargs), start_time=self._to_ns(start_time), - context=TraceContextTextMapPropagator().extract(carrier=MOCK_TRACE_PARENT), + context=self._get_span_context(kwargs), ) span.set_status(Status(StatusCode.ERROR)) - self._set_attributes(span, kwargs, response_obj) + self.set_attributes(span, kwargs, response_obj) span.end(end_time=self._to_ns(end_time)) + def set_attributes(self, span, kwargs, response_obj): + for key in ["model", "api_base", "api_version"]: + if key in kwargs: + span.set_attribute(key, kwargs[key]) + def _to_ns(self, dt): return int(dt.timestamp() * 1e9) - def _set_attributes(self, span, kwargs, response_obj): - keys = ["model", "api_base", "api_version"] - for key in keys: - if key in kwargs: - span.set_attribute("model", kwargs[key]) + def _get_span_name(self, kwargs): + f"litellm-{kwargs.get('call_type', 'completion')}" - def get_span_processor(self): - if litellm.set_verbose: - BatchSpanProcessor(ConsoleSpanExporter()) + def _get_span_context(self, kwargs): + litellm_params = kwargs.get("litellm_params", {}) or {} + proxy_server_request = litellm_params.get("proxy_server_request", {}) or {} + headers = proxy_server_request.get("headers", {}) or {} + traceparent = headers.get("traceparent", None) + + if traceparent is None: + return None else: - BatchSpanProcessor( + carrier = {"traceparent": traceparent} + return TraceContextTextMapPropagator().extract(carrier=carrier) + + def _get_span_processor(self): + if self.config.exporter == "console": + return BatchSpanProcessor(ConsoleSpanExporter()) + elif self.config.exporter == "otlp_http": + return BatchSpanProcessor( OTLPSpanExporter( - endpoint=os.getenv("OTEL_ENDPOINT"), - headers={ - "Authorization": f"Bearer {os.getenv('OTEL_BEARER_TOKEN')}" - }, + endpoint=self.OTEL_ENDPOINT, + headers={"Authorization": f"Bearer {self.OTEL_BEARER_TOKEN}"}, ) ) + else: + return BatchSpanProcessor(ConsoleSpanExporter()) diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index 77692e2ee..753c2ddc7 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -41,49 +41,39 @@ example_completion_result = { { "message": { "content": "Whispers of the wind carry dreams to me.", - "role": "assistant" + "role": "assistant", } } ], } example_embedding_result = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [ - -0.006929283495992422, - -0.005336422007530928, - -4.547132266452536e-05, - -0.024047505110502243, - -0.006929283495992422, - -0.005336422007530928, - -4.547132266452536e-05, - -0.024047505110502243, - -0.006929283495992422, - -0.005336422007530928, - -4.547132266452536e-05, - -0.024047505110502243, - ], - } - ], - "model": "text-embedding-3-small", - "usage": { - "prompt_tokens": 5, - "total_tokens": 5 - } + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": [ + -0.006929283495992422, + -0.005336422007530928, + -4.547132266452536e-05, + -0.024047505110502243, + -0.006929283495992422, + -0.005336422007530928, + -4.547132266452536e-05, + -0.024047505110502243, + -0.006929283495992422, + -0.005336422007530928, + -4.547132266452536e-05, + -0.024047505110502243, + ], + } + ], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 5, "total_tokens": 5}, } example_image_generation_result = { - "created": 1589478378, - "data": [ - { - "url": "https://..." - }, - { - "url": "https://..." - } - ] + "created": 1589478378, + "data": [{"url": "https://..."}, {"url": "https://..."}], } @@ -185,7 +175,9 @@ def test_engines_model_chat_completions(mock_acompletion, client_no_auth): } print("testing proxy server with chat completions") - response = client_no_auth.post("/engines/gpt-3.5-turbo/chat/completions", json=test_data) + response = client_no_auth.post( + "/engines/gpt-3.5-turbo/chat/completions", json=test_data + ) mock_acompletion.assert_called_once_with( model="gpt-3.5-turbo", messages=[ @@ -249,7 +241,9 @@ def test_chat_completion_azure(mock_acompletion, client_no_auth): @mock_patch_acompletion() -def test_openai_deployments_model_chat_completions_azure(mock_acompletion, client_no_auth): +def test_openai_deployments_model_chat_completions_azure( + mock_acompletion, client_no_auth +): global headers try: # Your test data @@ -388,10 +382,10 @@ def test_img_gen(mock_aimage_generation, client_no_auth): response = client_no_auth.post("/v1/images/generations", json=test_data) mock_aimage_generation.assert_called_once_with( - model='dall-e-3', - prompt='A cute baby sea otter', + model="dall-e-3", + prompt="A cute baby sea otter", n=1, - size='1024x1024', + size="1024x1024", metadata=mock.ANY, proxy_server_request=mock.ANY, ) @@ -608,3 +602,28 @@ def test_load_router_config(mock_cache, fake_env_vars): # test_load_router_config() + +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig + + +@mock_patch_acompletion() +def test_otel_with_proxy_server(mock_acompletion, client_no_auth, capsys): + litellm.callbacks = [OpenTelemetry(OpenTelemetryConfig(exporter="console"))] + + data = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]} + + response = client_no_auth.post("/v1/chat/completions", json=data) + mock_acompletion.assert_called_once_with( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id=mock.ANY, + litellm_logging_obj=mock.ANY, + request_timeout=mock.ANY, + specific_deployment=True, + metadata=mock.ANY, + proxy_server_request=mock.ANY, + ) + assert response.status_code == 200 + assert response.json() == example_completion_result + + print(capsys.readouterr())