mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
bba85f0b6c
* chore(lint): remove the LIT002 mutable-construction rule
Drop LIT002 from scripts/check_type_discipline.py along with its helpers,
its budget entry, its unit tests, and the AGENTS.md and gate docstring
mentions. `# mutable-ok` now only suppresses LIT001, so the markers that
only existed to silence LIT002 became LIT013 stale suppressions and are
removed. The files whose layout depended on those trailing comments are
reformatted with ruff format.
Every other LIT rule count is unchanged and the ASTs of all touched
litellm/ files match main apart from one docstring.
* chore(lint): keep the leftover mutable-ok markers for a follow-up
Restore the ~1.4k `# mutable-ok` markers stripped in the previous commit
so this PR only touches the checker, its tests, the budget, and docs.
Those markers no longer suppress anything, so `# mutable-ok` is exempt
from LIT013 until a follow-up strips them.
* Revert "chore(lint): keep the leftover mutable-ok markers for a follow-up"
This reverts commit c35bc0b84e.
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
890 lines
33 KiB
Python
890 lines
33 KiB
Python
# pyright: reportUnknownArgumentType=false
|
|
# a2a-sdk (and its protobuf-generated compat conversions) ships no usable types for
|
|
# the call surface used here, so SDK calls take Unknown-typed arguments. This module
|
|
# is dedicated to the A2A SDK boundary; the rule is off file-wide instead of
|
|
# scattering per-line ignores across every SDK call.
|
|
"""
|
|
LiteLLM A2A SDK functions.
|
|
|
|
Provides standalone functions with @client decorator for LiteLLM logging integration.
|
|
"""
|
|
|
|
import asyncio
|
|
import datetime
|
|
import uuid
|
|
from collections.abc import AsyncIterator, Coroutine, Mapping
|
|
from types import MappingProxyType, ModuleType
|
|
from typing import TYPE_CHECKING, Any, Final, Optional, cast
|
|
|
|
import litellm
|
|
from litellm._logging import verbose_logger, verbose_proxy_logger
|
|
from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator
|
|
from litellm.a2a_protocol.utils import A2ARequestUtils
|
|
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
|
|
from litellm.litellm_core_utils.asyncify import asyncify
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
from litellm.llms.custom_httpx.http_handler import (
|
|
get_async_httpx_client,
|
|
httpxSpecialProvider,
|
|
)
|
|
from litellm.types.agents import LiteLLMSendMessageResponse
|
|
from litellm.utils import client
|
|
|
|
if TYPE_CHECKING:
|
|
from a2a.client import Client as A2AClientType
|
|
from a2a.client import ClientCallContext as A2ACallContextType
|
|
from a2a.compat.v0_3.types import (
|
|
AgentCard,
|
|
Message,
|
|
SendMessageRequest,
|
|
SendMessageResponse,
|
|
SendStreamingMessageRequest,
|
|
SendStreamingMessageResponse,
|
|
SendStreamingMessageSuccessResponse,
|
|
Task,
|
|
)
|
|
from a2a.types.a2a_pb2 import SendMessageRequest as CoreSendMessageRequest
|
|
from a2a.types.a2a_pb2 import StreamResponse as CoreStreamResponse
|
|
|
|
# Runtime imports — requires a2a-sdk>=1.1.0
|
|
A2A_SDK_AVAILABLE = False
|
|
_a2a_conversions: ModuleType | None = None
|
|
|
|
try:
|
|
from a2a.client import Client, ClientCallContext, ClientConfig, create_client
|
|
from a2a.compat.v0_3 import conversions as _a2a_conversions
|
|
from a2a.compat.v0_3.types import (
|
|
Message,
|
|
SendMessageRequest,
|
|
SendMessageResponse,
|
|
SendMessageSuccessResponse,
|
|
SendStreamingMessageRequest,
|
|
SendStreamingMessageResponse,
|
|
Task,
|
|
)
|
|
|
|
A2A_SDK_AVAILABLE = True
|
|
except ImportError:
|
|
Client = None
|
|
ClientCallContext = None
|
|
ClientConfig = None
|
|
create_client = None
|
|
|
|
# Import our custom card resolver that supports multiple well-known paths
|
|
from litellm.a2a_protocol.card_resolver import (
|
|
AGENT_CARD_PATH_PARAM,
|
|
LiteLLMA2ACardResolver,
|
|
get_agent_card_url,
|
|
normalize_agent_card_interfaces,
|
|
)
|
|
from litellm.a2a_protocol.exception_mapping_utils import (
|
|
handle_a2a_localhost_retry,
|
|
map_a2a_exception,
|
|
)
|
|
from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
|
|
|
|
# Use our custom resolver instead of the default A2A SDK resolver
|
|
A2ACardResolver: Final = LiteLLMA2ACardResolver
|
|
|
|
|
|
def _set_usage_on_logging_obj(
|
|
kwargs: Mapping[str, object],
|
|
prompt_tokens: int,
|
|
completion_tokens: int,
|
|
) -> None:
|
|
"""
|
|
Set usage on litellm_logging_obj for standard logging payload.
|
|
|
|
Args:
|
|
kwargs: The kwargs dict containing litellm_logging_obj
|
|
prompt_tokens: Number of input tokens
|
|
completion_tokens: Number of output tokens
|
|
"""
|
|
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
|
|
if isinstance(litellm_logging_obj, Logging):
|
|
usage: Final = litellm.Usage(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
total_tokens=prompt_tokens + completion_tokens,
|
|
)
|
|
litellm_logging_obj.model_call_details["usage"] = usage
|
|
|
|
|
|
def _set_agent_id_on_logging_obj(
|
|
kwargs: Mapping[str, object],
|
|
agent_id: str | None,
|
|
) -> None:
|
|
"""
|
|
Set agent_id on litellm_logging_obj for SpendLogs tracking.
|
|
|
|
Args:
|
|
kwargs: The kwargs dict containing litellm_logging_obj
|
|
agent_id: The A2A agent ID
|
|
"""
|
|
if agent_id is None:
|
|
return
|
|
|
|
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
|
|
if isinstance(litellm_logging_obj, Logging):
|
|
# Set agent_id directly on model_call_details (same pattern as custom_llm_provider)
|
|
litellm_logging_obj.model_call_details["agent_id"] = agent_id
|
|
|
|
|
|
_A2A_COST_PARAM_KEYS: Final = ("cost_per_query", "input_cost_per_token", "output_cost_per_token")
|
|
|
|
|
|
def _a2a_cost_params(litellm_params: Mapping[str, object] | None) -> Mapping[str, object]:
|
|
"""Only the agent's pricing keys reach the logging object; its credentials never do."""
|
|
return MappingProxyType(
|
|
{
|
|
key: litellm_params[key]
|
|
for key in _A2A_COST_PARAM_KEYS
|
|
if litellm_params is not None and litellm_params.get(key) is not None
|
|
}
|
|
)
|
|
|
|
|
|
def _card_http_kwargs(extra_headers: dict[str, str] | None) -> dict[str, object] | None:
|
|
return {"headers": extra_headers} if extra_headers else None
|
|
|
|
|
|
def _agent_card_path(litellm_params: Mapping[str, object]) -> str | None:
|
|
configured_path: Final = litellm_params.get(AGENT_CARD_PATH_PARAM)
|
|
return configured_path if isinstance(configured_path, str) and configured_path else None
|
|
|
|
|
|
def _set_litellm_params_on_logging_obj(
|
|
kwargs: Mapping[str, object],
|
|
litellm_params: Mapping[str, object],
|
|
) -> None:
|
|
"""
|
|
Merge the agent's pricing params into model_call_details["litellm_params"]
|
|
so A2ACostCalculator can read them.
|
|
|
|
The non-streaming path reuses the proxy-built logging object, whose
|
|
litellm_params already carries metadata / proxy_server_request / user-key
|
|
context, so merge the pricing keys in rather than replacing the dict.
|
|
"""
|
|
logging_obj: Final = kwargs.get("litellm_logging_obj")
|
|
if not isinstance(logging_obj, Logging):
|
|
return
|
|
|
|
cost_params: Final = _a2a_cost_params(litellm_params)
|
|
if not cost_params:
|
|
return
|
|
|
|
logging_obj.model_call_details["litellm_params"] = {
|
|
**(logging_obj.model_call_details.get("litellm_params") or {}),
|
|
**cost_params,
|
|
}
|
|
|
|
|
|
def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: Mapping[str, object]) -> str:
|
|
"""
|
|
Extract agent info and set model/custom_llm_provider for cost tracking.
|
|
|
|
Sets model info on the litellm_logging_obj if available.
|
|
Returns the agent name for logging.
|
|
"""
|
|
agent_name = "unknown"
|
|
|
|
agent_card: Final = _get_a2a_client_agent_card(a2a_client)
|
|
|
|
if agent_card is not None:
|
|
agent_name = getattr(agent_card, "name", "unknown") or "unknown"
|
|
|
|
# Build model string
|
|
model: Final = f"a2a_agent/{agent_name}"
|
|
custom_llm_provider: Final = "a2a_agent"
|
|
|
|
# Set on litellm_logging_obj if available (for standard logging payload)
|
|
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
|
|
if isinstance(litellm_logging_obj, Logging):
|
|
litellm_logging_obj.model = model
|
|
litellm_logging_obj.custom_llm_provider = custom_llm_provider
|
|
litellm_logging_obj.model_call_details["model"] = model
|
|
litellm_logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
|
|
|
|
return agent_name
|
|
|
|
|
|
def _get_a2a_client_agent_card(a2a_client: "A2AClientType") -> Optional["AgentCard"]:
|
|
agent_card = cast(Optional["AgentCard"], getattr(a2a_client, "_litellm_agent_card", None))
|
|
if agent_card is not None:
|
|
return agent_card
|
|
agent_card = cast(Optional["AgentCard"], getattr(a2a_client, "agent_card", None))
|
|
if agent_card is not None:
|
|
return agent_card
|
|
return cast(Optional["AgentCard"], getattr(a2a_client, "_card", None))
|
|
|
|
|
|
async def _send_message_via_completion_bridge(
|
|
request: "SendMessageRequest",
|
|
custom_llm_provider: object,
|
|
api_base: str | None,
|
|
litellm_params: dict[str, object],
|
|
agent_extra_headers: dict[str, str] | None = None,
|
|
) -> LiteLLMSendMessageResponse:
|
|
"""
|
|
Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore).
|
|
|
|
Requires request; api_base is optional for providers that derive endpoint from model.
|
|
"""
|
|
verbose_logger.info("A2A using completion bridge: provider=%s, api_base=%s", custom_llm_provider, api_base)
|
|
|
|
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
|
A2ACompletionBridgeHandler,
|
|
)
|
|
|
|
params = request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params)
|
|
|
|
response_dict: Final = await A2ACompletionBridgeHandler.handle_non_streaming(
|
|
request_id=str(request.id),
|
|
params=params,
|
|
litellm_params=litellm_params,
|
|
api_base=api_base,
|
|
agent_extra_headers=agent_extra_headers,
|
|
)
|
|
|
|
return LiteLLMSendMessageResponse.from_dict(response_dict, request_id=str(request.id))
|
|
|
|
|
|
def _get_a2a_call_context(a2a_client: "A2AClientType") -> Optional["A2ACallContextType"]:
|
|
return getattr(a2a_client, "_litellm_call_context", None)
|
|
|
|
|
|
def _to_core_send_message_request(request: "SendMessageRequest") -> "CoreSendMessageRequest":
|
|
from a2a.compat.v0_3 import conversions
|
|
|
|
return conversions.to_core_send_message_request(request)
|
|
|
|
|
|
def _to_compat_stream_response(
|
|
event: "CoreStreamResponse", request_id: str | int
|
|
) -> "SendStreamingMessageSuccessResponse":
|
|
from a2a.compat.v0_3 import conversions
|
|
|
|
return conversions.to_compat_stream_response(event, request_id=request_id)
|
|
|
|
|
|
async def _send_message(a2a_client: "A2AClientType", request: "SendMessageRequest") -> "SendMessageResponse":
|
|
"""Send a non-streaming message via a2a-sdk 1.x and return JSON-RPC response."""
|
|
if _a2a_conversions is None:
|
|
raise ImportError(
|
|
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
|
|
)
|
|
|
|
pb_request: Final = _to_core_send_message_request(request)
|
|
last_event = None
|
|
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
|
|
last_event = event
|
|
if last_event is None:
|
|
raise RuntimeError("A2A send_message failed: no response received from agent.")
|
|
|
|
stream_compat: Final = _to_compat_stream_response(last_event, request_id=request.id)
|
|
result: Final = stream_compat.result
|
|
if not isinstance(result, (Message, Task)):
|
|
raise RuntimeError(
|
|
"A2A send_message failed: non-streaming message/send expects the "
|
|
"agent's final event to be a Message or Task result."
|
|
)
|
|
return SendMessageResponse(
|
|
root=SendMessageSuccessResponse(
|
|
id=request.id,
|
|
result=result,
|
|
)
|
|
)
|
|
|
|
|
|
async def _execute_a2a_send_with_retry(
|
|
a2a_client: "A2AClientType",
|
|
request: "SendMessageRequest",
|
|
agent_card: Optional["AgentCard"],
|
|
card_url: str | None,
|
|
api_base: str | None,
|
|
agent_name: str | None,
|
|
) -> "SendMessageResponse":
|
|
"""Send an A2A message with retry logic for localhost URL errors."""
|
|
a2a_response = None
|
|
for _ in range(2): # max 2 attempts: original + 1 retry
|
|
try:
|
|
a2a_response = await _send_message(a2a_client, request)
|
|
break # success, exit retry loop
|
|
except A2ALocalhostURLError as e:
|
|
a2a_client = await handle_a2a_localhost_retry(
|
|
error=e,
|
|
agent_card=agent_card,
|
|
a2a_client=a2a_client,
|
|
is_streaming=False,
|
|
)
|
|
card_url = get_agent_card_url(agent_card) if agent_card else None
|
|
except Exception as e:
|
|
try:
|
|
map_a2a_exception(e, card_url, api_base, model=agent_name)
|
|
except A2ALocalhostURLError as localhost_err:
|
|
a2a_client = await handle_a2a_localhost_retry(
|
|
error=localhost_err,
|
|
agent_card=agent_card,
|
|
a2a_client=a2a_client,
|
|
is_streaming=False,
|
|
)
|
|
card_url = get_agent_card_url(agent_card) if agent_card else None
|
|
continue
|
|
except Exception:
|
|
raise
|
|
if a2a_response is None:
|
|
raise RuntimeError("A2A send_message failed: no response received after retry attempts.")
|
|
return a2a_response
|
|
|
|
|
|
async def _stream_messages(
|
|
a2a_client: "A2AClientType", request: "SendStreamingMessageRequest"
|
|
) -> AsyncIterator["SendStreamingMessageResponse"]:
|
|
"""Stream message events via a2a-sdk 1.x and yield JSON-RPC chunks."""
|
|
if _a2a_conversions is None:
|
|
raise ImportError(
|
|
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
|
|
)
|
|
|
|
pb_request: Final[CoreSendMessageRequest] = _a2a_conversions.to_core_send_message_request(request)
|
|
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
|
|
compat_chunk = _to_compat_stream_response(event, request_id=request.id)
|
|
yield SendStreamingMessageResponse(root=compat_chunk)
|
|
|
|
|
|
async def _execute_a2a_stream_with_retry(
|
|
a2a_client: "A2AClientType",
|
|
request: "SendStreamingMessageRequest",
|
|
agent_card: Optional["AgentCard"],
|
|
card_url: str | None,
|
|
api_base: str | None,
|
|
agent_name: str | None,
|
|
) -> AsyncIterator["SendStreamingMessageResponse"]:
|
|
"""Stream an A2A message with retry logic for localhost URL errors."""
|
|
response_started = False
|
|
stream_succeeded = False
|
|
for _ in range(2): # max 2 attempts: original + 1 retry
|
|
try:
|
|
async for chunk in _stream_messages(a2a_client, request):
|
|
response_started = True
|
|
yield chunk
|
|
stream_succeeded = True
|
|
return
|
|
except A2ALocalhostURLError as e:
|
|
if response_started:
|
|
raise
|
|
a2a_client = await handle_a2a_localhost_retry(
|
|
error=e,
|
|
agent_card=agent_card,
|
|
a2a_client=a2a_client,
|
|
is_streaming=True,
|
|
)
|
|
card_url = get_agent_card_url(agent_card) if agent_card else None
|
|
continue
|
|
except Exception as e:
|
|
if response_started:
|
|
raise
|
|
try:
|
|
map_a2a_exception(e, card_url, api_base, model=agent_name)
|
|
except A2ALocalhostURLError as localhost_err:
|
|
a2a_client = await handle_a2a_localhost_retry(
|
|
error=localhost_err,
|
|
agent_card=agent_card,
|
|
a2a_client=a2a_client,
|
|
is_streaming=True,
|
|
)
|
|
card_url = get_agent_card_url(agent_card) if agent_card else None
|
|
continue
|
|
raise
|
|
if not stream_succeeded:
|
|
raise RuntimeError("A2A send_message_streaming failed: no response received after retry attempts.")
|
|
|
|
|
|
@client
|
|
async def asend_message(
|
|
a2a_client: Optional["A2AClientType"] = None,
|
|
request: Optional["SendMessageRequest"] = None,
|
|
api_base: str | None = None,
|
|
litellm_params: dict[str, object] | None = None,
|
|
agent_id: str | None = None,
|
|
agent_extra_headers: dict[str, str] | None = None,
|
|
**kwargs: object,
|
|
) -> LiteLLMSendMessageResponse:
|
|
"""
|
|
Async: Send a message to an A2A agent.
|
|
|
|
Uses the @client decorator for LiteLLM logging and tracking.
|
|
If litellm_params contains custom_llm_provider, routes through the completion bridge.
|
|
|
|
Args:
|
|
a2a_client: An initialized a2a.client.A2AClient instance (optional if using completion bridge)
|
|
request: SendMessageRequest from a2a.types (optional if using completion bridge with api_base)
|
|
api_base: API base URL (required for completion bridge, optional for standard A2A)
|
|
litellm_params: Optional dict with custom_llm_provider, model, etc. for completion bridge
|
|
agent_id: Optional agent ID for tracking in SpendLogs
|
|
**kwargs: Additional arguments passed to the client decorator
|
|
|
|
Returns:
|
|
LiteLLMSendMessageResponse (wraps a2a SendMessageResponse with _hidden_params)
|
|
|
|
Example (standard A2A):
|
|
```python
|
|
from litellm.a2a_protocol import asend_message, create_a2a_client
|
|
from a2a.types import SendMessageRequest, MessageSendParams
|
|
from uuid import uuid4
|
|
|
|
a2a_client = await create_a2a_client(base_url="http://localhost:10001")
|
|
request = SendMessageRequest(
|
|
id=str(uuid4()),
|
|
params=MessageSendParams(
|
|
message={"role": "user", "parts": [{"kind": "text", "text": "Hello!"}], "messageId": uuid4().hex}
|
|
)
|
|
)
|
|
response = await asend_message(a2a_client=a2a_client, request=request)
|
|
```
|
|
|
|
Example (completion bridge with LangGraph):
|
|
```python
|
|
from litellm.a2a_protocol import asend_message
|
|
from a2a.types import SendMessageRequest, MessageSendParams
|
|
from uuid import uuid4
|
|
|
|
request = SendMessageRequest(
|
|
id=str(uuid4()),
|
|
params=MessageSendParams(
|
|
message={"role": "user", "parts": [{"kind": "text", "text": "Hello!"}], "messageId": uuid4().hex}
|
|
)
|
|
)
|
|
response = await asend_message(
|
|
request=request,
|
|
api_base="http://localhost:2024",
|
|
litellm_params={"custom_llm_provider": "langgraph", "model": "agent"},
|
|
)
|
|
```
|
|
"""
|
|
litellm_params = litellm_params or {}
|
|
logging_obj: Final = kwargs.get("litellm_logging_obj")
|
|
trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None
|
|
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
|
|
|
|
# Route through completion bridge if custom_llm_provider is set
|
|
if custom_llm_provider:
|
|
if request is None:
|
|
raise ValueError("request is required for completion bridge")
|
|
return await _send_message_via_completion_bridge(
|
|
request=request,
|
|
custom_llm_provider=custom_llm_provider,
|
|
api_base=api_base,
|
|
litellm_params=litellm_params,
|
|
agent_extra_headers=agent_extra_headers,
|
|
)
|
|
|
|
# Standard A2A client flow
|
|
if request is None:
|
|
raise ValueError("request is required")
|
|
|
|
# Create A2A client if not provided but api_base is available
|
|
if a2a_client is None:
|
|
if api_base is None:
|
|
raise ValueError("Either a2a_client or api_base is required for standard A2A flow")
|
|
trace_id = trace_id or str(uuid.uuid4())
|
|
extra_headers: Final[dict[str, str]] = {"X-LiteLLM-Trace-Id": trace_id}
|
|
if agent_id:
|
|
extra_headers["X-LiteLLM-Agent-Id"] = agent_id
|
|
# Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones)
|
|
if agent_extra_headers:
|
|
extra_headers.update(agent_extra_headers)
|
|
a2a_client = await create_a2a_client(
|
|
base_url=api_base,
|
|
extra_headers=extra_headers,
|
|
relative_card_path=_agent_card_path(litellm_params),
|
|
)
|
|
|
|
# Type assertion: a2a_client is guaranteed to be non-None here
|
|
assert a2a_client is not None
|
|
|
|
agent_name: Final = _get_a2a_model_info(a2a_client, kwargs)
|
|
|
|
verbose_logger.info("A2A send_message request_id=%s, agent=%s", request.id, agent_name)
|
|
|
|
# Get agent card URL for localhost retry logic
|
|
agent_card: Final = _get_a2a_client_agent_card(a2a_client)
|
|
card_url: Final = get_agent_card_url(agent_card) if agent_card else None
|
|
|
|
a2a_response: Final = await _execute_a2a_send_with_retry(
|
|
a2a_client=a2a_client,
|
|
request=request,
|
|
agent_card=agent_card,
|
|
card_url=card_url,
|
|
api_base=api_base,
|
|
agent_name=agent_name,
|
|
)
|
|
|
|
verbose_logger.info("A2A send_message completed, request_id=%s", request.id)
|
|
|
|
# Wrap in LiteLLM response type for _hidden_params support
|
|
response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id))
|
|
|
|
# Calculate token usage from request and response
|
|
response_dict: Final[dict[str, object]] = a2a_response.root.model_dump(mode="json", exclude_none=True)
|
|
(
|
|
prompt_tokens,
|
|
completion_tokens,
|
|
_,
|
|
) = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)(
|
|
request=request,
|
|
response_dict=response_dict,
|
|
)
|
|
|
|
# Set usage on logging obj for standard logging payload
|
|
_set_usage_on_logging_obj(
|
|
kwargs=kwargs,
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
)
|
|
|
|
# Merge agent pricing params into the logging obj so cost is calculated
|
|
_set_litellm_params_on_logging_obj(kwargs=kwargs, litellm_params=litellm_params)
|
|
|
|
# Set agent_id on logging obj for SpendLogs tracking
|
|
_set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id)
|
|
|
|
return response
|
|
|
|
|
|
@client
|
|
def send_message(
|
|
a2a_client: "A2AClientType",
|
|
request: "SendMessageRequest",
|
|
**kwargs: Any,
|
|
) -> LiteLLMSendMessageResponse | Coroutine[object, object, LiteLLMSendMessageResponse]:
|
|
"""
|
|
Sync: Send a message to an A2A agent.
|
|
|
|
Uses the @client decorator for LiteLLM logging and tracking.
|
|
|
|
Args:
|
|
a2a_client: An initialized a2a.client.A2AClient instance
|
|
request: SendMessageRequest from a2a.types
|
|
**kwargs: Additional arguments passed to the client decorator
|
|
|
|
Returns:
|
|
LiteLLMSendMessageResponse (wraps a2a SendMessageResponse with _hidden_params)
|
|
"""
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
loop = None
|
|
|
|
if loop is not None:
|
|
return asend_message(a2a_client=a2a_client, request=request, **kwargs)
|
|
else:
|
|
return asyncio.run(asend_message(a2a_client=a2a_client, request=request, **kwargs))
|
|
|
|
|
|
def _build_streaming_logging_obj(
|
|
request: "SendStreamingMessageRequest",
|
|
agent_name: str,
|
|
agent_id: str | None,
|
|
litellm_params: dict[str, object] | None,
|
|
metadata: dict[str, object] | None,
|
|
proxy_server_request: dict[str, object] | None,
|
|
) -> Logging:
|
|
"""Build logging object for streaming A2A requests."""
|
|
start_time: Final = datetime.datetime.now()
|
|
model: Final = f"a2a_agent/{agent_name}"
|
|
|
|
logging_obj: Final = Logging(
|
|
model=model,
|
|
messages=[{"role": "user", "content": "streaming-request"}],
|
|
stream=False,
|
|
call_type="asend_message_streaming",
|
|
start_time=start_time,
|
|
litellm_call_id=str(request.id),
|
|
function_id=str(request.id),
|
|
)
|
|
logging_obj.model = model
|
|
logging_obj.custom_llm_provider = "a2a_agent"
|
|
logging_obj.model_call_details["model"] = model
|
|
logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
|
|
logging_obj.model_call_details["call_type"] = logging_obj.call_type
|
|
if agent_id:
|
|
logging_obj.model_call_details["agent_id"] = agent_id
|
|
|
|
_request_context: Final = (("metadata", metadata), ("proxy_server_request", proxy_server_request))
|
|
_litellm_params: Final = dict(
|
|
(*_a2a_cost_params(litellm_params).items(), *((key, value) for key, value in _request_context if value))
|
|
)
|
|
|
|
logging_obj.litellm_params = _litellm_params
|
|
logging_obj.optional_params = _litellm_params
|
|
logging_obj.model_call_details["litellm_params"] = _litellm_params
|
|
logging_obj.model_call_details["metadata"] = metadata or {}
|
|
|
|
return logging_obj
|
|
|
|
|
|
async def asend_message_streaming(
|
|
a2a_client: Optional["A2AClientType"] = None,
|
|
request: Optional["SendStreamingMessageRequest"] = None,
|
|
api_base: str | None = None,
|
|
litellm_params: dict[str, object] | None = None,
|
|
agent_id: str | None = None,
|
|
metadata: dict[str, object] | None = None,
|
|
proxy_server_request: dict[str, object] | None = None,
|
|
agent_extra_headers: dict[str, str] | None = None,
|
|
**kwargs: object,
|
|
) -> AsyncIterator[Any]:
|
|
"""
|
|
Async: Send a streaming message to an A2A agent.
|
|
|
|
If litellm_params contains custom_llm_provider, routes through the completion bridge.
|
|
|
|
Args:
|
|
a2a_client: An initialized a2a.client.A2AClient instance (optional if using completion bridge)
|
|
request: SendStreamingMessageRequest from a2a.types
|
|
api_base: API base URL (required for completion bridge)
|
|
litellm_params: Optional dict with custom_llm_provider, model, etc. for completion bridge
|
|
agent_id: Optional agent ID for tracking in SpendLogs
|
|
metadata: Optional metadata dict (contains user_api_key, user_id, team_id, etc.)
|
|
proxy_server_request: Optional proxy server request data
|
|
|
|
Yields:
|
|
SendStreamingMessageResponse chunks from the agent
|
|
|
|
Example (completion bridge with LangGraph):
|
|
```python
|
|
from litellm.a2a_protocol import asend_message_streaming
|
|
from a2a.types import SendStreamingMessageRequest, MessageSendParams
|
|
from uuid import uuid4
|
|
|
|
request = SendStreamingMessageRequest(
|
|
id=str(uuid4()),
|
|
params=MessageSendParams(
|
|
message={"role": "user", "parts": [{"kind": "text", "text": "Hello!"}], "messageId": uuid4().hex}
|
|
)
|
|
)
|
|
async for chunk in asend_message_streaming(
|
|
request=request,
|
|
api_base="http://localhost:2024",
|
|
litellm_params={"custom_llm_provider": "langgraph", "model": "agent"},
|
|
):
|
|
print(chunk)
|
|
```
|
|
"""
|
|
litellm_params = litellm_params or {}
|
|
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
|
|
|
|
# Route through completion bridge if custom_llm_provider is set
|
|
if custom_llm_provider:
|
|
if request is None:
|
|
raise ValueError("request is required for completion bridge")
|
|
# api_base is optional for providers that derive endpoint from model (e.g., bedrock/agentcore)
|
|
|
|
verbose_logger.info("A2A streaming using completion bridge: provider=%s", custom_llm_provider)
|
|
|
|
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
|
A2ACompletionBridgeHandler,
|
|
)
|
|
|
|
# Extract params from request
|
|
params: Final = (
|
|
request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params)
|
|
)
|
|
|
|
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
|
|
request_id=str(request.id),
|
|
params=params,
|
|
litellm_params=litellm_params,
|
|
api_base=api_base,
|
|
agent_extra_headers=agent_extra_headers,
|
|
):
|
|
yield chunk
|
|
return
|
|
|
|
if request is None:
|
|
raise ValueError("request is required")
|
|
|
|
_raw_logging_obj: Final = kwargs.get("litellm_logging_obj")
|
|
logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None
|
|
|
|
if a2a_client is None:
|
|
if api_base is None:
|
|
raise ValueError("Either a2a_client or api_base is required for standard A2A flow")
|
|
logging_trace_id: Final = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None
|
|
trace_id: Final = logging_trace_id or (str(request.id) if request.id else str(uuid.uuid4()))
|
|
extra_headers: Final[dict[str, str]] = {"X-LiteLLM-Trace-Id": trace_id}
|
|
if agent_id:
|
|
extra_headers["X-LiteLLM-Agent-Id"] = agent_id
|
|
if agent_extra_headers:
|
|
extra_headers.update(agent_extra_headers)
|
|
a2a_client = await create_a2a_client(
|
|
base_url=api_base,
|
|
extra_headers=extra_headers,
|
|
streaming=True,
|
|
relative_card_path=_agent_card_path(litellm_params),
|
|
)
|
|
|
|
assert a2a_client is not None
|
|
|
|
agent_name: Final = _get_a2a_model_info(a2a_client, kwargs)
|
|
|
|
if logging_obj is None:
|
|
logging_obj = _build_streaming_logging_obj(
|
|
request=request,
|
|
agent_name=agent_name,
|
|
agent_id=agent_id,
|
|
litellm_params=litellm_params,
|
|
metadata=metadata,
|
|
proxy_server_request=proxy_server_request,
|
|
)
|
|
|
|
verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name)
|
|
|
|
agent_card: Final = _get_a2a_client_agent_card(a2a_client)
|
|
card_url: Final = get_agent_card_url(agent_card) if agent_card else None
|
|
|
|
stream: Final = _execute_a2a_stream_with_retry(
|
|
a2a_client=a2a_client,
|
|
request=request,
|
|
agent_card=agent_card,
|
|
card_url=card_url,
|
|
api_base=api_base,
|
|
agent_name=agent_name,
|
|
)
|
|
|
|
_set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id)
|
|
|
|
async for chunk in A2AStreamingIterator(
|
|
stream=stream,
|
|
request=request,
|
|
logging_obj=logging_obj,
|
|
agent_name=agent_name,
|
|
):
|
|
yield chunk
|
|
|
|
|
|
async def create_a2a_client(
|
|
base_url: str,
|
|
timeout: float = DEFAULT_A2A_AGENT_TIMEOUT,
|
|
extra_headers: dict[str, str] | None = None,
|
|
streaming: bool = False,
|
|
relative_card_path: str | None = None,
|
|
) -> "A2AClientType":
|
|
"""
|
|
Create an A2A client for the given agent URL.
|
|
|
|
This resolves the agent card and returns a ready-to-use A2A client.
|
|
The client can be reused for multiple requests.
|
|
|
|
Args:
|
|
base_url: The base URL of the A2A agent (e.g., "http://localhost:10001")
|
|
timeout: Request timeout in seconds (default: ``DEFAULT_A2A_AGENT_TIMEOUT`` / env ``DEFAULT_A2A_AGENT_TIMEOUT``)
|
|
extra_headers: Optional additional headers to include in requests
|
|
relative_card_path: Optional card path relative to ``base_url`` (e.g. ``agentCard/v1.0`` for a
|
|
Microsoft Foundry agent); when None the well-known paths are probed in order
|
|
|
|
Returns:
|
|
An initialized a2a.client.A2AClient instance
|
|
|
|
Example:
|
|
```python
|
|
from litellm.a2a_protocol import create_a2a_client, asend_message
|
|
|
|
# Create client once
|
|
client = await create_a2a_client(base_url="http://localhost:10001")
|
|
|
|
# Reuse for multiple requests
|
|
response1 = await asend_message(a2a_client=client, request=request1)
|
|
response2 = await asend_message(a2a_client=client, request=request2)
|
|
```
|
|
"""
|
|
if not A2A_SDK_AVAILABLE:
|
|
raise ImportError(
|
|
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
|
|
)
|
|
|
|
verbose_logger.info("Creating A2A client for %s", base_url)
|
|
|
|
_async_handler: Final = get_async_httpx_client(
|
|
llm_provider=httpxSpecialProvider.A2AProvider,
|
|
params={"timeout": timeout},
|
|
)
|
|
httpx_client: Final = _async_handler.client
|
|
if extra_headers:
|
|
verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys()))
|
|
|
|
resolver: Final = A2ACardResolver(httpx_client=httpx_client, base_url=base_url)
|
|
agent_card: Final = normalize_agent_card_interfaces(
|
|
await resolver.get_agent_card(
|
|
relative_card_path=relative_card_path,
|
|
http_kwargs=_card_http_kwargs(extra_headers),
|
|
)
|
|
)
|
|
|
|
a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall]
|
|
agent_card,
|
|
client_config=ClientConfig( # pyright: ignore[reportOptionalCall]
|
|
httpx_client=httpx_client,
|
|
streaming=streaming,
|
|
),
|
|
)
|
|
# Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse
|
|
# the configured httpx client and this agent's headers without excavating
|
|
# a2a-sdk private internals.
|
|
a2a_client._litellm_httpx_client = httpx_client
|
|
a2a_client._litellm_call_context = ( # pyright: ignore[reportAttributeAccessIssue] # LiteLLM-owned stash
|
|
ClientCallContext(service_parameters=extra_headers) # pyright: ignore[reportOptionalCall] # SDK checked above
|
|
if extra_headers
|
|
else None
|
|
)
|
|
a2a_client._litellm_agent_card = agent_card
|
|
|
|
verbose_logger.info("A2A client created for %s", base_url)
|
|
|
|
return a2a_client
|
|
|
|
|
|
async def aget_agent_card(
|
|
base_url: str,
|
|
timeout: float = DEFAULT_A2A_AGENT_TIMEOUT,
|
|
extra_headers: dict[str, str] | None = None,
|
|
relative_card_path: str | None = None,
|
|
) -> "AgentCard":
|
|
"""
|
|
Fetch the agent card from an A2A agent.
|
|
|
|
Args:
|
|
base_url: The base URL of the A2A agent (e.g., "http://localhost:10001")
|
|
timeout: Request timeout in seconds (default: ``DEFAULT_A2A_AGENT_TIMEOUT`` / env ``DEFAULT_A2A_AGENT_TIMEOUT``)
|
|
extra_headers: Optional additional headers to include in requests
|
|
relative_card_path: Optional card path relative to ``base_url``; when None the well-known paths are probed
|
|
|
|
Returns:
|
|
AgentCard from the A2A agent
|
|
"""
|
|
if not A2A_SDK_AVAILABLE:
|
|
raise ImportError(
|
|
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
|
|
)
|
|
|
|
verbose_logger.info("Fetching agent card from %s", base_url)
|
|
|
|
# Use LiteLLM's cached httpx client
|
|
http_handler: Final = get_async_httpx_client(
|
|
llm_provider=httpxSpecialProvider.A2A,
|
|
params={"timeout": timeout},
|
|
)
|
|
httpx_client: Final = http_handler.client
|
|
|
|
resolver: Final = A2ACardResolver(
|
|
httpx_client=httpx_client,
|
|
base_url=base_url,
|
|
)
|
|
agent_card: Final = await resolver.get_agent_card(
|
|
relative_card_path=relative_card_path,
|
|
http_kwargs=_card_http_kwargs(extra_headers),
|
|
)
|
|
|
|
verbose_logger.info("Fetched agent card: %s", agent_card.name if hasattr(agent_card, "name") else "unknown")
|
|
return agent_card
|