Merge pull request #3974 from yujonglee/simple-otel

Add simple OpenTelemetry tracer
2026-10-05 02:41:56 +00:00 · 2024-06-04 08:30:28 -07:00 · 2024-06-04 08:30:28 -07:00 · 009c63f194
commit 009c63f194
parent 5d3674b63d 5b4b4d9a67
5 changed files with 203 additions and 51 deletions
--- a/litellm/integrations/opentelemetry.py
+++ b/litellm/integrations/opentelemetry.py
@ -0,0 +1,139 @@
+import os
+from typing import Optional
+from dataclasses import dataclass
+
+from litellm.integrations.custom_logger import CustomLogger
+
+
+LITELLM_TRACER_NAME = "litellm"
+LITELLM_RESOURCE = {"service.name": "litellm"}
+
+
+@dataclass
+class OpenTelemetryConfig:
+    from opentelemetry.sdk.trace.export import SpanExporter
+
+    exporter: str | SpanExporter = "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, config=OpenTelemetryConfig.from_env()):
+        from opentelemetry import trace
+        from opentelemetry.sdk.resources import Resource
+        from opentelemetry.sdk.trace import TracerProvider
+
+        self.config = config
+        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)
+
+    def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self._handle_sucess(kwargs, response_obj, start_time, end_time)
+
+    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self._handle_failure(kwargs, response_obj, start_time, end_time)
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self._handle_sucess(kwargs, response_obj, start_time, end_time)
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self._handle_failure(kwargs, response_obj, start_time, end_time)
+
+    def _handle_sucess(self, kwargs, response_obj, start_time, end_time):
+        from opentelemetry.trace import Status, StatusCode
+
+        span = self.tracer.start_span(
+            name=self._get_span_name(kwargs),
+            start_time=self._to_ns(start_time),
+            context=self._get_span_context(kwargs),
+        )
+        span.set_status(Status(StatusCode.OK))
+        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):
+        from opentelemetry.trace import Status, StatusCode
+
+        span = self.tracer.start_span(
+            name=self._get_span_name(kwargs),
+            start_time=self._to_ns(start_time),
+            context=self._get_span_context(kwargs),
+        )
+        span.set_status(Status(StatusCode.ERROR))
+        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 _get_span_name(self, kwargs):
+        return f"litellm-{kwargs.get('call_type', 'completion')}"
+
+    def _get_span_context(self, kwargs):
+        from opentelemetry.trace.propagation.tracecontext import (
+            TraceContextTextMapPropagator,
+        )
+
+        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:
+            carrier = {"traceparent": traceparent}
+            return TraceContextTextMapPropagator().extract(carrier=carrier)
+
+    def _get_span_processor(self):
+        from opentelemetry.sdk.trace.export import (
+            SpanExporter,
+            SimpleSpanProcessor,
+            BatchSpanProcessor,
+            ConsoleSpanExporter,
+        )
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+            OTLPSpanExporter as OTLPSpanExporterHTTP,
+        )
+        from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
+            OTLPSpanExporter as OTLPSpanExporterGRPC,
+        )
+
+        if isinstance(self.config.exporter, SpanExporter):
+            return SimpleSpanProcessor(self.config.exporter)
+
+        if self.config.exporter == "console":
+            return BatchSpanProcessor(ConsoleSpanExporter())
+        elif self.config.exporter == "otlp_http":
+            return BatchSpanProcessor(
+                OTLPSpanExporterHTTP(
+                    endpoint=self.OTEL_ENDPOINT,
+                    headers={"Authorization": f"Bearer {self.OTEL_BEARER_TOKEN}"},
+                )
+            )
+        elif self.config.exporter == "otlp_grpc":
+            return BatchSpanProcessor(
+                OTLPSpanExporterGRPC(
+                    endpoint=self.OTEL_ENDPOINT,
+                    headers={"Authorization": f"Bearer {self.OTEL_BEARER_TOKEN}"},
+                )
+            )
+        else:
+            return BatchSpanProcessor(ConsoleSpanExporter())
--- a/litellm/tests/test_opentelemetry.py
+++ b/litellm/tests/test_opentelemetry.py
@ -0,0 +1,26 @@
+import asyncio
+import litellm
+
+from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+
+def test_otel_callback():
+    exporter = InMemorySpanExporter()
+
+    litellm.callbacks = [OpenTelemetry(OpenTelemetryConfig(exporter=exporter))]
+
+    litellm.completion(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+    )
+
+    asyncio.run(
+        litellm.acompletion(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "hi"}],
+        )
+    )
+
+    spans = exporter.get_finished_spans()
+    assert len(spans) == 1 + 1
--- 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,
        )
--- a/poetry.lock
+++ b/poetry.lock
@ -2114,7 +2114,6 @@ files = [
    {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"},
    {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"},
    {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"},
-    {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"},
    {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"},
    {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"},
    {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"},
@ -2122,15 +2121,8 @@ files = [
    {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"},
    {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"},
    {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"},
-    {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"},
    {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"},
    {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
-    {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
-    {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
-    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
-    {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
-    {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
-    {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"},
    {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"},
    {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"},
    {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"},
@ -2147,7 +2139,6 @@ files = [
    {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"},
    {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"},
    {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"},
-    {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"},
    {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"},
    {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"},
    {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"},
@ -2155,7 +2146,6 @@ files = [
    {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"},
    {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"},
    {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"},
-    {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"},
    {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"},
    {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"},
    {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"},
--- a/requirements.txt
+++ b/requirements.txt
@ -28,6 +28,9 @@ python-multipart==0.0.9 # admin UI
 Pillow==10.3.0
 azure-ai-contentsafety==1.0.0 # for azure content safety
 azure-identity==1.15.0 # for azure content safety
+opentelemetry-api = "^1.25.0"
+opentelemetry-sdk = "^1.25.0"
+opentelemetry-exporter-otlp = "^1.25.0"

 ### LITELLM PACKAGE DEPENDENCIES
 python-dotenv==1.0.0 # for env