test - propogate trace IDs across services

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)