scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Merge 0cbf8802c1 into 6a3333d3c8

Browse source
This commit is contained in:
Edward Yi 2026-08-29 07:16:20 -04:00 • committed by GitHub
commit f84869da03
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 2680 additions and 143 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

243
litellm/a2a_protocol/litellm_completion_bridge/handler.py
View file

@ -37,6 +37,7 @@ _AGENT_ONLY_PARAMS: Final = frozenset(
"agent_id",
"agent_card_params",
A2A_USER_API_KEY_HASH_PARAM,
"databricks_oauth",
}
)
@ -46,6 +47,19 @@ class A2ACompletionBridgeHandler:
Static methods for handling A2A requests via LiteLLM completion.
"""
@staticmethod
def _merge_stream_values(previous: object, current: object) -> object:
if isinstance(previous, Mapping) and isinstance(current, Mapping):
merged = dict(previous)
for key, value in current.items():
merged[key] = (
A2ACompletionBridgeHandler._merge_stream_values(merged[key], value) if key in merged else value
)
return merged
if isinstance(previous, list) and isinstance(current, list):
return [*previous, *current]
return current
@staticmethod
def _build_completion_params(
params: dict[str, Any],
@ -59,7 +73,12 @@ class A2ACompletionBridgeHandler:
message: Final = params.get("message", {})
# Transform A2A message to OpenAI format
openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
supplied_messages: Final = params.get("messages")
openai_messages: Final = (
supplied_messages
if isinstance(supplied_messages, list)
else A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
@ -84,11 +103,13 @@ class A2ACompletionBridgeHandler:
"api_base": api_base,
"stream": stream,
}
configured_headers: Final[object] = litellm_params.get("extra_headers") or litellm_params.get("headers")
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add: Final = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
if k not in ("model", "custom_llm_provider", "extra_headers", "headers", "api_base", "stream")
and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so the helper
@ -100,10 +121,10 @@ class A2ACompletionBridgeHandler:
params=params,
)
if agent_extra_headers:
if agent_extra_headers or configured_headers:
completion_params["extra_headers"] = merge_agent_headers(
dynamic_headers=agent_extra_headers,
static_headers=completion_params.get("extra_headers"),
static_headers=configured_headers if isinstance(configured_headers, Mapping) else None,
)
return completion_params
@ -122,6 +143,7 @@ class A2ACompletionBridgeHandler:
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
agent_static_headers: Mapping[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> dict[str, object]:
@ -135,6 +157,7 @@ class A2ACompletionBridgeHandler:
api_base: API base URL from agent_card_params
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
agent_static_headers: Configured headers for provider-specific routing.
Returns:
A2A SendMessageResponse dict
@ -149,13 +172,24 @@ class A2ACompletionBridgeHandler:
if a2a_provider_config is not None:
verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider)
return await a2a_provider_config.handle_non_streaming(
request_id=request_id,
params=params,
api_base=api_base,
litellm_params=litellm_params,
agent_extra_headers=agent_extra_headers,
)
provider_params: Final = dict(params)
if custom_llm_provider == "pydantic_ai_agents" or (
custom_llm_provider == "bedrock"
and isinstance(litellm_params.get("model"), str)
and "agentcore" in litellm_params["model"]
):
provider_params.pop("messages", None)
provider_kwargs: Final[dict[str, Any]] = {
"request_id": request_id,
"params": provider_params,
"api_base": api_base,
"litellm_params": litellm_params,
"agent_extra_headers": agent_extra_headers,
"agent_static_headers": agent_static_headers,
}
if litellm_params.get("timeout") is not None:
provider_kwargs["timeout"] = litellm_params["timeout"]
return await a2a_provider_config.handle_non_streaming(**provider_kwargs)
completion_params: Final = A2ACompletionBridgeHandler._build_completion_params(
params=params,
@ -185,6 +219,7 @@ class A2ACompletionBridgeHandler:
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
agent_static_headers: Mapping[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> AsyncIterator[dict[str, object]]:
@ -204,6 +239,7 @@ class A2ACompletionBridgeHandler:
api_base: API base URL from agent_card_params
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
agent_static_headers: Configured headers for provider-specific routing.
Yields:
A2A streaming response events
@ -218,14 +254,31 @@ class A2ACompletionBridgeHandler:
if a2a_provider_config is not None:
verbose_logger.info("A2A: Using provider config for %s (streaming)", custom_llm_provider)
async for chunk in a2a_provider_config.handle_streaming(
request_id=request_id,
params=params,
api_base=api_base,
litellm_params=litellm_params,
agent_extra_headers=agent_extra_headers,
provider_params: Final = dict(params)
if custom_llm_provider == "pydantic_ai_agents" or (
custom_llm_provider == "bedrock"
and isinstance(litellm_params.get("model"), str)
and "agentcore" in litellm_params["model"]
):
yield chunk
provider_params.pop("messages", None)
provider_kwargs: Final[dict[str, Any]] = {
"request_id": request_id,
"params": provider_params,
"api_base": api_base,
"litellm_params": litellm_params,
"agent_extra_headers": agent_extra_headers,
"agent_static_headers": agent_static_headers,
}
if litellm_params.get("timeout") is not None:
provider_kwargs["timeout"] = litellm_params["timeout"]
provider_stream: Final = a2a_provider_config.handle_streaming(**provider_kwargs)
try:
async for chunk in provider_stream:
yield chunk
finally:
close_provider_stream = getattr(provider_stream, "aclose", None)
if close_provider_stream is not None:
await close_provider_stream()
return
@ -259,29 +312,104 @@ class A2ACompletionBridgeHandler:
# Call litellm.acompletion with streaming
response: Final = await A2ACompletionBridgeHandler._acompletion(completion_params)
# 3. Accumulate content and emit artifact update
accumulated_text = ""
# 3. Forward content as artifact updates
accumulated_tool_calls: Final[list[object]] = [] # mutable-ok: collect streaming tool-call deltas
choice_texts: dict[int, str] = {}
choice_tool_calls: dict[int, list[object]] = {}
choice_delta_fields: dict[int, dict[str, object]] = {}
choice_logprobs: dict[int, dict[str, object]] = {}
choice_finish_reasons: dict[int, str] = {}
stream_metadata: dict[str, str] = {}
stream_usage: object | None = None
stream_finish_reason: str | None = None
chunk_count = 0
async for chunk in response:
chunk_count += 1
try:
async for chunk in response:
chunk_count += 1
# Extract delta content
content = ""
if chunk is not None and hasattr(chunk, "choices") and chunk.choices:
choice = chunk.choices[0]
if hasattr(choice, "delta") and choice.delta:
content = choice.delta.content or ""
raw_usage = getattr(chunk, "usage", None)
if isinstance(raw_usage, Mapping):
stream_usage = raw_usage
else:
dump_usage = getattr(raw_usage, "model_dump", None)
if callable(dump_usage):
dumped_usage = dump_usage(exclude_none=True)
if isinstance(dumped_usage, Mapping):
stream_usage = dumped_usage
else:
dict_usage = getattr(raw_usage, "dict", None)
if callable(dict_usage):
dumped_usage = dict_usage(exclude_none=True)
if isinstance(dumped_usage, Mapping):
stream_usage = dumped_usage
if content:
accumulated_text += content
for metadata_name in ("system_fingerprint", "service_tier"):
metadata_value = getattr(chunk, metadata_name, None)
if not isinstance(metadata_value, str):
chunk_fields = A2ACompletionBridgeTransformation._model_dump(chunk)
metadata_value = chunk_fields.get(metadata_name)
if isinstance(metadata_value, str) and metadata_value:
stream_metadata[metadata_name] = metadata_value
# Emit artifact update with accumulated content
if accumulated_text:
artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event(
ctx=ctx,
text=accumulated_text,
)
yield artifact_event
# Extract delta content
choices = getattr(chunk, "choices", None) if chunk is not None else None
if isinstance(choices, (list, tuple)):
for choice_position, choice in enumerate(choices):
raw_index = getattr(choice, "index", choice_position)
choice_index = raw_index if isinstance(raw_index, int) else choice_position
choice_texts.setdefault(choice_index, "")
raw_finish_reason = getattr(choice, "finish_reason", None)
if isinstance(raw_finish_reason, str) and raw_finish_reason:
choice_finish_reasons[choice_index] = raw_finish_reason
if choice_index == 0 or stream_finish_reason is None:
stream_finish_reason = raw_finish_reason
content = ""
delta = getattr(choice, "delta", None)
if delta:
raw_content = getattr(delta, "content", None)
content = raw_content if isinstance(raw_content, str) else ""
choice_texts[choice_index] += content
tool_calls = getattr(delta, "tool_calls", None)
if isinstance(tool_calls, (list, tuple)):
accumulated_tool_calls.extend(tool_calls)
choice_tool_calls.setdefault(choice_index, []).extend(tool_calls)
delta_fields = A2ACompletionBridgeTransformation._model_dump(delta)
if delta_fields:
choice_fields = choice_delta_fields.setdefault(choice_index, {})
for field, value in delta_fields.items():
if field in {"content", "role", "tool_calls"} or value is None:
continue
previous = choice_fields.get(field)
if (isinstance(previous, str) and isinstance(value, str)) or (
isinstance(previous, list) and isinstance(value, list)
):
choice_fields[field] = previous + value
elif isinstance(previous, Mapping) and isinstance(value, Mapping):
choice_fields[field] = {**previous, **value}
else:
choice_fields[field] = value
raw_logprobs = getattr(choice, "logprobs", None)
serialized_logprobs = A2ACompletionBridgeTransformation._model_dump(raw_logprobs)
if serialized_logprobs:
previous_logprobs = choice_logprobs.get(choice_index, {})
merged_logprobs = A2ACompletionBridgeHandler._merge_stream_values(
previous_logprobs, serialized_logprobs
)
if isinstance(merged_logprobs, dict):
choice_logprobs[choice_index] = merged_logprobs
if content:
artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event(
ctx=ctx,
text=content,
index=choice_index,
)
yield artifact_event
finally:
close_response = getattr(response, "aclose", None)
if close_response is not None:
await close_response()
# 4. Emit final status update (kind: "status-update", status: "completed", final: true)
completed_event: Final = A2ACompletionBridgeTransformation.create_status_update_event(
@ -289,6 +417,47 @@ class A2ACompletionBridgeHandler:
state="completed",
final=True,
)
if accumulated_tool_calls:
completed_event["result"]["tool_calls"] = accumulated_tool_calls
if stream_finish_reason:
completed_event["result"]["finish_reason"] = stream_finish_reason
if stream_usage is not None:
completed_event["usage"] = stream_usage
for metadata_name, metadata_value in stream_metadata.items():
completed_event[metadata_name] = metadata_value
choice_indices = sorted(
set(choice_texts)
| set(choice_tool_calls)
| set(choice_delta_fields)
| set(choice_logprobs)
| set(choice_finish_reasons)
)
if choice_indices:
choice_payloads: list[dict[str, object]] = []
for choice_index in choice_indices:
choice_payload: dict[str, object] = {
"index": choice_index,
"message": {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": ""}],
**(
{"tool_calls": choice_tool_calls[choice_index]}
if choice_tool_calls.get(choice_index)
else {}
),
},
**(
{"finish_reason": choice_finish_reasons[choice_index]}
if choice_index in choice_finish_reasons
else {}
),
**({"logprobs": choice_logprobs[choice_index]} if choice_index in choice_logprobs else {}),
}
if choice_delta_fields.get(choice_index):
choice_payload["delta"] = choice_delta_fields[choice_index]
choice_payloads.append(choice_payload)
completed_event["result"]["choices"] = choice_payloads
yield completed_event
verbose_logger.info(

125
litellm/a2a_protocol/litellm_completion_bridge/transformation.py
View file

@ -34,6 +34,7 @@ class A2AStreamingContext:
self.request_id = request_id
self.task_id = str(uuid4())
self.context_id = str(uuid4())
self.artifact_id = str(uuid4())
self.input_message = input_message
self.accumulated_text = ""
self.has_emitted_task = False
@ -151,6 +152,22 @@ class A2ACompletionBridgeTransformation:
return [openai_message]
@staticmethod
def _model_dump(value: Any) -> dict[str, Any]:
if isinstance(value, dict):
return value
dump = getattr(value, "model_dump", None)
if callable(dump):
dumped = dump(exclude_none=True)
if isinstance(dumped, dict):
return dumped
dump = getattr(value, "dict", None)
if callable(dump):
dumped = dump(exclude_none=True)
if isinstance(dumped, dict):
return dumped
return {}
@staticmethod
def openai_response_to_a2a_response(
response: Any,
@ -166,20 +183,78 @@ class A2ACompletionBridgeTransformation:
Returns:
A2A SendMessageResponse dict
"""
# Extract content from response
content = ""
if hasattr(response, "choices") and response.choices:
choice: Final = response.choices[0]
if hasattr(choice, "message") and choice.message:
content = choice.message.content or ""
serialized_choices: list[dict[str, Any]] = []
raw_choices: Final = getattr(response, "choices", None)
if raw_choices:
for choice in raw_choices:
raw_message = getattr(choice, "message", None)
message_fields: Final = A2ACompletionBridgeTransformation._model_dump(raw_message)
raw_content = message_fields.get("content")
if raw_content is None:
raw_content = getattr(raw_message, "content", None)
content: Final = raw_content if isinstance(raw_content, str) else ""
message: Final = {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": content}],
"messageId": uuid4().hex,
}
raw_tool_calls = message_fields.get("tool_calls")
if raw_tool_calls:
message["tool_calls"] = [
call.model_dump(exclude_none=True)
if hasattr(call, "model_dump")
else call.dict(exclude_none=True)
if hasattr(call, "dict")
else call
for call in raw_tool_calls
]
for field in (
"annotations",
"audio",
"function_call",
"images",
"provider_specific_fields",
"reasoning_content",
"reasoning_items",
"refusal",
"thinking_blocks",
):
value = message_fields.get(field)
if value is not None:
message[field] = value
choice_fields: Final = A2ACompletionBridgeTransformation._model_dump(choice)
finish_reason: Final = choice_fields.get("finish_reason")
if finish_reason is None:
raw_finish_reason = getattr(choice, "finish_reason", None)
finish_reason = raw_finish_reason if isinstance(raw_finish_reason, str) else None
if finish_reason:
message["finish_reason"] = finish_reason
choice_payload: Final[dict[str, Any]] = {
"index": len(serialized_choices),
"message": message,
}
logprobs = choice_fields.get("logprobs")
if logprobs is None:
raw_logprobs = getattr(choice, "logprobs", None)
logprobs = raw_logprobs if isinstance(raw_logprobs, dict) else None
if logprobs is not None:
choice_payload["logprobs"] = logprobs
message["logprobs"] = logprobs
serialized_choices.append(choice_payload)
# Build A2A message
a2a_message: Final = {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": content}],
"messageId": uuid4().hex,
}
a2a_message: Final = (
serialized_choices[0]["message"]
if serialized_choices
else {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": ""}],
"messageId": uuid4().hex,
}
)
usage: Final = getattr(response, "usage", None)
# Build A2A response
a2a_response: Final = {
@ -187,8 +262,16 @@ class A2ACompletionBridgeTransformation:
"id": request_id,
"result": a2a_message,
}
if usage is not None:
a2a_response["usage"] = usage.model_dump(exclude_none=True) if hasattr(usage, "model_dump") else usage
for field in ("system_fingerprint", "service_tier"):
value = getattr(response, field, None)
if value is not None:
a2a_response[field] = value
if len(serialized_choices) > 1:
a2a_response["choices"] = serialized_choices
verbose_logger.debug("OpenAI -> A2A transform: content_length=%s", len(content))
verbose_logger.debug("OpenAI -> A2A transform: content_length=%s", len(a2a_message["parts"][0]["text"]))
return a2a_response
@ -277,6 +360,7 @@ class A2ACompletionBridgeTransformation:
def create_artifact_update_event(
ctx: A2AStreamingContext,
text: str,
index: int | None = None,
) -> dict[str, Any]:
"""
Create an artifact update event with content.
@ -285,15 +369,18 @@ class A2ACompletionBridgeTransformation:
ctx: Streaming context
text: The text content for the artifact
"""
artifact: Final[dict[str, Any]] = {
"artifactId": ctx.artifact_id,
"name": "response",
"parts": [{"kind": "text", "text": text}],
}
if index is not None:
artifact["index"] = index
return {
"id": ctx.request_id,
"jsonrpc": "2.0",
"result": {
"artifact": {
"artifactId": str(uuid4()),
"name": "response",
"parts": [{"kind": "text", "text": text}],
},
"artifact": artifact,
"contextId": ctx.context_id,
"kind": "artifact-update",
"taskId": ctx.task_id,

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

22
litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
View file

@ -6,6 +6,7 @@ This module provides fake streaming by converting non-streaming responses into s
"""
import asyncio
import time
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Any, Final, Protocol, cast, runtime_checkable
from uuid import uuid4
@ -101,6 +102,7 @@ class PydanticAITransformation:
max_attempts: int = 30,
poll_interval: float = 0.5,
agent_extra_headers: dict[str, str] | None = None,
timeout: float | None = None,
) -> dict[str, object]:
"""
Poll for task completion using tasks/get method.
@ -112,11 +114,16 @@ class PydanticAITransformation:
request_id: JSON-RPC request ID
max_attempts: Maximum polling attempts
poll_interval: Seconds between poll attempts
timeout: Total polling timeout in seconds
Returns:
Completed task response
"""
deadline: Final[float | None] = time.monotonic() + max(timeout, 0.0) if timeout is not None else None
for attempt in range(max_attempts):
remaining: float | None = deadline - time.monotonic() if deadline is not None else None
if remaining is not None and remaining <= 0:
break
poll_request = {
"jsonrpc": "2.0",
"id": f"{request_id}-poll-{attempt}",
@ -131,6 +138,7 @@ class PydanticAITransformation:
**(agent_extra_headers or {}),
"Content-Type": "application/json",
},
timeout=remaining,
)
response.raise_for_status()
poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
@ -146,9 +154,16 @@ class PydanticAITransformation:
elif state in ("failed", "canceled"):
raise Exception(f"Task {task_id} ended with state: {state}")
await asyncio.sleep(poll_interval)
if deadline is None:
await asyncio.sleep(poll_interval)
continue
remaining = deadline - time.monotonic()
if remaining <= 0:
break
await asyncio.sleep(min(poll_interval, remaining))
raise TimeoutError(f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds")
timeout_description = timeout if timeout is not None else max_attempts * poll_interval
raise TimeoutError(f"Task {task_id} did not complete within {timeout_description} seconds")
@staticmethod
async def _send_and_poll_raw(
@ -203,6 +218,7 @@ class PydanticAITransformation:
llm_provider=cast(Any, "pydantic_ai_agent"),
params={"timeout": timeout},
)
deadline: Final[float | None] = time.monotonic() + max(timeout, 0.0) if timeout is not None else None
response: Final = await client.post(
endpoint,
json=a2a_request,
@ -210,6 +226,7 @@ class PydanticAITransformation:
**(agent_extra_headers or {}),
"Content-Type": "application/json",
},
timeout=timeout,
)
response.raise_for_status()
response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
@ -230,6 +247,7 @@ class PydanticAITransformation:
task_id=task_id,
request_id=request_id,
agent_extra_headers=agent_extra_headers,
timeout=(max(deadline - time.monotonic(), 0.0) if deadline is not None else None),
)
verbose_logger.info("Pydantic AI: Received completed response for request_id=%s", request_id)

13
litellm/a2a_protocol/providers/watsonx_orchestrate/config.py
View file

@ -9,6 +9,7 @@ from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import (
WatsonxOrchestrateHandler,
)
from litellm.interactions.agents.utils import merge_agent_headers
class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
@ -28,10 +29,16 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
"litellm_params is required for WatsonxOrchestrateA2AConfig "
"(must contain cp4d_host, instance_id, wxo_agent_id, api_key)"
)
forwarded_headers: Final = merge_agent_headers(
dynamic_headers=kwargs.get("agent_extra_headers"),
static_headers=kwargs.get("agent_static_headers"),
)
return await WatsonxOrchestrateHandler.handle_non_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
static_headers=forwarded_headers,
timeout=kwargs.get("timeout"),
)
async def handle_streaming(
@ -48,9 +55,15 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
"litellm_params is required for WatsonxOrchestrateA2AConfig "
"(must contain cp4d_host, instance_id, wxo_agent_id, api_key)"
)
forwarded_headers: Final = merge_agent_headers(
dynamic_headers=kwargs.get("agent_extra_headers"),
static_headers=kwargs.get("agent_static_headers"),
)
async for chunk in WatsonxOrchestrateHandler.handle_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
static_headers=forwarded_headers,
timeout=kwargs.get("timeout"),
):
yield chunk

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

@ -6,7 +6,7 @@ import asyncio
import hashlib
import json
import time
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping
from typing import Any, Final, NamedTuple, Protocol
import httpx
@ -26,9 +26,34 @@ _IBM_CLOUD_IAM_URL: Final = "https://iam.cloud.ibm.com/identity/token"
_POLL_INTERVAL_S: Final = 2.0
_MAX_POLL_ATTEMPTS: Final = 90
_TOKEN_CACHE_TTL_BUFFER_S: Final = 60
_WXO_RESERVED_HEADERS: Final = frozenset({"accept", "authorization", "content-type"})
_token_cache: Final[dict[str, tuple[str, float]]] = {}
def _build_wxo_headers(
token: str,
accept: str,
static_headers: Mapping[str, str] | None = None,
) -> dict[str, str]:
headers: dict[str, str] = (
{
key: value
for key, value in static_headers.items()
if isinstance(key, str) and isinstance(value, str) and key.lower() not in _WXO_RESERVED_HEADERS
}
if isinstance(static_headers, Mapping)
else {}
)
headers.update(
{
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"Accept": accept,
}
)
return headers
class WXORequestParams(NamedTuple):
cp4d_host: str
instance_id: str
@ -185,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", "")
@ -208,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:
@ -219,6 +258,7 @@ class WatsonxOrchestrateHandler:
run_id=run_id,
auth_headers=auth_headers,
client=client,
timeout=timeout,
)
status = run_data.get("status", "")
@ -228,23 +268,28 @@ class WatsonxOrchestrateHandler:
return run_data
@staticmethod
async def _accumulate_wxo_sse_text(response: Any) -> str:
source: Final[_WXOView] = {"sse_source": response}
accumulated_text = ""
async for line in source["sse_source"].aiter_lines():
if not line.startswith("data:"):
continue
data_str = line[5:].strip()
if not data_str or data_str == "[DONE]":
continue
try:
event = WatsonxOrchestrateHandler._decode_run_event(data_str)
except json.JSONDecodeError:
continue
chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event)
if chunk_text:
accumulated_text += chunk_text
return accumulated_text
async def _accumulate_wxo_sse_text(response: Any, timeout: float | None = None) -> str:
async def _collect() -> str:
source: Final[_WXOView] = {"sse_source": response}
accumulated_text = ""
async for line in source["sse_source"].aiter_lines():
if not line.startswith("data:"):
continue
data_str = line[5:].strip()
if not data_str or data_str == "[DONE]":
continue
try:
event = WatsonxOrchestrateHandler._decode_run_event(data_str)
except json.JSONDecodeError:
continue
chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event)
if chunk_text:
accumulated_text += chunk_text
return accumulated_text
if timeout is None:
return await _collect()
return await asyncio.wait_for(_collect(), timeout=max(timeout, 0))
@staticmethod
def _extract_litellm_params(litellm_params: WXOLitellmParams) -> WXORequestParams:
@ -277,10 +322,12 @@ class WatsonxOrchestrateHandler:
request_id: str,
params: dict[str, object],
litellm_params: WXOLitellmParams,
static_headers: Mapping[str, str] | None = None,
timeout: float | None = None,
) -> dict[str, object]:
wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client: Final = WatsonxOrchestrateHandler._http_client(timeout=90.0)
client: Final = WatsonxOrchestrateHandler._http_client(timeout=timeout if timeout is not None else 90.0)
token: Final = await WatsonxOrchestrateHandler._get_bearer_token(
cp4d_host=wxo.cp4d_host,
auth_mode=wxo.auth_mode,
@ -289,11 +336,11 @@ class WatsonxOrchestrateHandler:
client=client,
)
base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id)
auth_headers: Final = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"Accept": "application/json",
}
auth_headers: Final = _build_wxo_headers(
token=token,
accept="application/json",
static_headers=static_headers,
)
text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body(
@ -314,6 +361,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)
@ -326,10 +374,12 @@ class WatsonxOrchestrateHandler:
litellm_params: WXOLitellmParams,
chunk_size: int = 50,
delay_ms: int = 10,
static_headers: Mapping[str, str] | None = None,
timeout: float | None = None,
) -> AsyncIterator[dict[str, object]]:
wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client: Final = WatsonxOrchestrateHandler._http_client(timeout=120.0)
client: Final = WatsonxOrchestrateHandler._http_client(timeout=timeout if timeout is not None else 120.0)
token: Final = await WatsonxOrchestrateHandler._get_bearer_token(
cp4d_host=wxo.cp4d_host,
auth_mode=wxo.auth_mode,
@ -338,11 +388,11 @@ class WatsonxOrchestrateHandler:
client=client,
)
base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id)
auth_headers: Final = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"Accept": "text/event-stream, application/json",
}
auth_headers: Final = _build_wxo_headers(
token=token,
accept="text/event-stream, application/json",
static_headers=static_headers,
)
text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body(
wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id
@ -366,6 +416,8 @@ class WatsonxOrchestrateHandler:
request_id=request_id,
params=params,
litellm_params=litellm_params,
static_headers=static_headers,
timeout=timeout,
)
response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(result)
async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text(
@ -387,10 +439,11 @@ class WatsonxOrchestrateHandler:
base_url=base_url,
auth_headers=auth_headers,
client=client,
timeout=timeout,
)
accumulated_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result)
else:
accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response)
accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response, timeout=timeout)
async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text(
text=accumulated_text,

22
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1226,6 +1226,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
@ -1264,7 +1273,13 @@ class CustomStreamWrapper:
if not _chunk_has_content and (not isinstance(chunk, dict) or "provider_specific_fields" not in chunk):
raise StopIteration
anthropic_response_obj: Final[GChunk] = cast(GChunk, chunk)
chunk_id = anthropic_response_obj.get("id")
if isinstance(chunk_id, str) and chunk_id.strip():
model_response = self.set_model_id(chunk_id, model_response)
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"]
@ -1279,7 +1294,8 @@ class CustomStreamWrapper:
)
if "tool_use" in anthropic_response_obj and anthropic_response_obj["tool_use"] is not None:
completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]]
tool_use = anthropic_response_obj["tool_use"]
completion_obj["tool_calls"] = tool_use if isinstance(tool_use, list) else [tool_use]
if (
"provider_specific_fields" in anthropic_response_obj
@ -2611,6 +2627,8 @@ def convert_generic_chunk_to_model_response_stream(
) -> ModelResponseStream:
from litellm.types.utils import Delta
tool_use = chunk.get("tool_use", None)
tool_calls = tool_use if isinstance(tool_use, list) else [tool_use] if tool_use is not None else None
model_response_stream: Final = ModelResponseStream(
id=str(uuid.uuid4()),
model="",
@ -2619,7 +2637,7 @@ def convert_generic_chunk_to_model_response_stream(
index=chunk.get("index", 0),
delta=Delta(
content=chunk["text"],
tool_calls=chunk.get("tool_use", None),
tool_calls=tool_calls,
),
)
],

235
litellm/llms/a2a/chat/streaming_iterator.py
View file

@ -2,12 +2,16 @@
A2A Streaming Response Iterator
"""
from typing import Final
from collections.abc import Mapping
from typing import Any, Final
from uuid import uuid4
import litellm
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
from litellm.types.llms.openai import ChatCompletionToolCallChunk
from litellm.types.utils import Delta, GenericStreamingChunk, ModelResponseStream, StreamingChoices
from ..common_utils import extract_text_from_a2a_response
from ..common_utils import A2AError, extract_text_from_a2a_response
class A2AModelResponseIterator(BaseModelResponseIterator):
@ -30,6 +34,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator):
json_mode=json_mode,
)
self.model = model
self.response_id: str | None = None
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream:
"""
@ -56,26 +61,161 @@ class A2AModelResponseIterator(BaseModelResponseIterator):
}
}
"""
if "error" in chunk:
error_value: Final = chunk["error"]
error_message: Final = (
error_value.get("message")
if isinstance(error_value, dict) and isinstance(error_value.get("message"), str)
else str(error_value)
)
raise A2AError(status_code=500, message=f"A2A error: {error_message}")
try:
if self.response_id is None:
raw_response_id = chunk.get("id")
raw_result = chunk.get("result")
if not isinstance(raw_response_id, str) and isinstance(raw_result, Mapping):
raw_response_id = raw_result.get("id")
self.response_id = (
raw_response_id
if isinstance(raw_response_id, str) and raw_response_id.strip()
else f"chatcmpl-{uuid4().hex}"
)
# Extract text from A2A response
text: Final = extract_text_from_a2a_response(chunk)
result: Final = chunk.get("result", {})
chunk_index = 0
if isinstance(result, Mapping):
artifact = result.get("artifact")
if isinstance(artifact, Mapping) and isinstance(artifact.get("index"), int):
chunk_index = artifact["index"]
choices = result.get("choices")
if isinstance(choices, list) and choices and isinstance(choices[0], Mapping):
raw_index = choices[0].get("index")
if isinstance(raw_index, int):
chunk_index = raw_index
status: Final = result.get("status", {}) if isinstance(result, Mapping) else {}
is_working_status: Final = (
isinstance(result, Mapping)
and result.get("kind") == "status-update"
and isinstance(status, Mapping)
and status.get("state") == "working"
)
text: Final = "" if is_working_status else extract_text_from_a2a_response(chunk)
provider_fields: dict[str, object] = {}
provider_fields.update(
{
key: value
for key, value in chunk.items()
if key in {"system_fingerprint", "service_tier"} and value is not None
}
)
if isinstance(result, Mapping) and not is_working_status:
control_fields = {
"artifacts",
"choices",
"contextId",
"final",
"finish_reason",
"history",
"id",
"kind",
"message",
"parts",
"status",
"taskId",
"tool_calls",
"usage",
}
provider_fields.update(
{key: value for key, value in result.items() if key not in control_fields and value is not None}
)
choices = result.get("choices")
if isinstance(choices, list) and choices:
first_choice = choices[0]
if isinstance(first_choice, Mapping):
provider_fields.update(
{
key: value
for key, value in first_choice.items()
if key not in {"index", "message", "finish_reason"} and value is not None
}
)
first_message = first_choice.get("message")
if isinstance(first_message, Mapping):
provider_fields.update(
{
key: value
for key, value in first_message.items()
if key not in {"kind", "role", "parts", "tool_calls"} and value is not None
}
)
# Determine finish reason
finish_reason: Final = self._get_finish_reason(chunk)
tool_calls: Final = self._get_tool_calls(chunk)
usage: Final = self._get_usage(chunk)
if isinstance(result, Mapping):
choices = result.get("choices")
if isinstance(choices, list) and choices:
streaming_choices: list[StreamingChoices] = []
for choice_position, raw_choice in enumerate(choices):
if not isinstance(raw_choice, Mapping):
continue
raw_index = raw_choice.get("index", choice_position)
choice_index = raw_index if isinstance(raw_index, int) else choice_position
delta_fields: dict[str, Any] = {}
raw_delta = raw_choice.get("delta")
if isinstance(raw_delta, Mapping):
delta_fields.update(raw_delta)
raw_message = raw_choice.get("message")
if isinstance(raw_message, Mapping):
message_text = extract_text_from_a2a_response({"result": {"message": raw_message}})
if message_text and "content" not in delta_fields:
delta_fields["content"] = message_text
message_tool_calls = raw_message.get("tool_calls")
if message_tool_calls and "tool_calls" not in delta_fields:
delta_fields["tool_calls"] = message_tool_calls
raw_finish_reason = raw_choice.get("finish_reason")
choice_finish_reason = (
raw_finish_reason
if isinstance(raw_finish_reason, str) and raw_finish_reason
else finish_reason
)
streaming_choices.append(
StreamingChoices(
index=choice_index,
delta=Delta(**delta_fields),
finish_reason=choice_finish_reason,
logprobs=raw_choice.get("logprobs"),
)
)
if streaming_choices:
return ModelResponseStream(
choices=streaming_choices,
id=self.response_id,
usage=usage,
provider_specific_fields=provider_fields or None,
)
# Return generic streaming chunk
return GenericStreamingChunk(
text=text,
is_finished=bool(finish_reason),
finish_reason=finish_reason or "",
usage=None,
index=0,
tool_use=None,
id=self.response_id,
is_finished=bool(finish_reason or tool_calls),
finish_reason=finish_reason or ("tool_calls" if tool_calls else ""),
usage=usage,
index=chunk_index,
tool_use=tool_calls,
provider_specific_fields=provider_fields or None,
)
except Exception:
# Return empty chunk on parse error
if self.response_id is None:
self.response_id = f"chatcmpl-{uuid4().hex}"
return GenericStreamingChunk(
text="",
id=self.response_id,
is_finished=False,
finish_reason="",
usage=None,
@ -83,12 +223,25 @@ class A2AModelResponseIterator(BaseModelResponseIterator):
tool_use=None,
)
def _handle_string_chunk(self, str_line: str | dict) -> GenericStreamingChunk | ModelResponseStream:
if isinstance(str_line, dict):
return self.chunk_parser(chunk=str_line)
return super()._handle_string_chunk(str_line=str_line)
def _get_finish_reason(self, chunk: dict) -> str | None:
"""Extract finish reason from A2A chunk"""
result: Final = chunk.get("result", {})
# Check for task completion
if isinstance(result, dict):
explicit_finish_reason: Final = result.get("finish_reason")
if isinstance(explicit_finish_reason, str) and explicit_finish_reason:
return explicit_finish_reason
message: Final = result.get("message")
if isinstance(message, dict):
message_finish_reason: Final = message.get("finish_reason")
if isinstance(message_finish_reason, str) and message_finish_reason:
return message_finish_reason
status: Final = result.get("status", {})
if isinstance(status, dict):
state: Final = status.get("state")
@ -102,3 +255,67 @@ class A2AModelResponseIterator(BaseModelResponseIterator):
return "stop"
return None
def _get_usage(self, chunk: dict) -> object | None:
raw_usage: object | None = chunk.get("usage")
result: Final = chunk.get("result", {})
if raw_usage is None and isinstance(result, dict):
raw_usage = result.get("usage")
if raw_usage is None:
return None
if isinstance(raw_usage, Mapping):
try:
return litellm.Usage(**raw_usage)
except Exception:
return raw_usage
if hasattr(raw_usage, "model_dump"):
try:
return litellm.Usage(**raw_usage.model_dump(exclude_none=True))
except Exception:
return raw_usage
return raw_usage
def _get_tool_calls(self, chunk: dict) -> ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None:
result: Final = chunk.get("result", {})
if not isinstance(result, dict):
return None
tool_calls = result.get("tool_calls")
if isinstance(tool_calls, list) and tool_calls:
return self._serialize_tool_calls(tool_calls)
message = result.get("message")
if isinstance(message, dict) and isinstance(message.get("tool_calls"), list) and message["tool_calls"]:
return self._serialize_tool_calls(message["tool_calls"])
return None
@staticmethod
def _serialize_tool_call(tool_call: object) -> ChatCompletionToolCallChunk | None:
if isinstance(tool_call, dict):
return tool_call
if hasattr(tool_call, "model_dump"):
return tool_call.model_dump(exclude_none=True)
if hasattr(tool_call, "dict"):
return tool_call.dict(exclude_none=True)
return None
@classmethod
def _serialize_tool_calls(
cls, tool_calls: list[object]
) -> ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None:
serialized: Final = [
tool_call_value
for tool_call in tool_calls
if (tool_call_value := cls._serialize_tool_call(tool_call)) is not None
]
if len(serialized) == 1:
return serialized[0]
return serialized or None
async def aclose(self) -> None:
streaming_response = self.streaming_response
self.streaming_response = None
try:
await super().aclose()
finally:
close_stream = getattr(streaming_response, "aclose", None)
if close_stream is not None:
await close_stream()

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

@ -850,17 +850,30 @@ 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
except HTTPException as e:
try:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=data,
)
except Exception:
pass
raise
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()
if post_call_succeeded:
_enqueue_fn(response)
response_dict: Final[dict[str, Any]] = (
response.model_dump(mode="json", exclude_none=True)

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

@ -5,20 +5,507 @@ Handles routing for A2A agents (models with "a2a/<agent-name>" prefix).
Looks up agents in the registry and injects their API base URL.
"""
from typing import Any, Final
from __future__ import annotations
import asyncio
from collections.abc import Awaitable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from uuid import uuid4
from fastapi import HTTPException
from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.interactions.agents.utils import merge_agent_headers
from litellm.llms.a2a.common_utils import A2AError, convert_messages_to_prompt, extract_text_from_a2a_response
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, CustomPricingLiteLLMParams, Message, ModelResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str])
_MESSAGES_ADAPTER: Final = TypeAdapter(list[AllMessageValues])
_FORWARDED_REQUEST_PARAMS: Final = frozenset(
{
"audio",
"frequency_penalty",
"functions",
"function_call",
"guided_json",
"include_server_side_tool_invocations",
"logit_bias",
"logprobs",
"guardrails",
"max_completion_tokens",
"max_tokens",
"modalities",
"n",
"parallel_tool_calls",
"prediction",
"presence_penalty",
"reasoning_effort",
"response_format",
"seed",
"service_tier",
"safety_identifier",
"stop",
"store",
"stream_options",
"temperature",
"thinking",
"timeout",
"tool_choice",
"tools",
"top_logprobs",
"top_p",
"user",
"verbosity",
"web_search_options",
"output_config",
"prompt_cache_key",
}
)
_A2A_PRICING_PARAMS: Final = frozenset({"cost_per_query", "response_cost"}) | frozenset(
CustomPricingLiteLLMParams.model_fields
)
def _get_agent_request_headers(data: Mapping[str, object]) -> dict[str, str]:
proxy_request: Final = data.get("proxy_server_request")
raw_headers: object = proxy_request.get("headers") if isinstance(proxy_request, Mapping) else None
if not isinstance(raw_headers, Mapping):
metadata: Final = data.get("metadata")
if not isinstance(metadata, Mapping):
metadata = data.get("litellm_metadata")
raw_headers = metadata.get("headers") if isinstance(metadata, Mapping) else None
return (
{str(key).lower(): str(value) for key, value in raw_headers.items()} if isinstance(raw_headers, Mapping) else {}
)
class _A2ATextPart(TypedDict):
kind: ReadOnly[str]
text: ReadOnly[str]
class _A2AMessage(TypedDict):
role: ReadOnly[str]
parts: ReadOnly[tuple[_A2ATextPart, ...]]
messageId: ReadOnly[str]
contextId: ReadOnly[str | None]
class _A2AParams(TypedDict):
message: ReadOnly[_A2AMessage]
messages: ReadOnly[list[AllMessageValues]]
async def _route_registered_provider(
data: Mapping[str, object],
model_name: str,
api_base: str | None,
litellm_params: Mapping[str, object],
static_headers: Mapping[str, str] | None,
dynamic_headers: Mapping[str, str] | None = None,
) -> ModelResponse | CustomStreamWrapper:
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator
raw_messages: Final = data.get("messages")
messages: Final = _MESSAGES_ADAPTER.validate_python(raw_messages)
stream: Final = data.get("stream") is True
request_id: Final = str(uuid4())
raw_session_id: Final = data.get("litellm_session_id")
metadata: Final = data.get("metadata")
session_id: Final = (
raw_session_id
if isinstance(raw_session_id, str)
else metadata.get("session_id")
if isinstance(metadata, Mapping) and isinstance(metadata.get("session_id"), str)
else None
)
params: Final[_A2AParams] = {
"message": {
"role": "user",
"parts": ({"kind": "text", "text": convert_messages_to_prompt(messages)},),
"messageId": str(uuid4()),
"contextId": session_id,
},
"messages": messages,
}
provider_params: Final = {
**_OBJECT_DICT_ADAPTER.validate_python(litellm_params),
**{key: data[key] for key in _FORWARDED_REQUEST_PARAMS if key in data and data[key] is not None},
}
registered_provider: Final = litellm_params.get("custom_llm_provider")
registered_model: Final = litellm_params.get("model")
native_provider: Final = registered_provider == "pydantic_ai_agents" or (
registered_provider == "bedrock" and isinstance(registered_model, str) and "agentcore" in registered_model
)
bridge_params: Final = _OBJECT_DICT_ADAPTER.validate_python(
{"message": params["message"]} if native_provider else params
)
configured_headers: Final = litellm_params.get("extra_headers") or litellm_params.get("headers")
configured_headers_dict: Final = (
_HEADERS_ADAPTER.validate_python(configured_headers) if isinstance(configured_headers, dict) else None
)
agent_extra_headers: Final = merge_agent_headers(
dynamic_headers=merge_agent_headers(
dynamic_headers=dynamic_headers,
static_headers=configured_headers_dict,
),
static_headers=static_headers,
)
if agent_extra_headers:
provider_params["extra_headers"] = agent_extra_headers
logging_obj: Final = data.get("litellm_logging_obj")
if isinstance(logging_obj, Logging):
provider_params["no-log"] = True
provider_model: Final = litellm_params.get("model")
if isinstance(provider_model, str):
logging_obj.model_call_details["model"] = provider_model
logging_obj.model_call_details.setdefault("litellm_params", {})["model"] = provider_model
provider_name: Final = litellm_params.get("custom_llm_provider")
if isinstance(provider_name, str):
logging_obj.model_call_details["custom_llm_provider"] = provider_name
pricing_params = {
key: litellm_params[key]
for key in _A2A_PRICING_PARAMS
if key in litellm_params and litellm_params[key] is not None
}
if pricing_params:
logging_obj.litellm_params.update(pricing_params)
logging_obj.model_call_details["litellm_params"].update(pricing_params)
logging_obj.custom_pricing = True
if stream:
streaming_response: Final = A2ACompletionBridgeHandler.handle_streaming(
request_id=request_id,
params=bridge_params,
litellm_params=provider_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
agent_static_headers=static_headers,
)
completion_stream: Final = A2AModelResponseIterator(
streaming_response=streaming_response,
sync_stream=False,
model=model_name,
)
if not isinstance(logging_obj, Logging):
raise TypeError("litellm_logging_obj is required for streaming A2A requests")
return CustomStreamWrapper(
completion_stream=completion_stream,
model=model_name,
custom_llm_provider="a2a",
logging_obj=logging_obj,
stream_options=data.get("stream_options"),
)
response: Final = await A2ACompletionBridgeHandler.handle_non_streaming(
request_id=request_id,
params=bridge_params,
litellm_params=provider_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
agent_static_headers=static_headers,
)
error_value: Final = response.get("error")
if isinstance(error_value, dict):
error: Final = _OBJECT_DICT_ADAPTER.validate_python(error_value)
error_message: Final = error.get("message")
raise A2AError(
status_code=500,
message=f"A2A error: {error_message if isinstance(error_message, str) else 'Unknown error'}",
)
result: Final = response.get("result")
result_dict: Final = result if isinstance(result, Mapping) else {}
nested_message: Final = result_dict.get("message")
response_message: Final = nested_message if isinstance(nested_message, Mapping) else result_dict
response_choices: Final = response.get("choices")
choice_payloads: Final = response_choices if isinstance(response_choices, list) else result_dict.get("choices")
def _serialize_value(value: object) -> object:
if hasattr(value, "model_dump"):
return value.model_dump(exclude_none=True)
if hasattr(value, "dict"):
return value.dict(exclude_none=True)
return value
def _build_message(message_payload: Mapping[str, object], content: str) -> Message:
message_kwargs: dict[str, object] = {
"content": content,
"role": "assistant",
}
raw_tool_calls = message_payload.get("tool_calls")
if isinstance(raw_tool_calls, list):
message_kwargs["tool_calls"] = raw_tool_calls
for field in (
"audio",
"annotations",
"function_call",
"images",
"provider_specific_fields",
"reasoning_content",
"reasoning_items",
"refusal",
"thinking_blocks",
):
value = message_payload.get(field)
if value is not None:
message_kwargs[field] = _serialize_value(value)
return Message(**message_kwargs)
if isinstance(choice_payloads, list) and choice_payloads:
model_choices = []
for choice_index, choice in enumerate(choice_payloads):
choice_mapping: Mapping[str, object] = choice if isinstance(choice, Mapping) else {}
raw_message = choice_mapping.get("message")
message_payload: Mapping[str, object] = raw_message if isinstance(raw_message, Mapping) else choice_mapping
choice_kwargs: dict[str, object] = {
"finish_reason": (
choice_mapping.get("finish_reason")
if isinstance(choice_mapping.get("finish_reason"), str)
else message_payload.get("finish_reason")
if isinstance(message_payload.get("finish_reason"), str)
else "stop"
),
"index": choice_mapping.get("index", choice_index)
if isinstance(choice_mapping.get("index", choice_index), int)
else choice_index,
"message": _build_message(
message_payload,
extract_text_from_a2a_response({"result": message_payload}),
),
}
raw_logprobs = choice_mapping.get("logprobs", message_payload.get("logprobs"))
if raw_logprobs is not None:
choice_kwargs["logprobs"] = _serialize_value(raw_logprobs)
model_choices.append(Choices(**choice_kwargs))
else:
tool_calls: Final = response_message.get("tool_calls")
normalized_tool_calls: Final = tool_calls if isinstance(tool_calls, list) else None
finish_reason: Final = response_message.get("finish_reason")
text: Final = extract_text_from_a2a_response(response)
choice_kwargs = {
"finish_reason": (
finish_reason if isinstance(finish_reason, str) else "tool_calls" if normalized_tool_calls else "stop"
),
"index": 0,
"message": _build_message(response_message, text),
}
raw_logprobs = response_message.get("logprobs")
if raw_logprobs is not None:
choice_kwargs["logprobs"] = _serialize_value(raw_logprobs)
model_choices = [Choices(**choice_kwargs)]
model_response: Final = ModelResponse(
id=str(response.get("id") or request_id),
model=model_name,
choices=model_choices,
system_fingerprint=response.get("system_fingerprint")
if isinstance(response.get("system_fingerprint"), str)
else None,
service_tier=response.get("service_tier") if isinstance(response.get("service_tier"), str) else None,
)
raw_usage: Final = response.get("usage")
usage = litellm.Usage(**raw_usage) if isinstance(raw_usage, Mapping) else raw_usage
if usage is None and native_provider:
try:
from litellm.utils import token_counter
prompt_tokens: Final = token_counter(model="gpt-3.5-turbo", messages=messages)
completion_tokens: Final = token_counter(
model="gpt-3.5-turbo",
text=extract_text_from_a2a_response(response),
count_response_tokens=True,
)
usage = litellm.Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
except Exception: # noqa: BLE001 - token estimation must not fail the response
pass
if usage is not None:
model_response.usage = usage
if isinstance(logging_obj, Logging):
logging_obj.model_call_details["usage"] = usage
if isinstance(logging_obj, Logging):
def _enqueue_logging(final_response: ModelResponse | None = None) -> None:
asyncio.create_task(
logging_obj.dispatch_success_handlers(
final_response if final_response is not None else model_response,
cache_hit=False,
prefer_async_handlers=True,
)
)
logging_obj._enqueue_deferred_logging = _enqueue_logging
return model_response
def _merge_agent_guardrails(
data: Mapping[str, object],
agent_guardrails: object,
) -> Mapping[str, object]:
if not agent_guardrails:
return data
configured_guardrails: list[object] = agent_guardrails if isinstance(agent_guardrails, list) else [agent_guardrails]
metadata_key: Final = "litellm_metadata" if "litellm_metadata" in data else "metadata"
metadata = data.get(metadata_key)
metadata_guardrails = metadata.get("guardrails") if isinstance(metadata, dict) else None
root_guardrails = data.get("guardrails")
existing_guardrails: list[object] = []
for value in (metadata_guardrails, root_guardrails):
if isinstance(value, list):
existing_guardrails.extend(value)
elif value:
existing_guardrails.append(value)
merged_guardrails = existing_guardrails + [
guardrail for guardrail in configured_guardrails if guardrail not in existing_guardrails
]
if isinstance(data, dict):
data["guardrails"] = merged_guardrails
if isinstance(metadata, dict):
data[metadata_key] = {**metadata, "guardrails": merged_guardrails}
return data
merged_data = dict(data)
merged_data["guardrails"] = merged_guardrails
if isinstance(metadata, dict):
merged_data[metadata_key] = {**metadata, "guardrails": merged_guardrails}
return merged_data
async def merge_a2a_agent_guardrails_before_hooks(data: Mapping[str, object]) -> Mapping[str, object]:
model_name: Final = data.get("model")
if not isinstance(model_name, str) or not model_name.startswith("a2a/"):
return data
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
agent = await get_agent_with_read_through(model_name[4:])
if agent is None or not agent.litellm_params:
return data
return _merge_agent_guardrails(data, agent.litellm_params.get("guardrails"))
async def authorize_a2a_agent_before_hooks(
data: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth | None,
) -> Mapping[str, object]:
model_name: Final = data.get("model")
if not isinstance(model_name, str) or not model_name.startswith("a2a/"):
return data
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentRequestHandler
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
agent = await get_agent_with_read_through(model_name[4:])
if agent is None:
return data
is_admin: Final = user_api_key_dict is not None and (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
if not is_admin:
is_allowed: Final = await AgentRequestHandler.is_agent_allowed(
agent_id=agent.agent_id,
user_api_key_auth=user_api_key_dict,
)
if not is_allowed:
raise HTTPException(
status_code=403,
detail=f"Agent '{agent.agent_name}' is not allowed for your key/team. Contact proxy admin for access.",
)
if (agent.litellm_params or {}).get("require_trace_id_on_calls_to_agent"):
_enforce_inbound_trace_id(data, agent.agent_id)
if isinstance(data, dict):
data["agent_id"] = agent.agent_id
metadata = data.get("metadata")
if not isinstance(metadata, dict):
metadata = {}
data["metadata"] = metadata
metadata["agent_id"] = agent.agent_id
return data
def _get_agent_dynamic_headers(
data: Mapping[str, object],
agent_id: str,
agent_name: str,
extra_headers: list[str] | None,
) -> dict[str, str]:
normalized_headers: Final = _get_agent_request_headers(data)
dynamic_headers: dict[str, str] = {}
for header_name in extra_headers or []:
header_name_str: Final = str(header_name)
if header_name_str.lower().startswith("x-litellm-"):
continue
value: Final = normalized_headers.get(header_name_str.lower())
if value is not None:
dynamic_headers[header_name_str] = value
for alias in (agent_id.lower(), agent_name.lower()):
prefix: Final = f"x-a2a-{alias}-"
for key, value in normalized_headers.items():
if key.startswith(prefix):
header_name: Final = key[len(prefix) :]
if header_name and not header_name.lower().startswith("x-litellm-"):
dynamic_headers[header_name] = value
return dynamic_headers
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 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 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
def _enforce_inbound_trace_id(data: Mapping[str, object], agent_id: str) -> None:
from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
if not get_chain_id_from_headers(_get_agent_request_headers(data)):
raise HTTPException(
status_code=400,
detail=f"Agent '{agent_id}' requires x-litellm-trace-id header on all inbound requests.",
)
async def route_a2a_agent_request(
data: dict,
data: Mapping[str, object],
route_type: str,
user_api_key_dict: UserAPIKeyAuth | None = None,
) -> Any | None:
) -> Awaitable[object] | None:
"""
Route A2A agent requests directly to litellm with injected API base.
@ -68,14 +555,75 @@ async def route_a2a_agent_request(
detail=f"Agent '{agent_name}' is not allowed for your key/team. Contact proxy admin for access.",
)
if (agent.litellm_params or {}).get("require_trace_id_on_calls_to_agent"):
_enforce_inbound_trace_id(data, agent.agent_id)
# Get API base URL from agent config
if not agent.agent_card_params or "url" not in agent.agent_card_params:
agent_card_params: Final = agent.agent_card_params
agent_url: Final = agent_card_params.get("url") if agent_card_params else None
registered_params_value: Final = agent.litellm_params
registered_provider_value: Final = (
registered_params_value.get("custom_llm_provider") if registered_params_value else None
)
registered_provider: Final = registered_provider_value if isinstance(registered_provider_value, str) else None
from litellm.a2a_protocol.litellm_completion_bridge.handler import A2A_USER_API_KEY_HASH_PARAM
registered_params_for_route: Final[Mapping[str, object]] = (
{
**registered_params_value,
A2A_USER_API_KEY_HASH_PARAM: user_api_key_dict.api_key,
}
if registered_params_value and user_api_key_dict is not None and user_api_key_dict.api_key
else registered_params_value or {}
)
configured_api_base: Final = registered_params_value.get("api_base") if registered_params_value else None
api_base: Final = configured_api_base if isinstance(configured_api_base, str) and configured_api_base else agent_url
cardless_provider: Final = registered_provider is not None and registered_provider != "a2a"
has_configured_api_base: Final = isinstance(configured_api_base, str) and bool(configured_api_base)
if (not isinstance(agent_url, str) or not agent_url) and not has_configured_api_base and not cardless_provider:
verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name)
route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False)
# Inject API base and route to litellm
data["api_base"] = agent.agent_card_params["url"]
verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, data["api_base"])
routed_data: Final = _merge_agent_guardrails(
data=data,
agent_guardrails=registered_params_value.get("guardrails") if registered_params_value else None,
)
registered_dynamic_headers: Final = _get_agent_dynamic_headers(
data=routed_data,
agent_id=agent.agent_id,
agent_name=agent.agent_name,
extra_headers=agent.extra_headers,
)
registered_static_headers: Mapping[str, str] | None = agent.static_headers
if registered_params_value and registered_params_value.get("databricks_oauth"):
from litellm.proxy.agent_endpoints.databricks_oauth import resolve_databricks_app_auth_header
return getattr(litellm, f"{route_type}")(**data)
databricks_headers = await resolve_databricks_app_auth_header(dict(registered_params_value))
registered_static_headers = merge_agent_headers(
dynamic_headers=registered_static_headers,
static_headers=databricks_headers,
)
registered_static_headers = merge_agent_headers(
dynamic_headers=registered_static_headers,
static_headers=_get_agent_identity_headers(user_api_key_dict, data.get("litellm_trace_id")),
)
if (
registered_provider
and registered_provider != "a2a"
and route_type == "acompletion"
and registered_params_value is not None
):
verbose_proxy_logger.debug("[A2A] Routing %s through %s", model_name, registered_provider)
return _route_registered_provider(
data=routed_data,
model_name=model_name,
api_base=api_base,
litellm_params=registered_params_for_route,
static_headers=registered_static_headers,
dynamic_headers=registered_dynamic_headers,
)
completion_data: Final = MappingProxyType({**routed_data, "api_base": api_base})
verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, api_base)
return getattr(litellm, f"{route_type}")(**completion_data) # pyright: ignore[reportAny] # dynamic SDK route

12
litellm/proxy/agent_endpoints/model_list_helpers.py
View file

@ -4,6 +4,8 @@ Helper functions for appending A2A agents to model lists.
Used by proxy model endpoints to make agents appear in UI alongside models.
"""
from typing import Final
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
@ -33,11 +35,14 @@ async def append_agents_to_model_group(
for agent_id in allowed_agent_ids:
agent = global_agent_registry.get_agent_by_id(agent_id)
if agent is not None:
agent_params: Final = agent.litellm_params
provider_value: Final = agent_params.get("custom_llm_provider") if agent_params else None
custom_llm_provider: Final = provider_value if isinstance(provider_value, str) else "a2a"
model_groups.append(
ModelGroupInfoProxy(
model_group=f"a2a/{agent.agent_name}",
mode="chat",
providers=["a2a"],
providers=[custom_llm_provider],
)
)
case _:
@ -70,12 +75,15 @@ async def append_agents_to_model_info(
for agent_id in allowed_agent_ids:
agent = global_agent_registry.get_agent_by_id(agent_id)
if agent is not None:
agent_params: Final = agent.litellm_params
provider_value: Final = agent_params.get("custom_llm_provider") if agent_params else None
custom_llm_provider: Final = provider_value if isinstance(provider_value, str) else "a2a"
models.append(
{
"model_name": f"a2a/{agent.agent_name}",
"litellm_params": {
"model": f"a2a/{agent.agent_name}",
"custom_llm_provider": "a2a",
"custom_llm_provider": custom_llm_provider,
},
"model_info": {
"id": agent.agent_id,

36
litellm/proxy/common_request_processing.py
View file

@ -1876,6 +1876,11 @@ class ProxyBaseLLMRequestProcessing:
## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call
## IMPORTANT Note: - initialize this before running pre-call checks. Ensures we log rejected requests to langfuse.
from litellm.proxy.agent_endpoints.a2a_routing import (
authorize_a2a_agent_before_hooks,
merge_a2a_agent_guardrails_before_hooks,
)
logging_obj, self.data = litellm.utils.function_setup(
original_function=route_type,
rules_obj=litellm.utils.Rules(),
@ -1885,6 +1890,13 @@ class ProxyBaseLLMRequestProcessing:
self.data["litellm_logging_obj"] = logging_obj
self.data = await authorize_a2a_agent_before_hooks(
data=self.data,
user_api_key_dict=user_api_key_dict,
)
self.data = await merge_a2a_agent_guardrails_before_hooks(self.data)
# Merge model-level guardrails before pre_call_hook so DB/UI-configured
# guardrails actually execute on pre_call. Without this, guardrails set
# via litellm_params.guardrails are only honored on post_call paths
@ -1900,12 +1912,32 @@ class ProxyBaseLLMRequestProcessing:
trust_client_model_info=False,
)
authorized_model = self.data.get("model")
self.data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=self.data,
call_type=route_type,
)
if self.data.get("model") != authorized_model:
self.data = await authorize_a2a_agent_before_hooks(
data=self.data,
user_api_key_dict=user_api_key_dict,
)
self.data = await merge_a2a_agent_guardrails_before_hooks(self.data)
self.data = _check_and_merge_model_level_guardrails(
data=self.data,
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
# add_litellm_data_to_request predates that mutation.
@ -2539,6 +2571,7 @@ class ProxyBaseLLMRequestProcessing:
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(
logging_obj=logging_obj,
exception_raised=_exception_raised,
response=response,
)
# Streaming cleanup: if an exception occurred AND the deferred
@ -3023,6 +3056,7 @@ class ProxyBaseLLMRequestProcessing:
def _flush_deferred_async_logging(
logging_obj: Any,
exception_raised: bool,
response: Any | None = None,
) -> None:
"""
Fire the deferred async-success closure stored by wrapper_async, then
@ -3053,7 +3087,7 @@ class ProxyBaseLLMRequestProcessing:
if exception_raised:
return
try:
_enqueue_fn()
_enqueue_fn(response) if response is not None else _enqueue_fn()
except Exception as e:
verbose_proxy_logger.exception("Error firing deferred logging: %s", e)

3
litellm/types/utils.py
View file

@ -323,7 +323,8 @@ class ModelInfo(ModelInfoBase, total=False):
class GenericStreamingChunk(TypedDict, total=False):
text: Required[str]
tool_use: ChatCompletionToolCallChunk | None
id: str
tool_use: ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None
is_finished: Required[bool]
finish_reason: Required[str]
usage: Required[ChatCompletionUsageBlock | None]

4
litellm/utils.py
View file

@ -1928,11 +1928,11 @@ def client(original_function):
if not _is_litellm_internal_call:
if getattr(logging_obj, "_defer_async_logging", False):
def _enqueue_deferred_logging() -> None:
def _enqueue_deferred_logging(final_response=None) -> None:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
result=final_response if final_response is not None else result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,

4
tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
View file

@ -10,10 +10,9 @@ Verifies that:
"""
import json
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
SAMPLE_ARN = "arn:aws:bedrock-agentcore:us-west-2:123456789:runtime/my_agent"
SAMPLE_MODEL = f"bedrock/agentcore/{SAMPLE_ARN}"
@ -482,6 +481,7 @@ class TestHandlerIntegration:
api_base=None,
litellm_params=SAMPLE_LITELLM_PARAMS,
agent_extra_headers=None,
agent_static_headers=None,
)
@pytest.mark.asyncio

30
tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py
View file

@ -4,10 +4,10 @@ Tests for Pydantic AI agents transformation.
Tests the helper functions and response transformation without making real API calls.
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
PydanticAITransformation,
)
@ -94,3 +94,31 @@ class TestPydanticAITransformation:
assert result["result"]["kind"] == "message"
assert result["result"]["role"] == "agent"
assert result["result"]["parts"][0]["text"] == "The answer is 4."
@pytest.mark.asyncio
async def test_poll_for_completion_honors_request_timeout():
client = MagicMock()
response = MagicMock()
response.json.return_value = {
"result": {"status": {"state": "working"}, "id": "task-1"}
}
response.raise_for_status.return_value = None
client.post = AsyncMock(return_value=response)
with patch(
"litellm.a2a_protocol.providers.pydantic_ai_agents.transformation.time.monotonic",
side_effect=[100.0, 100.0, 100.02],
):
with pytest.raises(TimeoutError, match=r"0\.01 seconds"):
await PydanticAITransformation._poll_for_completion(
client=client,
endpoint="http://example.test",
task_id="task-1",
request_id="req-1",
poll_interval=1.0,
timeout=0.01,
)
client.post.assert_awaited_once()
assert client.post.await_args.kwargs["timeout"] == pytest.approx(0.01)

60
tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py
View file

@ -6,9 +6,11 @@ from pathlib import Path
import httpx
import pytest
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
from litellm.a2a_protocol.providers.watsonx_orchestrate import handler as wxo_handler
from litellm.a2a_protocol.providers.watsonx_orchestrate.config import (
WatsonxOrchestrateA2AConfig,
)
from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import (
WatsonxOrchestrateHandler,
)
@ -46,6 +48,12 @@ class _SSELines:
yield line
class _HangingSSELines:
async def aiter_lines(self):
await asyncio.Event().wait()
yield ""
class _InvalidJsonStreamResponse:
headers = {"content-type": "application/json"}
@ -231,6 +239,12 @@ async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events():
)
@pytest.mark.asyncio
async def test_accumulate_wxo_sse_text_respects_timeout():
with pytest.raises(asyncio.TimeoutError):
await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(_HangingSSELines(), timeout=0.001)
@pytest.mark.asyncio
async def test_short_lived_tokens_are_not_served_from_cache():
client = _ShortTtlTokenClient()
@ -343,6 +357,50 @@ async def test_poll_run_raises_asyncio_timeout_when_never_terminal():
assert client.get_calls == 2
def test_build_wxo_headers_preserves_auth_headers():
headers = wxo_handler._build_wxo_headers(
token="token",
accept="application/json",
static_headers={
"x-tenant-id": "tenant-1",
"Authorization": "caller-token",
"content-type": "text/plain",
},
)
assert headers["x-tenant-id"] == "tenant-1"
assert headers["Authorization"] == "Bearer token"
assert headers["Content-Type"] == "application/json"
assert headers["Accept"] == "application/json"
assert "content-type" not in headers
@pytest.mark.asyncio
async def test_wxo_config_forwards_headers_and_timeout(monkeypatch):
captured = {}
async def fake_handle_non_streaming(**kwargs):
captured.update(kwargs)
return {"result": {}}
monkeypatch.setattr(WatsonxOrchestrateHandler, "handle_non_streaming", fake_handle_non_streaming)
await WatsonxOrchestrateA2AConfig().handle_non_streaming(
request_id="req-1",
params={},
litellm_params={"model": "agent"},
agent_extra_headers={"x-request-id": "request-1"},
agent_static_headers={"x-tenant-id": "tenant-1"},
timeout=12,
)
assert captured["static_headers"] == {
"x-request-id": "request-1",
"x-tenant-id": "tenant-1",
}
assert captured["timeout"] == 12
@pytest.mark.asyncio
async def test_handle_streaming_polls_non_sse_json_until_complete(monkeypatch):
client = _JsonStreamClient({"status": "running", "run_id": "run-1"})

314
tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py
View file

@ -12,6 +12,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.types.utils import Choices, Message, ModelResponse
class TestA2AStreamingTransformation:
"""Test the A2A streaming transformation creates proper events."""
@ -26,9 +28,7 @@ class TestA2AStreamingTransformation:
"parts": [{"text": "Reply to ticket #4823"}],
"metadata": {"skillId": "draft_reply"},
}
openai_messages = (
A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
# Metadata is forwarded on the run payload only, not duplicated on messages.
assert "metadata" not in openai_messages[0]
@ -174,10 +174,7 @@ class TestA2AStreamingTransformation:
assert "artifactId" in event["result"]["artifact"]
assert event["result"]["artifact"]["name"] == "response"
assert event["result"]["artifact"]["parts"][0]["kind"] == "text"
assert (
event["result"]["artifact"]["parts"][0]["text"]
== "Hello, I am an AI assistant."
)
assert event["result"]["artifact"]["parts"][0]["text"] == "Hello, I am an AI assistant."
@pytest.mark.asyncio
@ -197,6 +194,8 @@ async def test_handle_streaming_emits_proper_events():
mock_chunk2.choices = [MagicMock()]
mock_chunk2.choices[0].delta = MagicMock()
mock_chunk2.choices[0].delta.content = " world"
mock_chunk2.choices[0].finish_reason = "length"
mock_chunk2.usage = {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}
async def mock_streaming_response():
yield mock_chunk1
@ -222,8 +221,8 @@ async def test_handle_streaming_emits_proper_events():
):
events.append(event)
# Should have 4 events: task, working, artifact, completed
assert len(events) == 4
# Should have 5 events: task, working, two artifacts, completed
assert len(events) == 5
# Event 1: task submitted
assert events[0]["result"]["kind"] == "task"
@ -234,14 +233,299 @@ async def test_handle_streaming_emits_proper_events():
assert events[1]["result"]["status"]["state"] == "working"
assert events[1]["result"]["final"] is False
# Event 3: artifact update with accumulated content
# Event 3: first artifact update
assert events[2]["result"]["kind"] == "artifact-update"
assert events[2]["result"]["artifact"]["parts"][0]["text"] == "Hello world"
assert events[2]["result"]["artifact"]["parts"][0]["text"] == "Hello"
# Event 4: status completed
assert events[3]["result"]["kind"] == "status-update"
assert events[3]["result"]["status"]["state"] == "completed"
assert events[3]["result"]["final"] is True
# Event 4: second artifact update
assert events[3]["result"]["kind"] == "artifact-update"
assert events[3]["result"]["artifact"]["parts"][0]["text"] == " world"
assert (
events[2]["result"]["artifact"]["artifactId"]
== events[3]["result"]["artifact"]["artifactId"]
)
# Event 5: status completed
assert events[4]["result"]["kind"] == "status-update"
assert events[4]["result"]["status"]["state"] == "completed"
assert events[4]["result"]["final"] is True
assert events[4]["result"]["finish_reason"] == "length"
assert events[4]["usage"]["total_tokens"] == 5
def test_build_completion_params_keeps_bridge_routing_fields():
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
params = A2ACompletionBridgeHandler._build_completion_params(
params={"message": {"role": "user", "parts": []}},
litellm_params={
"custom_llm_provider": "openai",
"model": "agent",
"api_base": "https://untrusted.example",
"stream": False,
},
api_base="https://configured.example",
agent_extra_headers=None,
stream=True,
)
assert params["api_base"] == "https://configured.example"
assert params["stream"] is True
def test_build_completion_params_drops_proxy_only_databricks_oauth():
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
params = A2ACompletionBridgeHandler._build_completion_params(
params={"message": {"role": "user", "parts": []}},
litellm_params={
"custom_llm_provider": "databricks",
"model": "agent",
"databricks_oauth": {"client_id": "id"},
},
api_base="https://configured.example",
agent_extra_headers=None,
stream=False,
)
assert "databricks_oauth" not in params
@pytest.mark.asyncio
async def test_handle_streaming_accumulates_logprobs_and_provider_metadata():
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
chunks = []
for token in ("a", "b"):
choice = MagicMock()
choice.index = 0
choice.finish_reason = None
choice.delta.content = token
choice.logprobs = {"content": [{"token": token}]}
chunk = MagicMock()
chunk.choices = [choice]
chunk.system_fingerprint = "fp-1"
chunk.service_tier = "scale"
chunks.append(chunk)
chunks[-1].choices[0].finish_reason = "stop"
async def mock_streaming_response():
for chunk in chunks:
yield chunk
with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
mock_acompletion.return_value = mock_streaming_response()
events = [
event
async for event in A2ACompletionBridgeHandler.handle_streaming(
request_id="req-metadata",
params={"message": {"role": "user", "parts": []}},
litellm_params={"custom_llm_provider": "openai", "model": "agent"},
)
]
result = events[-1]
assert result["system_fingerprint"] == "fp-1"
assert result["service_tier"] == "scale"
assert result["result"]["choices"][0]["logprobs"]["content"] == [
{"token": "a"},
{"token": "b"},
]
@pytest.mark.asyncio
async def test_handle_streaming_preserves_multiple_choices():
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
mock_chunk = MagicMock()
first_choice = MagicMock()
first_choice.index = 0
first_choice.finish_reason = None
first_choice.delta.content = "first"
second_choice = MagicMock()
second_choice.index = 1
second_choice.finish_reason = "length"
second_choice.delta.content = "second"
mock_chunk.choices = [first_choice, second_choice]
async def mock_streaming_response():
yield mock_chunk
with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
mock_acompletion.return_value = mock_streaming_response()
events = [
event
async for event in A2ACompletionBridgeHandler.handle_streaming(
request_id="req-choices",
params={"message": {"role": "user", "parts": []}},
litellm_params={"custom_llm_provider": "langgraph", "model": "agent", "n": 2},
)
]
choices = events[-1]["result"]["choices"]
assert [choice["index"] for choice in choices] == [0, 1]
assert [choice["message"]["parts"][0]["text"] for choice in choices] == ["", ""]
assert choices[1]["finish_reason"] == "length"
@pytest.mark.asyncio
async def test_handle_streaming_preserves_non_text_delta_fields():
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
delta = MagicMock()
delta.content = ""
delta.tool_calls = None
delta.model_dump.return_value = {
"audio": {"data": "abc"},
"reasoning_content": "thinking",
"provider_specific_fields": {"trace_id": "trace-1"},
}
choice = MagicMock()
choice.index = 0
choice.finish_reason = "stop"
choice.delta = delta
choice.logprobs = {"content": []}
chunk = MagicMock()
chunk.choices = [choice]
async def mock_streaming_response():
yield chunk
with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
mock_acompletion.return_value = mock_streaming_response()
events = [
event
async for event in A2ACompletionBridgeHandler.handle_streaming(
request_id="req-fields",
params={"message": {"role": "user", "parts": []}},
litellm_params={"custom_llm_provider": "langgraph", "model": "agent"},
)
]
result = events[-1]["result"]
choice_result = result["choices"][0]
assert choice_result["delta"] == {
"audio": {"data": "abc"},
"reasoning_content": "thinking",
"provider_specific_fields": {"trace_id": "trace-1"},
}
assert choice_result["logprobs"] == {"content": []}
@pytest.mark.asyncio
async def test_provider_config_receives_full_message_history():
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
provider_config = MagicMock()
provider_config.handle_non_streaming = AsyncMock(return_value={"result": {}})
messages = [
{"role": "system", "content": "Be concise"},
{"role": "user", "content": "Hello"},
]
params = {
"message": {"role": "user", "parts": []},
"messages": messages,
}
with patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config",
return_value=provider_config,
):
await A2ACompletionBridgeHandler.handle_non_streaming(
request_id="req-1",
params=params,
litellm_params={"custom_llm_provider": "langflow", "model": "flow"},
)
assert provider_config.handle_non_streaming.await_args.kwargs["params"]["messages"] == messages
@pytest.mark.asyncio
async def test_native_provider_config_drops_internal_message_history():
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
provider_config = MagicMock()
provider_config.handle_non_streaming = AsyncMock(return_value={"result": {}})
params = {
"message": {"role": "user", "parts": []},
"messages": [{"role": "user", "content": "Hello"}],
}
with patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config",
return_value=provider_config,
):
await A2ACompletionBridgeHandler.handle_non_streaming(
request_id="req-native",
params=params,
litellm_params={"custom_llm_provider": "pydantic_ai_agents", "model": "agent"},
)
assert provider_config.handle_non_streaming.await_args.kwargs["params"] == {
"message": params["message"]
}
def test_response_transform_preserves_audio_and_logprobs():
from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
A2ACompletionBridgeTransformation,
)
response = ModelResponse(
id="resp-1",
model="test-model",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
content="hello",
role="assistant",
audio={"data": "abc", "expires_at": 1, "transcript": "hello"},
),
logprobs={"content": []},
)
],
)
transformed = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(response)
assert transformed["result"]["audio"]["data"] == "abc"
assert transformed["result"]["logprobs"] == {"content": []}
def test_response_transform_preserves_refusal():
from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
A2ACompletionBridgeTransformation,
)
message = MagicMock()
message.model_dump.return_value = {
"content": None,
"refusal": "I cannot help with that request.",
}
choice = MagicMock(message=message)
choice.model_dump.return_value = {"finish_reason": "stop"}
response = MagicMock(choices=[choice], usage=None)
transformed = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(response)
assert transformed["result"]["refusal"] == "I cannot help with that request."
assert transformed["result"]["parts"] == [{"kind": "text", "text": ""}]
@pytest.mark.asyncio

223
tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py Normal file
View file

@ -0,0 +1,223 @@
"""Tests for the A2A chat streaming iterator."""
import pytest
from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator
from litellm.llms.a2a.common_utils import A2AError
from litellm.types.utils import Delta
@pytest.mark.asyncio
async def test_async_iterator_accepts_decoded_a2a_events():
async def _events():
yield {
"jsonrpc": "2.0",
"result": {
"kind": "artifact-update",
"artifact": {"parts": [{"kind": "text", "text": "Hello"}]},
},
}
iterator = A2AModelResponseIterator(
streaming_response=_events(),
sync_stream=False,
)
chunk = await iterator.__aiter__().__anext__()
assert chunk["text"] == "Hello"
def test_chunk_parser_reuses_response_id_for_idless_artifacts():
iterator = A2AModelResponseIterator(streaming_response=[], sync_stream=False)
first = iterator.chunk_parser(
{"result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "one"}]}}}
)
second = iterator.chunk_parser(
{"result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "two"}]}}}
)
assert first["id"] == second["id"]
@pytest.mark.asyncio
async def test_async_iterator_ignores_status_message_text():
async def _events():
yield {
"jsonrpc": "2.0",
"result": {
"kind": "status-update",
"status": {
"state": "working",
"message": {"parts": [{"kind": "text", "text": "Processing request..."}]},
},
},
}
iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False)
chunk = await iterator.__aiter__().__anext__()
assert chunk["text"] == ""
@pytest.mark.asyncio
async def test_async_iterator_preserves_non_text_fields():
async def _events():
yield {
"jsonrpc": "2.0",
"result": {
"kind": "status-update",
"status": {"state": "completed"},
"audio": {"data": "abc"},
"reasoning_content": "thinking",
"logprobs": {"content": []},
},
}
iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False)
chunk = await iterator.__aiter__().__anext__()
assert chunk["provider_specific_fields"] == {
"audio": {"data": "abc"},
"reasoning_content": "thinking",
"logprobs": {"content": []},
}
@pytest.mark.asyncio
async def test_async_iterator_preserves_tool_calls():
tool_calls = [
{
"id": "call-1",
"type": "function",
"function": {"name": "lookup", "arguments": "{}"},
}
]
async def _events():
yield {"jsonrpc": "2.0", "result": {"tool_calls": tool_calls}}
iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False)
chunk = await iterator.__aiter__().__anext__()
assert chunk["tool_use"] == tool_calls[0]
assert chunk["finish_reason"] == "tool_calls"
@pytest.mark.asyncio
async def test_async_iterator_preserves_parallel_tool_calls():
tool_calls = [
{
"id": "call-1",
"type": "function",
"function": {"name": "lookup", "arguments": "{}"},
},
{
"id": "call-2",
"type": "function",
"function": {"name": "write", "arguments": "{}"},
},
]
async def _events():
yield {"jsonrpc": "2.0", "result": {"tool_calls": tool_calls}}
iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False)
chunk = await iterator.__aiter__().__anext__()
assert chunk["tool_use"] == tool_calls
@pytest.mark.asyncio
async def test_async_iterator_preserves_every_terminal_choice():
async def _events():
yield {
"jsonrpc": "2.0",
"result": {
"kind": "status-update",
"status": {"state": "completed"},
"choices": [
{
"index": 0,
"message": {"parts": [{"kind": "text", "text": "first"}]},
"finish_reason": "stop",
},
{
"index": 1,
"message": {"parts": [{"kind": "text", "text": "second"}]},
"finish_reason": "length",
},
],
},
}
iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False)
chunk = await iterator.__aiter__().__anext__()
assert [choice.index for choice in chunk.choices] == [0, 1]
assert [choice.delta.content for choice in chunk.choices] == ["first", "second"]
assert [choice.finish_reason for choice in chunk.choices] == ["stop", "length"]
@pytest.mark.asyncio
async def test_async_iterator_serializes_delta_tool_calls_and_usage():
delta = Delta(
tool_calls=[
{
"id": "call-1",
"type": "function",
"function": {"name": "lookup", "arguments": "{}"},
}
]
)
async def _events():
yield {
"jsonrpc": "2.0",
"usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5},
"result": {
"tool_calls": [delta.tool_calls[0]],
"finish_reason": "length",
},
}
iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False)
chunk = await iterator.__aiter__().__anext__()
assert chunk["tool_use"]["id"] == "call-1"
assert chunk["finish_reason"] == "length"
assert chunk["usage"].total_tokens == 5
@pytest.mark.asyncio
async def test_async_iterator_closes_nested_stream():
closed = False
async def _events():
nonlocal closed
try:
yield {"jsonrpc": "2.0", "result": {"kind": "artifact-update"}}
raise AssertionError("stream should be closed before a second event")
finally:
closed = True
iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False)
await iterator.__aiter__().__anext__()
await iterator.aclose()
assert closed is True
@pytest.mark.asyncio
async def test_async_iterator_propagates_jsonrpc_errors():
async def _events():
yield {"jsonrpc": "2.0", "error": {"code": -32000, "message": "agent failed"}}
iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False)
with pytest.raises(A2AError, match="agent failed"):
await iterator.__aiter__().__anext__()

57
tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py
View file

@ -4,8 +4,6 @@ Test appending A2A agents to model lists.
Maps to: litellm/proxy/agent_endpoints/model_list_helpers.py
"""
from unittest.mock import AsyncMock, Mock, patch
import pytest
@ -66,6 +64,32 @@ async def test_append_agents_to_model_group():
assert result[0].providers == ["a2a"]
@pytest.mark.asyncio
async def test_append_agents_to_model_group_preserves_registered_provider():
agent = AgentResponse(
agent_id="agent-123",
agent_name="test-agent",
agent_card_params={"url": "http://example.com"},
litellm_params={"custom_llm_provider": "pydantic_ai_agents"},
)
registry = Mock()
registry.get_agent_by_id = Mock(return_value=agent)
with (
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access",
AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-123"}))),
),
patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry),
):
result = await append_agents_to_model_group(
model_groups=[],
user_api_key_dict=Mock(spec=UserAPIKeyAuth),
)
assert result[0].providers == ["pydantic_ai_agents"]
@pytest.mark.asyncio
async def test_append_agents_to_model_info():
"""Test agents are converted to model info format with a2a/ prefix"""
@ -109,3 +133,32 @@ async def test_append_agents_to_model_info():
assert result[0]["litellm_params"]["custom_llm_provider"] == "a2a"
assert result[0]["model_info"]["id"] == "agent-123"
assert result[0]["model_info"]["mode"] == "chat"
@pytest.mark.asyncio
async def test_append_agents_to_model_info_preserves_registered_provider():
agent = AgentResponse(
agent_id="agent-123",
agent_name="test-agent",
agent_card_params={"url": "http://example.com"},
litellm_params={"custom_llm_provider": "pydantic_ai_agents"},
)
registry = Mock()
registry.get_agent_by_id = Mock(return_value=agent)
with (
patch( # test-quality-ok: access resolution is outside model-list assembly
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access",
AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-123"}))),
),
patch( # test-quality-ok: registry output drives model-list assembly
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
registry,
),
):
result = await append_agents_to_model_info(
models=[],
user_api_key_dict=Mock(spec=UserAPIKeyAuth),
)
assert result[0]["litellm_params"]["custom_llm_provider"] == "pydantic_ai_agents"

65
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -6505,6 +6505,71 @@ class TestPerRequestModelGroupAlias:
assert merged_for == ["group-b"]
@pytest.mark.asyncio
async def test_a2a_reroute_runs_target_guardrails(self, monkeypatch):
processing_obj = ProxyBaseLLMRequestProcessing(
data={"model": "source", "messages": [{"role": "user", "content": "hello"}]}
)
mock_request = MagicMock(spec=Request)
mock_request.headers = {}
async def mock_add_litellm_data_to_request(*args, **kwargs):
return kwargs.get("data", {})
hook_modes = []
async def reroute_once(user_api_key_dict, data, call_type, guardrails_only=False):
hook_modes.append(guardrails_only)
if len(hook_modes) == 1:
data["model"] = "a2a/agent"
return data
async def passthrough(data, user_api_key_dict):
return data
async def passthrough_data_only(data):
return data
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=reroute_once)
monkeypatch.setattr(
litellm.proxy.common_request_processing,
"add_litellm_data_to_request",
mock_add_litellm_data_to_request,
)
monkeypatch.setattr(
litellm.utils,
"function_setup",
lambda original_function, rules_obj, start_time, **kwargs: (MagicMock(), kwargs),
)
monkeypatch.setattr(
litellm.proxy.common_request_processing,
"_check_and_merge_model_level_guardrails",
lambda data, llm_router, trust_client_model_info=True: data,
)
monkeypatch.setattr(
"litellm.proxy.agent_endpoints.a2a_routing.authorize_a2a_agent_before_hooks",
passthrough,
)
monkeypatch.setattr(
"litellm.proxy.agent_endpoints.a2a_routing.merge_a2a_agent_guardrails_before_hooks",
passthrough_data_only,
)
returned_data, _ = await processing_obj.common_processing_pre_call_logic(
request=mock_request,
general_settings={},
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash"),
proxy_logging_obj=mock_proxy_logging_obj,
proxy_config=None,
route_type="acompletion",
llm_router=None,
)
assert returned_data["model"] == "a2a/agent"
assert hook_modes == [False, True]
class TestInjectCostIntoUsageDict:
@staticmethod
def _expected_cost(model, prompt_tokens, completion_tokens):

640
tests/test_litellm/proxy/test_route_a2a_models.py
View file

@ -4,13 +4,17 @@ Test A2A model routing in proxy.
Maps to: litellm/proxy/agent_endpoints/a2a_routing.py
"""
from unittest.mock import AsyncMock, Mock, patch
import pytest
from fastapi import HTTPException
from litellm.proxy.agent_endpoints.a2a_routing import route_a2a_agent_request
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.agent_endpoints.a2a_routing import (
_route_registered_provider,
merge_a2a_agent_guardrails_before_hooks,
route_a2a_agent_request,
)
from litellm.proxy.route_llm_request import route_request
@ -58,7 +62,7 @@ async def test_route_a2a_model_bypasses_router():
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
mock_registry,
):
result = await route_request(
await route_request(
data=data,
llm_router=mock_router,
user_model=None,
@ -72,6 +76,632 @@ async def test_route_a2a_model_bypasses_router():
assert call_kwargs["api_base"] == "http://agent.example.com"
@pytest.mark.asyncio
async def test_route_a2a_model_uses_registered_provider():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params={
"custom_llm_provider": "pydantic_ai_agents",
"guardrails": ["agent-guardrail"],
},
static_headers={"Authorization": "Bearer static"},
extra_headers=["X-Tenant"],
)
data = {
"model": "a2a/test-agent",
"messages": [{"role": "user", "content": "Hello"}],
"guardrails": ["request-guardrail"],
"max_tokens": 32,
"temperature": 0.2,
"timeout": 12.0,
"tools": [{"type": "function", "function": {"name": "lookup"}}],
"output_config": {"format": "json"},
"prompt_cache_key": "cache-key",
"safety_identifier": "safety-id",
"proxy_server_request": {
"headers": {
"x-tenant": "tenant-1",
"x-a2a-test-agent-x-run": "run-1",
}
},
}
bridge_response = {
"jsonrpc": "2.0",
"id": "request-id",
"result": {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": "Hello back"}],
"refusal": "I cannot complete that request.",
"messageId": "message-id",
},
}
with (
patch( # test-quality-ok: registry lookup is the routing seam
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch( # test-quality-ok: access control is outside this routing test
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch( # test-quality-ok: provider dispatch is the tested seam
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=bridge_response),
) as bridge,
patch( # test-quality-ok: generic dispatch must stay unused
"litellm.acompletion", AsyncMock()
) as generic_completion,
):
call = await route_a2a_agent_request(data, "acompletion")
response = await call
bridge.assert_awaited_once()
generic_completion.assert_not_called()
assert response.choices[0].message.content == "Hello back"
assert response.choices[0].message.refusal == "I cannot complete that request."
bridge_kwargs = bridge.await_args.kwargs
assert bridge_kwargs["litellm_params"]["max_tokens"] == 32
assert bridge_kwargs["litellm_params"]["temperature"] == 0.2
assert bridge_kwargs["litellm_params"]["timeout"] == 12.0
assert bridge_kwargs["litellm_params"]["tools"] == data["tools"]
assert bridge_kwargs["litellm_params"]["output_config"] == data["output_config"]
assert bridge_kwargs["litellm_params"]["prompt_cache_key"] == data["prompt_cache_key"]
assert bridge_kwargs["litellm_params"]["safety_identifier"] == data["safety_identifier"]
assert bridge_kwargs["litellm_params"]["guardrails"] == ["request-guardrail", "agent-guardrail"]
assert bridge_kwargs["litellm_params"]["extra_headers"] == {
"X-Tenant": "tenant-1",
"x-run": "run-1",
"Authorization": "Bearer static",
}
@pytest.mark.asyncio
async def test_route_a2a_cardless_bedrock_agentcore_uses_registered_model():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={},
litellm_params={
"custom_llm_provider": "bedrock",
"model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123:runtime/test",
},
)
bridge_response = {
"jsonrpc": "2.0",
"id": "request-id",
"result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]},
}
with (
patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=bridge_response),
) as bridge,
):
call = await route_a2a_agent_request(
{"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]},
"acompletion",
)
await call
assert bridge.await_args.kwargs["api_base"] is None
@pytest.mark.asyncio
async def test_route_a2a_registered_provider_uses_configured_api_base_without_card_url():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={},
litellm_params={
"custom_llm_provider": "langflow",
"model": "flow",
"api_base": "https://flow.example.com",
},
)
bridge_response = {
"jsonrpc": "2.0",
"id": "request-id",
"result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]},
}
with (
patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=bridge_response),
) as bridge,
):
call = await route_a2a_agent_request(
{"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]},
"acompletion",
)
await call
assert bridge.await_args.kwargs["api_base"] == "https://flow.example.com"
@pytest.mark.asyncio
async def test_registered_provider_response_preserves_multiple_choices():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params={"custom_llm_provider": "pydantic_ai_agents"},
)
bridge_response = {
"jsonrpc": "2.0",
"id": "request-id",
"choices": [
{"index": 0, "message": {"parts": [{"kind": "text", "text": "first"}]}},
{"index": 1, "message": {"parts": [{"kind": "text", "text": "second"}]}},
],
"result": {},
}
with (
patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=bridge_response),
),
):
call = await route_a2a_agent_request(
{"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]},
"acompletion",
)
response = await call
assert [choice.message.content for choice in response.choices] == ["first", "second"]
@pytest.mark.asyncio
async def test_route_a2a_cardless_watsonx_orchestrate_uses_registered_model():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={},
litellm_params={
"custom_llm_provider": "watsonx_orchestrate",
"model": "agent",
"cp4d_host": "https://wxo.example.com",
"instance_id": "instance",
},
)
bridge_response = {
"jsonrpc": "2.0",
"id": "request-id",
"result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]},
}
with (
patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=bridge_response),
) as bridge,
):
call = await route_a2a_agent_request(
{"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]},
"acompletion",
)
await call
assert bridge.await_args.kwargs["api_base"] is None
@pytest.mark.asyncio
async def test_route_a2a_registered_provider_preserves_identity_headers():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params={"custom_llm_provider": "pydantic_ai_agents"},
)
bridge_response = {
"jsonrpc": "2.0",
"id": "request-id",
"result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]},
}
with (
patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=bridge_response),
) as bridge,
):
call = await route_a2a_agent_request(
{
"model": "a2a/test-agent",
"messages": [{"role": "user", "content": "Hello"}],
"proxy_server_request": {
"headers": {
"x-a2a-test-agent-x-litellm-user-id": "attacker",
"x-a2a-test-agent-x-litellm-team-id": "attacker-team",
}
},
},
"acompletion",
user_api_key_dict=UserAPIKeyAuth(user_id="trusted-user", team_id="trusted-team"),
)
await call
headers = bridge.await_args.kwargs["agent_extra_headers"]
assert headers["X-LiteLLM-User-Id"] == "trusted-user"
assert headers["X-LiteLLM-Team-Id"] == "trusted-team"
assert "x-litellm-user-id" not in {key.lower() for key in headers if key != "X-LiteLLM-User-Id"}
@pytest.mark.asyncio
async def test_route_a2a_registered_provider_preserves_messages_and_session():
from litellm.a2a_protocol.litellm_completion_bridge.handler import A2A_USER_API_KEY_HASH_PARAM
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params={"custom_llm_provider": "langflow", "model": "flow"},
)
data = {
"model": "a2a/test-agent",
"messages": [
{"role": "system", "content": "Be concise"},
{"role": "user", "content": "Hello"},
],
"litellm_session_id": "session-1",
}
bridge_response = {
"jsonrpc": "2.0",
"id": "request-id",
"result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]},
}
with (
patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=bridge_response),
) as bridge,
):
call = await route_a2a_agent_request(
data,
"acompletion",
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
)
await call
bridge_kwargs = bridge.await_args.kwargs
assert bridge_kwargs["params"]["messages"] == data["messages"]
assert bridge_kwargs["params"]["message"]["contextId"] == "session-1"
assert bridge_kwargs["litellm_params"][A2A_USER_API_KEY_HASH_PARAM] == "hashed-key"
@pytest.mark.asyncio
async def test_route_a2a_requires_inbound_trace_id():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params={
"custom_llm_provider": "pydantic_ai_agents",
"require_trace_id_on_calls_to_agent": True,
},
)
with (
patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
):
with pytest.raises(HTTPException, match="requires x-litellm-trace-id"):
await route_a2a_agent_request(
{"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]},
"acompletion",
)
@pytest.mark.asyncio
async def test_route_a2a_resolves_databricks_oauth_headers():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params={"custom_llm_provider": "databricks", "databricks_oauth": {"client_id": "id"}},
)
bridge_response = {
"jsonrpc": "2.0",
"id": "request-id",
"result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]},
}
with (
patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch(
"litellm.proxy.agent_endpoints.databricks_oauth.resolve_databricks_app_auth_header",
AsyncMock(return_value={"Authorization": "Bearer minted"}),
),
patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=bridge_response),
) as bridge,
):
call = await route_a2a_agent_request(
{"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]},
"acompletion",
)
await call
assert bridge.await_args.kwargs["agent_extra_headers"]["Authorization"] == "Bearer minted"
@pytest.mark.asyncio
async def test_registered_provider_response_preserves_tool_calls():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params={"custom_llm_provider": "pydantic_ai_agents"},
)
bridge_response = {
"jsonrpc": "2.0",
"id": "request-id",
"result": {
"kind": "message",
"parts": [],
"tool_calls": [
{
"id": "call-1",
"type": "function",
"function": {"name": "lookup", "arguments": "{}"},
}
],
"finish_reason": "tool_calls",
},
}
with (
patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch(
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=bridge_response),
),
):
call = await route_a2a_agent_request(
{"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]},
"acompletion",
)
response = await call
assert response.choices[0].finish_reason == "tool_calls"
assert response.choices[0].message.tool_calls[0].id == "call-1"
@pytest.mark.asyncio
async def test_a2a_agent_guardrails_merge_before_hooks():
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params={"guardrails": ["agent-guardrail"]},
)
with patch(
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
):
merged = await merge_a2a_agent_guardrails_before_hooks(
{"model": "a2a/test-agent", "guardrails": ["request-guardrail"]}
)
assert merged["guardrails"] == ["request-guardrail", "agent-guardrail"]
@pytest.mark.asyncio
async def test_route_a2a_stream_uses_registered_provider():
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.agents import AgentResponse
agent = AgentResponse(
agent_id="test-agent-id",
agent_name="test-agent",
agent_card_params={"url": "http://agent.example.com"},
litellm_params={"custom_llm_provider": "pydantic_ai_agents"},
)
logging_obj = Mock(spec=Logging)
logging_obj.model_call_details = {}
data = {
"model": "a2a/test-agent",
"messages": [{"role": "user", "content": "Hello"}],
"stream": True,
"litellm_logging_obj": logging_obj,
}
provider_stream = object()
completion_stream = object()
wrapper = object()
with (
patch( # test-quality-ok: registry lookup is the routing seam
"litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
AsyncMock(return_value=agent),
),
patch( # test-quality-ok: access control is outside this routing test
"litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
AsyncMock(return_value=True),
),
patch( # test-quality-ok: provider dispatch is the tested seam
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_streaming",
Mock(return_value=provider_stream),
) as bridge,
patch( # test-quality-ok: iterator wiring is the tested seam
"litellm.llms.a2a.chat.streaming_iterator.A2AModelResponseIterator",
Mock(return_value=completion_stream),
),
patch( # test-quality-ok: wrapper wiring is the tested seam
"litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper",
Mock(return_value=wrapper),
) as stream_wrapper,
patch( # test-quality-ok: generic dispatch must stay unused
"litellm.acompletion", AsyncMock()
) as generic_completion,
):
call = await route_a2a_agent_request(data, "acompletion")
response = await call
bridge.assert_called_once()
stream_wrapper.assert_called_once_with(
completion_stream=completion_stream,
model="a2a/test-agent",
custom_llm_provider="a2a",
logging_obj=logging_obj,
stream_options=None,
)
generic_completion.assert_not_called()
assert response is wrapper
@pytest.mark.asyncio
async def test_registered_provider_logging_uses_provider_model_for_builtin_pricing():
class FakeLogging:
def __init__(self) -> None:
self.model_call_details = {"litellm_params": {}}
self.litellm_params = self.model_call_details["litellm_params"]
self.custom_pricing = False
logging_obj = FakeLogging()
response = {"result": {"message": {"parts": [{"kind": "text", "text": "hello"}]}}}
with (
patch("litellm.litellm_core_utils.litellm_logging.Logging", FakeLogging),
patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=response),
),
):
await _route_registered_provider(
data={
"messages": [{"role": "user", "content": "hello"}],
"litellm_logging_obj": logging_obj,
},
model_name="a2a/agent",
api_base="https://provider.example",
litellm_params={"model": "gpt-4o", "custom_llm_provider": "openai"},
static_headers=None,
)
assert logging_obj.model_call_details["model"] == "gpt-4o"
assert logging_obj.model_call_details["custom_llm_provider"] == "openai"
assert logging_obj.model_call_details["litellm_params"]["model"] == "gpt-4o"
@pytest.mark.asyncio
async def test_native_registered_provider_estimates_usage_when_missing(monkeypatch):
response = {"result": {"message": {"parts": [{"kind": "text", "text": "hello"}]}}}
counter = Mock(side_effect=[3, 2])
monkeypatch.setattr("litellm.utils.token_counter", counter)
with patch(
"litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming",
AsyncMock(return_value=response),
):
result = await _route_registered_provider(
data={"messages": [{"role": "user", "content": "hello"}]},
model_name="a2a/agent",
api_base="https://provider.example",
litellm_params={"model": "agent", "custom_llm_provider": "pydantic_ai_agents"},
static_headers=None,
)
assert result.usage.prompt_tokens == 3
assert result.usage.completion_tokens == 2
assert result.usage.total_tokens == 5
@pytest.mark.asyncio
async def test_route_non_a2a_model_raises_error_if_not_in_router():
"""Test that non-a2a models that aren't in router raise an error"""
@ -141,7 +771,7 @@ def _router_without_models():
@pytest.mark.asyncio
async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_replica(monkeypatch):
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy import proxy_server
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
agent_name = "a2a-sibling-replica-agent"