diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index dcbe7a6506e..8e13cc1b8a6 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -608,6 +608,8 @@ def _build_streaming_logging_obj( request: "SendStreamingMessageRequest", agent_name: str, agent_id: str | None, + litellm_call_id: str, + litellm_trace_id: str, litellm_params: dict[str, object] | None, metadata: dict[str, object] | None, proxy_server_request: dict[str, object] | None, @@ -622,8 +624,9 @@ def _build_streaming_logging_obj( stream=False, call_type="asend_message_streaming", start_time=start_time, - litellm_call_id=str(request.id), - function_id=str(request.id), + litellm_call_id=litellm_call_id, + litellm_trace_id=litellm_trace_id, + function_id=litellm_call_id, ) logging_obj.model = model logging_obj.custom_llm_provider = "a2a_agent" @@ -656,6 +659,8 @@ async def asend_message_streaming( proxy_server_request: dict[str, object] | None = None, agent_extra_headers: dict[str, str] | None = None, agent_name: str | None = None, + litellm_call_id: str | None = None, + litellm_trace_id: str | None = None, **kwargs: object, ) -> AsyncIterator["SendStreamingMessageResponse | Mapping[str, object]"]: """ @@ -725,13 +730,22 @@ async def asend_message_streaming( agent_extra_headers=agent_extra_headers, ) caller_logging_obj: Final = kwargs.get("litellm_logging_obj") + bridge_logging_obj_from_caller: Final = caller_logging_obj if isinstance(caller_logging_obj, Logging) else None + bridge_call_id: Final = litellm_call_id or str(uuid.uuid4()) + bridge_trace_id: Final = ( + (bridge_logging_obj_from_caller.litellm_trace_id if bridge_logging_obj_from_caller else None) + or litellm_trace_id + or str(uuid.uuid4()) + ) bridge_logging_obj: Final = ( - caller_logging_obj - if isinstance(caller_logging_obj, Logging) + bridge_logging_obj_from_caller + if bridge_logging_obj_from_caller is not None else _build_streaming_logging_obj( request=request, agent_name=agent_name or "unknown", agent_id=agent_id, + litellm_call_id=bridge_call_id, + litellm_trace_id=bridge_trace_id, litellm_params=litellm_params, metadata=metadata, proxy_server_request=proxy_server_request, @@ -760,17 +774,17 @@ async def asend_message_streaming( _raw_logging_obj: Final = kwargs.get("litellm_logging_obj") logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None + call_id: Final = litellm_call_id or str(uuid.uuid4()) + trace_id: Final = (logging_obj.litellm_trace_id if logging_obj else None) or litellm_trace_id or str(uuid.uuid4()) if a2a_client is None: if api_base is None: raise ValueError("Either a2a_client or api_base is required for standard A2A flow") - logging_trace_id: Final = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None - trace_id: Final = logging_trace_id or (str(request.id) if request.id else str(uuid.uuid4())) - extra_headers: Final[dict[str, str]] = {"X-LiteLLM-Trace-Id": trace_id} - if agent_id: - extra_headers["X-LiteLLM-Agent-Id"] = agent_id - if agent_extra_headers: - extra_headers.update(agent_extra_headers) + extra_headers: Final = { + **({"X-LiteLLM-Agent-Id": agent_id} if agent_id else {}), + **(agent_extra_headers or {}), + "X-LiteLLM-Trace-Id": trace_id, + } a2a_client = await create_a2a_client( base_url=api_base, extra_headers=extra_headers, @@ -788,6 +802,8 @@ async def asend_message_streaming( request=request, agent_name=card_agent_name, agent_id=agent_id, + litellm_call_id=call_id, + litellm_trace_id=trace_id, litellm_params=litellm_params, metadata=metadata, proxy_server_request=proxy_server_request, diff --git a/litellm/constants.py b/litellm/constants.py index e21751c7fd1..e412284512c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -676,6 +676,7 @@ request_timeout: float = ( ) request_timeout_explicitly_set: bool = _request_timeout_from_env is not None DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = get_env_float("DEFAULT_A2A_AGENT_TIMEOUT", 6000.0) # 10 minutes +A2A_CALL_TYPES: Final = frozenset({"asend_message", "send_message", "asend_message_streaming"}) AGENT_KILL_SWITCH_TIMEOUT_SECONDS: Final = 10.0 AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS: Final = 2000 # Patterns that indicate a localhost/internal URL in A2A agent cards that should be diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b6574a7787d..d79d67df384 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -41,6 +41,7 @@ from litellm.caching.caching import DualCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.caching.redis_batch import flush_post_call_redis_batches from litellm.constants import ( + A2A_CALL_TYPES, DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, @@ -363,6 +364,7 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object]) _RESPONSE_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_STANDARD_LOGGING_ID_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS _UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( ("user_api_key_auth", "user_api_key_budget_reservation") @@ -7742,7 +7744,12 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - id = response_obj.get("id") or kwargs.get("litellm_call_id") + litellm_call_id: Final[str | None] = _STANDARD_LOGGING_ID_ADAPTER.validate_python(kwargs.get("litellm_call_id")) + request_id: Final[str | None] = ( + litellm_call_id + if isinstance(call_type, str) and call_type in A2A_CALL_TYPES + else _STANDARD_LOGGING_ID_ADAPTER.validate_python(response_obj.get("id")) or litellm_call_id + ) _model_id: Final = metadata.get("model_info", {}).get("id", "") _model_group: Final = metadata.get("model_group", "") @@ -7780,7 +7787,7 @@ def get_standard_logging_object_payload( usage_object=usage_dict, proxy_server_request=proxy_server_request, start_time=start_time, - response_id=id, + response_id=request_id, custom_llm_provider=selected_provider if isinstance(selected_provider, str) else None, is_silent_experiment=logging_obj.is_silent_experiment is True, ) @@ -7789,8 +7796,8 @@ def get_standard_logging_object_payload( "user", None ) # maintain backwards compatibility with old request body check - if cache_hit is True or served_from_cache(init_response_obj): - id = f"{id}_cache_hit{time.time()}" # do not duplicate the request id + is_cached: Final = cache_hit is True or served_from_cache(init_response_obj) + id: Final[str | None] = f"{request_id}_cache_hit{time.time()}" if is_cached else request_id uncached_price: Final = logging_obj.uncached_price(result=init_response_obj) if cache_hit is True else None saved_cache_cost: Final = (uncached_price.cost if uncached_price is not None else None) or 0.0 diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 5e1d861e7e9..19b84b4b5be 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -436,6 +436,8 @@ async def _handle_stream_message( request_data: dict[str, object] | None = None, proxy_logging_obj: ProxyLogging | None = None, served_version: A2AVersion = "0.3", + litellm_call_id: str | None = None, + litellm_trace_id: str | None = None, ) -> StreamingResponse: """Handle message/stream method via SDK functions. @@ -506,6 +508,8 @@ async def _handle_stream_message( proxy_server_request=proxy_server_request, agent_extra_headers=agent_extra_headers, agent_name=agent_name, + litellm_call_id=litellm_call_id, + litellm_trace_id=litellm_trace_id, ) if ( @@ -939,6 +943,8 @@ async def invoke_agent_a2a( request_data=data, proxy_logging_obj=proxy_logging_obj, served_version=served_version, + litellm_call_id=logging_obj.litellm_call_id, + litellm_trace_id=logging_obj.litellm_trace_id, ) elif method in { "tasks/get", diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index ba5f873c4e9..781d9298b23 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -16,6 +16,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + A2A_CALL_TYPES, BILLED_RETRY_ATTEMPT_METADATA_KEY, CALLER_SESSION_ID_METADATA_KEY, CLI_SESSION_KEY_PREFIX, @@ -368,7 +369,7 @@ def get_spend_logs_id( ) -> str | None: standard_logging_payload = kwargs.get("standard_logging_object") candidate_ids: Final = ( - response_obj.get("id"), + None if call_type in A2A_CALL_TYPES else response_obj.get("id"), standard_logging_payload.get("id") if isinstance(standard_logging_payload, dict) else None, kwargs.get("litellm_call_id"), ) diff --git a/tests/integration/spend/test_agent_spend.py b/tests/integration/spend/test_agent_spend.py index 0a101fcb0c2..635e99531af 100644 --- a/tests/integration/spend/test_agent_spend.py +++ b/tests/integration/spend/test_agent_spend.py @@ -327,15 +327,18 @@ def _assert_reused_jsonrpc_id_has_distinct_spend_identity( ) for key in keys ) - for response in responses: - assert response.status_code == 200, response.text - if "stream" in method.lower(): - frames: Final = sse_frames(response.text) - assert frames and all(frame["id"] == "1" for frame in frames), response.text - assert all("error" not in frame for frame in frames), response.text - else: - response_body: Final = response.json() - assert response_body["id"] == "1" and "error" not in response_body, response.text + assert all(response.status_code == 200 for response in responses), responses + if "stream" in method.lower(): + frames_by_response: Final = tuple(sse_frames(response.text) for response in responses) + assert all( + frames and all(frame["id"] == "1" for frame in frames) and all("error" not in frame for frame in frames) + for frames in frames_by_response + ), responses + else: + response_bodies: Final = tuple(response.json() for response in responses) + assert all( + response_body["id"] == "1" and "error" not in response_body for response_body in response_bodies + ), responses assert agent.wire is not None upstream_requests: Final = tuple(item for item in agent.wire.drain() if item.method == "POST")