Merge pull request #4133 from BerriAI/litellm_work_with_traceparents

[Feat] OTEL - allow propagating traceparent in headers

docs/my-website/docs/proxy/logging.md

@@ -606,6 +606,52 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
 
 ** 🎉 Expect to see this trace logged in your OTEL collector**
 
+### Context propagation across Services `Traceparent HTTP Header`
+
+❓ Use this when you want to **pass information about the incoming request in a distributed tracing system**
+
+✅ Key change: Pass the **`traceparent` header** in your requests. [Read more about traceparent headers here](https://uptrace.dev/opentelemetry/opentelemetry-traceparent.html#what-is-traceparent-header)
+```curl
+traceparent: 00-80e1afed08e019fc1110464cfa66635c-7a085853722dc6d2-01
+```
+Example Usage
+1. Make Request to LiteLLM Proxy with `traceparent` header
+```python
+import openai
+import uuid
+
+client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
+example_traceparent = f"00-80e1afed08e019fc1110464cfa66635c-02e80198930058d4-01"
+extra_headers = {
+    "traceparent": example_traceparent
+}
+_trace_id = example_traceparent.split("-")[1]
+
+print("EXTRA HEADERS: ", extra_headers)
+print("Trace ID: ", _trace_id)
+
+response = client.chat.completions.create(
+    model="llama3",
+    messages=[
+        {"role": "user", "content": "this is a test request, write a short poem"}
+    ],
+    extra_headers=extra_headers,
+)
+
+print(response)
+
+```
+
+```shell
+# EXTRA HEADERS:  {'traceparent': '00-80e1afed08e019fc1110464cfa66635c-02e80198930058d4-01'}
+# Trace ID:  80e1afed08e019fc1110464cfa66635c
+```
+
+2. Lookup Trace ID on OTEL Logger
+
+Search for Trace=`80e1afed08e019fc1110464cfa66635c` on your OTEL Collector
+
+<Image img={require('../../img/otel_parent.png')} />
 
 
 

litellm/integrations/opentelemetry.py

@@ -454,6 +454,23 @@ class OpenTelemetry(CustomLogger):
     def _get_span_name(self, kwargs):
         return LITELLM_REQUEST_SPAN_NAME
 
+    def get_traceparent_from_header(self, headers):
+        if headers is None:
+            return None
+        _traceparent = headers.get("traceparent", None)
+        if _traceparent is None:
+            return None
+
+        from opentelemetry.trace.propagation.tracecontext import (
+            TraceContextTextMapPropagator,
+        )
+
+        verbose_logger.debug("OpenTelemetry: GOT A TRACEPARENT {}".format(_traceparent))
+        propagator = TraceContextTextMapPropagator()
+        _parent_context = propagator.extract(carrier={"traceparent": _traceparent})
+        verbose_logger.debug("OpenTelemetry: PARENT CONTEXT {}".format(_parent_context))
+        return _parent_context
+
     def _get_span_context(self, kwargs):
         from opentelemetry.trace.propagation.tracecontext import (
             TraceContextTextMapPropagator,

litellm/proxy/proxy_server.py

@@ -533,6 +533,9 @@ async def user_api_key_auth(
             parent_otel_span = open_telemetry_logger.tracer.start_span(
                 name="Received Proxy Server Request",
                 start_time=_to_ns(datetime.now()),
+                context=open_telemetry_logger.get_traceparent_from_header(
+                    headers=request.headers
+                ),
             )
         ### USER-DEFINED AUTH FUNCTION ###
         if user_custom_auth is not None:

litellm/proxy/tests/test_openai_request_with_traceparent.py

@@ -0,0 +1,41 @@
+# mypy: ignore-errors
+import openai
+from opentelemetry import trace
+from opentelemetry.context import Context
+from opentelemetry.trace import SpanKind
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+
+trace.set_tracer_provider(TracerProvider())
+memory_exporter = InMemorySpanExporter()
+span_processor = SimpleSpanProcessor(memory_exporter)
+trace.get_tracer_provider().add_span_processor(span_processor)
+tracer = trace.get_tracer(__name__)
+
+# create an otel traceparent header
+tracer = trace.get_tracer(__name__)
+with tracer.start_as_current_span("ishaan-local-dev-app") as span:
+    span.set_attribute("generation_name", "ishaan-generation-openai-client")
+    client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
+    extra_headers = {}
+    context = trace.set_span_in_context(span)
+    traceparent = TraceContextTextMapPropagator()
+    traceparent.inject(carrier=extra_headers, context=context)
+    print("EXTRA HEADERS: ", extra_headers)
+    _trace_parent = extra_headers.get("traceparent")
+    trace_id = _trace_parent.split("-")[1]
+    print("Trace ID: ", trace_id)
+
+    # # request sent to model set on litellm proxy, `litellm --model`
+    response = client.chat.completions.create(
+        model="llama3",
+        messages=[
+            {"role": "user", "content": "this is a test request, write a short poem"}
+        ],
+        extra_headers=extra_headers,
+    )
+
+    print(response)

litellm/proxy/tests/test_simple_traceparent_openai.py

@@ -0,0 +1,21 @@
+# mypy: ignore-errors
+import openai
+import uuid
+
+client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
+example_traceparent = f"00-80e1afed08e019fc1110464cfa66635c-02e80198930058d4-01"
+extra_headers = {"traceparent": example_traceparent}
+_trace_id = example_traceparent.split("-")[1]
+
+print("EXTRA HEADERS: ", extra_headers)
+print("Trace ID: ", _trace_id)
+
+response = client.chat.completions.create(
+    model="llama3",
+    messages=[
+        {"role": "user", "content": "this is a test request, write a short poem"}
+    ],
+    extra_headers=extra_headers,
+)
+
+print(response)