scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 67

feat: set stable contextid for a2a calls - allows for easily passing to downstream llm/mcp calls

Browse source
This commit is contained in:
Krrish Dholakia 2026-03-02 21:00:33 -08:00
parent dcb4a4a5eb
commit 00e861b90b
3 changed files with 65 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

89
litellm/a2a_protocol/main.py
View file

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

22
litellm/proxy/guardrails/guardrail_endpoints.py
View file

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

1
litellm/proxy/litellm_pre_call_utils.py
View file

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