scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix: close A2A provider review gaps

Browse source
This commit is contained in:
aiedwardyi 2026-08-26 13:19:41 +09:00
parent 97f01b6928
commit 4556ffc86d
No known key found for this signature in database
8 changed files with 68 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/a2a_protocol/providers/bedrock_agentcore/config.py
View file

@ -38,6 +38,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
params=params,
litellm_params=litellm_params,
agent_extra_headers=kwargs.get("agent_extra_headers"),
timeout=kwargs.get("timeout"),
)
async def handle_streaming(
@ -58,5 +59,6 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
params=params,
litellm_params=litellm_params,
agent_extra_headers=kwargs.get("agent_extra_headers"),
timeout=kwargs.get("timeout"),
):
yield chunk

4
litellm/a2a_protocol/providers/bedrock_agentcore/handler.py
View file

@ -31,6 +31,7 @@ class BedrockAgentCoreA2AHandler:
params: dict[str, Any],
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
timeout: float | None = None,
) -> dict[str, Any]:
"""
Handle non-streaming A2A request to AgentCore.
@ -62,6 +63,7 @@ class BedrockAgentCoreA2AHandler:
url,
headers=headers,
data=body,
timeout=timeout,
)
response.raise_for_status()
response_data: Final = response.json()
@ -77,6 +79,7 @@ class BedrockAgentCoreA2AHandler:
params: dict[str, Any],
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
timeout: float | None = None,
) -> AsyncIterator[dict[str, Any]]:
"""
Handle streaming A2A request to AgentCore.
@ -110,6 +113,7 @@ class BedrockAgentCoreA2AHandler:
headers=headers,
data=body,
stream=True,
timeout=timeout,
)
response.raise_for_status()

11
litellm/a2a_protocol/providers/langflow/config.py
View file

@ -6,6 +6,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.interactions.agents.utils import merge_agent_headers
from litellm.llms.langflow.a2a import merge_a2a_session_into_litellm_params
@ -28,11 +29,16 @@ class LangFlowA2AConfig(BaseA2AProviderConfig):
litellm_params = merge_a2a_session_into_litellm_params(
litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)
)
forwarded_headers = merge_agent_headers(
dynamic_headers=kwargs.get("agent_extra_headers"),
static_headers=kwargs.get("agent_static_headers"),
)
return await A2ACompletionBridgeHandler.handle_non_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=forwarded_headers,
_skip_a2a_provider_routing=True,
)
@ -51,11 +57,16 @@ class LangFlowA2AConfig(BaseA2AProviderConfig):
litellm_params = merge_a2a_session_into_litellm_params(
litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)
)
forwarded_headers = merge_agent_headers(
dynamic_headers=kwargs.get("agent_extra_headers"),
static_headers=kwargs.get("agent_static_headers"),
)
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=forwarded_headers,
_skip_a2a_provider_routing=True,
):
yield chunk

21
litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py
View file

@ -210,12 +210,25 @@ class WatsonxOrchestrateHandler:
client: AsyncHTTPHandler,
max_attempts: int = _MAX_POLL_ATTEMPTS,
interval_s: float = _POLL_INTERVAL_S,
timeout: float | None = None,
) -> _WXORun:
url: Final = f"{base_url}/v1/orchestrate/runs/{run_id}"
deadline: float | None = time.monotonic() + max(timeout, 0) if timeout is not None else None
for attempt in range(max_attempts):
await asyncio.sleep(interval_s)
response = await client.get(url, headers=auth_headers)
if deadline is not None:
remaining = deadline - time.monotonic()
if remaining <= 0:
raise asyncio.TimeoutError(f"WXO run '{run_id}' exceeded timeout of {timeout}s")
if interval_s > 0:
await asyncio.sleep(min(interval_s, remaining))
remaining = deadline - time.monotonic()
if remaining <= 0:
raise asyncio.TimeoutError(f"WXO run '{run_id}' exceeded timeout of {timeout}s")
response = await asyncio.wait_for(client.get(url, headers=auth_headers), timeout=remaining)
else:
await asyncio.sleep(interval_s)
response = await client.get(url, headers=auth_headers)
response.raise_for_status()
result = WatsonxOrchestrateHandler._run_body(response)
status = result.get("status", "")
@ -233,6 +246,7 @@ class WatsonxOrchestrateHandler:
base_url: str,
auth_headers: dict[str, str],
client: AsyncHTTPHandler,
timeout: float | None = None,
) -> _WXORun:
status = run_data.get("status", "")
if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES:
@ -244,6 +258,7 @@ class WatsonxOrchestrateHandler:
run_id=run_id,
auth_headers=auth_headers,
client=client,
timeout=timeout,
)
status = run_data.get("status", "")
@ -341,6 +356,7 @@ class WatsonxOrchestrateHandler:
base_url=base_url,
auth_headers=auth_headers,
client=client,
timeout=timeout,
)
response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(run_data)
@ -418,6 +434,7 @@ class WatsonxOrchestrateHandler:
base_url=base_url,
auth_headers=auth_headers,
client=client,
timeout=timeout,
)
accumulated_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result)
else:

12
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1180,6 +1180,15 @@ class CustomStreamWrapper:
completion_obj: dict[str, Any],
) -> _ProviderChunkResult:
response_obj: dict[str, Any] = {}
if isinstance(chunk, ModelResponseStream) and self.custom_llm_provider == "a2a":
model_response = chunk
model_response.model = self.model
finish_reasons = [getattr(choice, "finish_reason", None) for choice in chunk.choices]
if finish_reasons and all(isinstance(reason, str) and reason for reason in finish_reasons):
self.received_finish_reason = finish_reasons[0]
self.sent_last_chunk = True
return _ProviderChunkEarlyReturn(model_response)
if (
isinstance(chunk, ModelResponseStream)
and self.custom_llm_provider is not None
@ -1219,6 +1228,9 @@ class CustomStreamWrapper:
raise StopIteration
anthropic_response_obj: Final[GChunk] = cast(GChunk, chunk)
completion_obj["content"] = anthropic_response_obj["text"]
chunk_index = anthropic_response_obj.get("index")
if isinstance(chunk_index, int):
model_response.choices[0].index = chunk_index
if anthropic_response_obj["is_finished"]:
self.received_finish_reason = anthropic_response_obj["finish_reason"]

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

@ -850,17 +850,20 @@ async def invoke_agent_a2a(
agent_extra_headers=agent_extra_headers,
)
post_call_succeeded = False
try:
response = await proxy_logging_obj.post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=response,
)
post_call_succeeded = True
finally:
_enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None)
if _enqueue_fn is not None:
logging_obj._enqueue_deferred_logging = None
_enqueue_fn(response)
if post_call_succeeded:
_enqueue_fn(response)
response_dict: Final[dict[str, Any]] = (
response.model_dump(mode="json", exclude_none=True)

15
litellm/proxy/agent_endpoints/a2a_routing.py
View file

@ -459,14 +459,17 @@ def _get_agent_dynamic_headers(
return dynamic_headers
def _get_agent_identity_headers(user_api_key_dict: UserAPIKeyAuth | None) -> dict[str, str]:
if user_api_key_dict is None:
return {}
def _get_agent_identity_headers(
user_api_key_dict: UserAPIKeyAuth | None,
trace_id: object | None = None,
) -> dict[str, str]:
headers: dict[str, str] = {}
if user_api_key_dict.user_id:
if user_api_key_dict is not None and user_api_key_dict.user_id:
headers["X-LiteLLM-User-Id"] = user_api_key_dict.user_id
if user_api_key_dict.team_id:
if user_api_key_dict is not None and user_api_key_dict.team_id:
headers["X-LiteLLM-Team-Id"] = user_api_key_dict.team_id
if trace_id:
headers["X-LiteLLM-Trace-Id"] = str(trace_id)
return headers
@ -585,7 +588,7 @@ async def route_a2a_agent_request(
)
registered_static_headers = merge_agent_headers(
dynamic_headers=registered_static_headers,
static_headers=_get_agent_identity_headers(user_api_key_dict),
static_headers=_get_agent_identity_headers(user_api_key_dict, data.get("litellm_trace_id")),
)
if (
registered_provider

7
litellm/proxy/common_request_processing.py
View file

@ -1890,6 +1890,13 @@ class ProxyBaseLLMRequestProcessing:
llm_router=llm_router,
trust_client_model_info=False,
)
if isinstance(self.data.get("model"), str) and self.data["model"].startswith("a2a/"):
self.data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=self.data,
call_type=route_type,
guardrails_only=True,
)
# Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may
# have mutated `self.data` in place, and the audit-trail snapshot taken in