Manually set attributes

agentlightning/llm_proxy.py

@@ -306,6 +306,13 @@ class LightningSpanExporter(SpanExporter):
         # Buffer append under lock to protect against concurrent exporters.
         with self._ensure_lock():
             for span in spans:
+                from opentelemetry.trace import format_span_id
+
+                print(
+                    f">>> [SPAN] {span.name} {format_span_id(span.context.span_id)} "
+                    f"(parent: {format_span_id(span.parent.span_id) if span.parent is not None else 'N/A'}) "
+                    f"{[key for key in span.attributes.keys() if not key.startswith('llm.request.functions') and not key.startswith('gen_ai.prompt')]}"
+                )
                 self._buffer.append(span)
             default_endpoint = self._otlp_exporter._endpoint  # pyright: ignore[reportPrivateUsage]
             try:
@@ -521,12 +528,30 @@ class LightningOpenTelemetry(OpenTelemetry):
         """The root span is sometimes missing (e.g., when Anthropic endpoint is used).
         It is created in an auth module in LiteLLM. If it's missing, we create it here.
         """
+        print(">>> kwargs keys:", kwargs.keys())
+        headers = kwargs.get("headers", kwargs.get("proxy_server_request", {}).get("headers", {}))
+        print(">>> headers:", headers)
+
         if "metadata" not in kwargs or "litellm_parent_otel_span" not in kwargs["metadata"]:
             parent_otel_span = self.create_litellm_proxy_request_started_span(  # type: ignore
                 start_time=datetime.now(),
-                headers=kwargs.get("headers", {}),
+                headers=headers,
             )
-            updated_metadata = {**kwargs.get("metadata", {}), "litellm_parent_otel_span": parent_otel_span}
+
+            # Sometimes the headers (within metadata) are not set within the litellm internally.
+            # So why don't we just set it here?
+            metadata = kwargs.get("litellm_metadata", {})
+
+            # requester_custom_headers might be NOT set right now.
+            custom_headers = {
+                k: v for k, v in headers.items() if k.startswith("x-") and v is not None and isinstance(v, str)
+            }
+            metadata.setdefault("requester_custom_headers", custom_headers)
+            for key, value in metadata.items():
+                self.safe_set_attribute(span=parent_otel_span, key="metadata.{}".format(key), value=value)
+            print(f">>> [METADATA] {metadata}")
+
+            updated_metadata = {**metadata, "litellm_parent_otel_span": parent_otel_span}
 
             return {**kwargs, "metadata": updated_metadata}
         else:

tests/llm_proxy/test_llm_proxy_gpu.py

@@ -443,7 +443,9 @@ async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer, otlp_enable
     try:
         resource, rollout = await _new_resource(proxy, store)
         adapter = LlmProxyTraceToTriplet()
-        client = anthropic.Anthropic(base_url=resource.endpoint, api_key="token-abc123", timeout=120)
+        client = anthropic.Anthropic(
+            base_url=resource.endpoint, api_key="token-abc123", timeout=120, default_headers={"x-test": "test"}
+        )
 
         # non-stream
         response = client.messages.create(
@@ -487,6 +489,9 @@ async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer, otlp_enable
                 assert "llm.hosted_vllm.choices" in span.attributes
             if span.name == "litellm_request":
                 assert "gen_ai.completion.0.content" in span.attributes
+            if span.name == "Received Proxy Server Request":
+                assert "metadata.requester_custom_headers" in span.attributes
+                assert "x-test" in span.attributes["metadata.requester_custom_headers"]  # type: ignore
         assert len(spans) > 0
         triplets = adapter.adapt(spans)
         for i, triplet in enumerate(triplets):