From 00e861b90be3c7bd233fcaf0caa72da1a55b1d09 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 2 Mar 2026 21:00:33 -0800 Subject: [PATCH] feat: set stable contextid for a2a calls - allows for easily passing to downstream llm/mcp calls --- litellm/a2a_protocol/main.py | 89 +++++++++++-------- .../proxy/guardrails/guardrail_endpoints.py | 22 +++-- litellm/proxy/litellm_pre_call_utils.py | 1 + 3 files changed, 65 insertions(+), 47 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 642dfaf023c..a8f8fff2ba9 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -24,11 +24,7 @@ from litellm.utils import client if TYPE_CHECKING: from a2a.client import A2AClient as A2AClientType - from a2a.types import ( - AgentCard, - SendMessageRequest, - SendStreamingMessageRequest, - ) + from a2a.types import AgentCard, SendMessageRequest, SendStreamingMessageRequest # Runtime imports with availability check A2A_SDK_AVAILABLE = False @@ -124,13 +120,48 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: litellm_logging_obj.model = model litellm_logging_obj.custom_llm_provider = custom_llm_provider litellm_logging_obj.model_call_details["model"] = model - litellm_logging_obj.model_call_details[ - "custom_llm_provider" - ] = custom_llm_provider + litellm_logging_obj.model_call_details["custom_llm_provider"] = ( + custom_llm_provider + ) return agent_name +async def _send_message_via_completion_bridge( + request: "SendMessageRequest", + custom_llm_provider: str, + api_base: Optional[str], + litellm_params: Dict[str, Any], +) -> LiteLLMSendMessageResponse: + """ + Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore). + + Requires request; api_base is optional for providers that derive endpoint from model. + """ + verbose_logger.info( + f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}" + ) + + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + params = ( + request.params.model_dump(mode="json") + if hasattr(request.params, "model_dump") + else dict(request.params) + ) + + response_dict = await A2ACompletionBridgeHandler.handle_non_streaming( + request_id=str(request.id), + params=params, + litellm_params=litellm_params, + api_base=api_base, + ) + + return LiteLLMSendMessageResponse.from_dict(response_dict) + + @client async def asend_message( a2a_client: Optional["A2AClientType"] = None, @@ -193,39 +224,21 @@ async def asend_message( ``` """ litellm_params = litellm_params or {} + logging_obj = kwargs.get("litellm_logging_obj") + trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None custom_llm_provider = litellm_params.get("custom_llm_provider") # Route through completion bridge if custom_llm_provider is set if custom_llm_provider: if request is None: raise ValueError("request is required for completion bridge") - # api_base is optional for providers that derive endpoint from model (e.g., bedrock/agentcore) - - verbose_logger.info( - f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}" - ) - - from litellm.a2a_protocol.litellm_completion_bridge.handler import ( - A2ACompletionBridgeHandler, - ) - - # Extract params from request - params = ( - request.params.model_dump(mode="json") - if hasattr(request.params, "model_dump") - else dict(request.params) - ) - - response_dict = await A2ACompletionBridgeHandler.handle_non_streaming( - request_id=str(request.id), - params=params, - litellm_params=litellm_params, + return await _send_message_via_completion_bridge( + request=request, + custom_llm_provider=custom_llm_provider, api_base=api_base, + litellm_params=litellm_params, ) - # Convert to LiteLLMSendMessageResponse - return LiteLLMSendMessageResponse.from_dict(response_dict) - # Standard A2A client flow if request is None: raise ValueError("request is required") @@ -236,11 +249,13 @@ async def asend_message( raise ValueError( "Either a2a_client or api_base is required for standard A2A flow" ) - trace_id = str(uuid.uuid4()) + trace_id = trace_id or 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) + 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 @@ -255,6 +270,8 @@ async def asend_message( ) card_url = getattr(agent_card, "url", None) if agent_card else None + context_id = trace_id or str(uuid.uuid4()) + request.params.message.context_id = context_id # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL a2a_response = None for _ in range(2): # max 2 attempts: original + 1 retry @@ -606,7 +623,9 @@ async def create_a2a_client( if extra_headers: httpx_client.headers.update(extra_headers) - verbose_proxy_logger.debug(f"A2A client created with extra_headers={extra_headers}") + verbose_proxy_logger.debug( + f"A2A client created with extra_headers={extra_headers}" + ) # Resolve agent card resolver = A2ACardResolver( diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 5215fca0293..dd809e4882b 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1170,11 +1170,10 @@ def _build_field_dict( # Determine the field type from annotation field_type = _get_field_type_from_annotation(field_annotation) - # Check for custom UI type override (ui_type preferred; "type" leaks into OpenAPI and breaks schema) - field_json_schema_extra = getattr(field, "json_schema_extra", {}) or {} + # Check for custom UI type override + field_json_schema_extra = getattr(field, "json_schema_extra", {}) if field_json_schema_extra and "ui_type" in field_json_schema_extra: - ut = field_json_schema_extra["ui_type"] - field_type = ut if isinstance(ut, str) else getattr(ut, "value", ut) + field_type = field_json_schema_extra["ui_type"].value elif field_json_schema_extra and "type" in field_json_schema_extra: field_type = field_json_schema_extra["type"] @@ -1356,9 +1355,9 @@ async def get_provider_specific_params(): lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel) tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel) - tool_permission_fields[ - "ui_friendly_name" - ] = ToolPermissionGuardrailConfigModel.ui_friendly_name() + tool_permission_fields["ui_friendly_name"] = ( + ToolPermissionGuardrailConfigModel.ui_friendly_name() + ) # Return the provider-specific parameters provider_params = { @@ -1497,7 +1496,6 @@ async def test_custom_code_guardrail( ``` """ - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=403, @@ -1632,10 +1630,10 @@ async def apply_guardrail( from litellm.proxy.utils import handle_exception_on_proxy try: - active_guardrail: Optional[ - CustomGuardrail - ] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( - guardrail_name=request.guardrail_name + active_guardrail: Optional[CustomGuardrail] = ( + GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( + guardrail_name=request.guardrail_name + ) ) if active_guardrail is None: raise HTTPException( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index dbc4ce7668a..ca21d0430a5 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -577,6 +577,7 @@ class LiteLLMProxyRequestSetup: # Finally update the requests metadata with the `metadata_from_headers` ######################################################################################### + agent_id_from_header = headers.get("x-litellm-agent-id") # x-litellm-trace-id and x-litellm-session-id are interchangeable for call chaining chain_id = headers.get("x-litellm-trace-id") or headers.get(