scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(a2a): use LiteLLM-owned call IDs

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Moe Khalil 2026-10-31 19:35:09 +00:00
parent 90381f7ced
commit 7e53cd3c84
6 changed files with 59 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
litellm/a2a_protocol/main.py
View file

@ -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,

1
litellm/constants.py
View file

@ -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

15
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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

6
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -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",

3
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -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"),
)

21
tests/integration/spend/test_agent_spend.py
View file

@ -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")