diff --git a/docs/my-website/docs/a2a.md b/docs/my-website/docs/a2a.md index d7145e4b83c..7b863f185d0 100644 --- a/docs/my-website/docs/a2a.md +++ b/docs/my-website/docs/a2a.md @@ -193,6 +193,120 @@ The logs show: style={{width: '100%', display: 'block', margin: '2rem auto'}} /> + +## Forwarding LiteLLM Context Headers + +When LiteLLM invokes your A2A agent, it sends special headers that enable: +- **Trace Grouping**: All LLM calls from the same agent execution appear under one trace +- **Agent Spend Tracking**: Costs are attributed to the specific agent + +| Header | Purpose | +|--------|---------| +| `X-LiteLLM-Trace-Id` | Links all LLM calls to the same execution flow | +| `X-LiteLLM-Agent-Id` | Attributes spend to the correct agent | + + +To enable these features, your A2A server must **forward these headers** to any LLM calls it makes back to LiteLLM. + +### Implementation Steps + +**Step 1: Extract headers from incoming A2A request** +```python def get_litellm_headers(request) -> dict: + """Extract X-LiteLLM-* headers from incoming A2A request.""" + all_headers = request.call_context.state.get('headers', {}) + return { + k: v for k, v in all_headers.items() + if k.lower().startswith('x-litellm-') + } +``` + +**Step 2: Forward headers to your LLM calls** +Pass the extracted headers when making calls back to LiteLLM: + + + +```python from openai import OpenAI + +headers = get_litellm_headers(request) + +client = OpenAI( + api_key="sk-your-litellm-key", + base_url="http://localhost:4000", + default_headers=headers, # Forward headers +) + +response = client.chat.completions.create( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello"}] +) +``` + + + + +```python +from langchain_openai import ChatOpenAI + +headers = get_litellm_headers(request) + +llm = ChatOpenAI( + model="gpt-4o", + openai_api_key="sk-your-litellm-key", + base_url="http://localhost:4000", + default_headers=headers, # Forward headers +) +``` + + + +```python +import litellm + +headers = get_litellm_headers(request) + +response = litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + api_base="http://localhost:4000", + extra_headers=headers, # Forward headers +) +``` + + + +```python +import httpx + +headers = get_litellm_headers(request) +headers["Authorization"] = "Bearer sk-your-litellm-key" + +response = httpx.post( + "http://localhost:4000/v1/chat/completions", + headers=headers, + json={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]} +) +``` + + + +### Result + +With header forwarding enabled, you'll see: + +**Trace Grouping in Langfuse:** + + + +**Agent Spend Attribution:** + + + ## API Reference ### Endpoint diff --git a/docs/my-website/img/a2a_agent_spend.png b/docs/my-website/img/a2a_agent_spend.png new file mode 100644 index 00000000000..15ec769392a Binary files /dev/null and b/docs/my-website/img/a2a_agent_spend.png differ diff --git a/docs/my-website/img/a2a_trace_grouping.png b/docs/my-website/img/a2a_trace_grouping.png new file mode 100644 index 00000000000..05130420aae Binary files /dev/null and b/docs/my-website/img/a2a_trace_grouping.png differ diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 2d36dbeacda..ae95faede22 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -9,7 +9,7 @@ import datetime from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union import litellm -from litellm._logging import verbose_logger +from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.agents import LiteLLMSendMessageResponse from litellm.utils import client +import uuid if TYPE_CHECKING: from a2a.client import A2AClient as A2AClientType @@ -225,7 +226,11 @@ async def asend_message( raise ValueError( "Either a2a_client or api_base is required for standard A2A flow" ) - a2a_client = await create_a2a_client(base_url=api_base) + trace_id = str(uuid.uuid4()) + extra_headers = {"X-LiteLLM-Trace-Id": trace_id} + if agent_id: + extra_headers["X-LiteLLM-Agent-Id"] = agent_id + a2a_client = await create_a2a_client(base_url=api_base, extra_headers=extra_headers) # Type assertion: a2a_client is guaranteed to be non-None here assert a2a_client is not None @@ -490,6 +495,10 @@ async def create_a2a_client( ) httpx_client = http_handler.client + if extra_headers: + httpx_client.headers.update(extra_headers) + verbose_proxy_logger.debug(f"A2A client created with extra_headers={extra_headers}") + # Resolve agent card resolver = A2ACardResolver( httpx_client=httpx_client, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0516a4aaa66..c248082c263 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3421,6 +3421,8 @@ class LitellmMetadataFromRequestHeaders(TypedDict, total=False): """ spend_logs_metadata: Optional[dict] + agent_id: Optional[str] + trace_id: Optional[str] class JWTKeyItem(TypedDict, total=False): diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1d3ef2e10c2..60472cd7790 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -558,6 +558,16 @@ class LiteLLMProxyRequestSetup: ######################################################################################### # Finally update the requests metadata with the `metadata_from_headers` ######################################################################################### + agent_id_from_header = headers.get("x-litellm-agent-id") + trace_id_from_header = headers.get("x-litellm-trace-id") + if agent_id_from_header: + metadata_from_headers["agent_id"] = agent_id_from_header + verbose_proxy_logger.debug(f"Extracted agent_id from header: {agent_id_from_header}") + + if trace_id_from_header: + metadata_from_headers["trace_id"] = trace_id_from_header + verbose_proxy_logger.debug(f"Extracted trace_id from header: {trace_id_from_header}") + if isinstance(data[_metadata_variable_name], dict): data[_metadata_variable_name].update(metadata_from_headers) return data diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 5d348d35cb2..db4e4beec21 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -396,7 +396,7 @@ def get_logging_payload( # noqa: PLR0915 ) # Extract agent_id for A2A requests (set directly on model_call_details) - agent_id: Optional[str] = kwargs.get("agent_id") + agent_id: Optional[str] = kwargs.get("agent_id") or metadata.get("agent_id") custom_llm_provider = kwargs.get("custom_llm_provider") raw_model = cast(str, kwargs.get("model") or "") model_name = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})