diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 6d3b09aed68..1e003a6e4bc 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -22,7 +22,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( A2AStreamingContext, ) from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager -from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS +from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS, A2A_COMPLETION_BRIDGE_WORKING_MESSAGE from litellm.interactions.agents.utils import merge_agent_headers from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import ModelResponse @@ -268,7 +268,7 @@ class A2ACompletionBridgeHandler: ctx=ctx, state="working", final=False, - message_text="Processing request...", + message_text=A2A_COMPLETION_BRIDGE_WORKING_MESSAGE, ) yield working_event diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index f44e3f6a81b..dcbe7a6506e 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -191,6 +191,23 @@ def _set_litellm_params_on_logging_obj( } +async def _record_send_message_billing( + kwargs: Mapping[str, object], + request: "SendMessageRequest", + response_dict: Mapping[str, object], + litellm_params: Mapping[str, object], + agent_id: str | None, +) -> None: + prompt_tokens, completion_tokens, _ = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)( + request=request, + response_dict=response_dict, + ) + _set_usage_on_logging_obj(kwargs=kwargs, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens) + if "error" not in response_dict: + _set_litellm_params_on_logging_obj(kwargs=kwargs, litellm_params=litellm_params) + _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id) + + 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. @@ -483,13 +500,21 @@ async def asend_message( if custom_llm_provider: if request is None: raise ValueError("request is required for completion bridge") - return await _send_message_via_completion_bridge( + bridge_response: Final = 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, ) + await _record_send_message_billing( + kwargs=kwargs, + request=request, + response_dict=bridge_response.model_dump(mode="json", exclude_none=True), + litellm_params=litellm_params, + agent_id=agent_id, + ) + return bridge_response # Standard A2A client flow if request is None: @@ -538,30 +563,14 @@ async def asend_message( # 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( + await _record_send_message_billing( kwargs=kwargs, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, + request=request, + response_dict=a2a_response.root.model_dump(mode="json", exclude_none=True), + litellm_params=litellm_params, + agent_id=agent_id, ) - # 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 @@ -646,6 +655,7 @@ async def asend_message_streaming( metadata: dict[str, object] | None = None, proxy_server_request: dict[str, object] | None = None, agent_extra_headers: dict[str, str] | None = None, + agent_name: str | None = None, **kwargs: object, ) -> AsyncIterator["SendStreamingMessageResponse | Mapping[str, object]"]: """ @@ -661,6 +671,7 @@ async def asend_message_streaming( 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 + agent_name: Agent name logged as a2a_agent/ for completion bridge agents Yields: SendStreamingMessageResponse chunks from the agent @@ -685,6 +696,7 @@ async def asend_message_streaming( print(chunk) ``` """ + called_at: Final = datetime.datetime.now() litellm_params = litellm_params or {} custom_llm_provider: Final = litellm_params.get("custom_llm_provider") @@ -705,20 +717,47 @@ async def asend_message_streaming( request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params) ) - async for chunk in A2ACompletionBridgeHandler.handle_streaming( + bridge_stream: Final = 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 + ) + caller_logging_obj: Final = kwargs.get("litellm_logging_obj") + bridge_logging_obj: Final = ( + caller_logging_obj + if isinstance(caller_logging_obj, Logging) + else _build_streaming_logging_obj( + request=request, + agent_name=agent_name or "unknown", + agent_id=agent_id, + litellm_params=litellm_params, + metadata=metadata, + proxy_server_request=proxy_server_request, + start_time=called_at, + ) + ) + _set_litellm_params_on_logging_obj(kwargs=kwargs, litellm_params=litellm_params) + _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id) + + bridge_iterator: Final = A2AStreamingIterator( + stream=bridge_stream, + request=request, + logging_obj=bridge_logging_obj, + agent_name=agent_name or "unknown", + ) + try: + async for chunk in bridge_iterator: + yield chunk + except (GeneratorExit, asyncio.CancelledError): + await bridge_iterator.abill_on_disconnect() + raise return if request is None: raise ValueError("request is required") - called_at: Final = datetime.datetime.now() _raw_logging_obj: Final = kwargs.get("litellm_logging_obj") logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None @@ -742,12 +781,12 @@ async def asend_message_streaming( assert a2a_client is not None - agent_name: Final = _get_a2a_model_info(a2a_client, kwargs) + card_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_name=card_agent_name, agent_id=agent_id, litellm_params=litellm_params, metadata=metadata, @@ -755,7 +794,7 @@ async def asend_message_streaming( start_time=called_at, ) - verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name) + verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, card_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 @@ -766,18 +805,23 @@ async def asend_message_streaming( agent_card=agent_card, card_url=card_url, api_base=api_base, - agent_name=agent_name, + agent_name=card_agent_name, ) _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id) - async for chunk in A2AStreamingIterator( + url_agent_iterator: Final = A2AStreamingIterator( stream=stream, request=request, logging_obj=logging_obj, - agent_name=agent_name, - ): - yield chunk + agent_name=card_agent_name, + ) + try: + async for chunk in url_agent_iterator: + yield chunk + except (GeneratorExit, asyncio.CancelledError): + await url_agent_iterator.abill_on_disconnect() + raise async def create_a2a_client( diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 35fe6e76f4d..bc65bac1a7c 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -515,6 +515,8 @@ class PydanticAITransformation: "contextId": context_id, "kind": "artifact-update", "taskId": task_id, + "append": i > 0, + "lastChunk": is_last_chunk, "artifact": { "artifactId": artifact_id, "parts": [ diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py index c2d581321cd..75262bde392 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py @@ -206,6 +206,8 @@ class WatsonxOrchestrateTransformation: "contextId": context_id, "kind": "artifact-update", "taskId": task_id, + "append": i > 0, + "lastChunk": is_last, "artifact": { "artifactId": artifact_id, "parts": [{"kind": "text", "text": chunk_text}], diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index 328681b8dbe..7b2b6884cae 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -3,16 +3,18 @@ A2A Streaming Iterator with token tracking and logging support. """ import asyncio -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from datetime import datetime from typing import TYPE_CHECKING, Final +import anyio from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.a2a_protocol.cost_calculator import A2ACostCalculator from litellm.a2a_protocol.utils import A2ARequestUtils +from litellm.constants import A2A_COMPLETION_BRIDGE_WORKING_MESSAGE from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -23,6 +25,23 @@ if TYPE_CHECKING: _OPAQUE_PAYLOAD: Final = TypeAdapter(object) +def _chunk_payload(chunk: "SendStreamingMessageResponse | Mapping[str, object]") -> Mapping[str, object]: + if isinstance(chunk, Mapping): + return chunk + return chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} + + +def _artifact_update(payload: Mapping[str, object]) -> tuple[str, bool] | None: + result: Final = as_str_mapping(payload.get("result")) + if result is None or result.get("kind") != "artifact-update": + return None + artifact: Final = as_str_mapping(result.get("artifact")) + artifact_id: Final = artifact.get("artifactId") if artifact is not None else None + if not isinstance(artifact_id, str): + return None + return artifact_id, result.get("append") is True + + class A2AStreamingIterator: """ Async iterator for A2A streaming responses with token tracking. @@ -32,7 +51,7 @@ class A2AStreamingIterator: def __init__( self, - stream: AsyncIterator["SendStreamingMessageResponse"], + stream: AsyncIterator["SendStreamingMessageResponse | Mapping[str, object]"], request: "SendStreamingMessageRequest", logging_obj: LiteLLMLoggingObj, agent_name: str = "unknown", @@ -44,14 +63,16 @@ class A2AStreamingIterator: self.start_time = logging_obj.start_time # Collect chunks for token counting - self.chunks: list[SendStreamingMessageResponse] = [] - self.collected_text_parts: list[str] = [] - self.final_chunk: SendStreamingMessageResponse | None = None + self.chunks: list[SendStreamingMessageResponse | Mapping[str, object]] = [] + self.latest_message_text = "" + self.artifact_texts: Mapping[str, str] = {} + self.saw_error = False + self.final_chunk: SendStreamingMessageResponse | Mapping[str, object] | None = None def __aiter__(self): return self - async def __anext__(self) -> "SendStreamingMessageResponse": + async def __anext__(self) -> "SendStreamingMessageResponse | Mapping[str, object]": try: chunk: Final = await self.stream.__anext__() if self.logging_obj.completion_start_time is None: @@ -59,38 +80,54 @@ class A2AStreamingIterator: # Store chunk self.chunks.append(chunk) + payload: Final = _chunk_payload(chunk) + self.saw_error = self.saw_error or payload.get("error") is not None # Extract text from chunk for token counting - self._collect_text_from_chunk(chunk) + self._collect_text_from_chunk(payload) # Check if this is the final chunk (completed status) - if self._is_completed_chunk(chunk): + if self._is_completed_chunk(payload): self.final_chunk = chunk return chunk except StopAsyncIteration: - # Stream ended - handle logging - if self.final_chunk is None and self.chunks: - self.final_chunk = self.chunks[-1] await self._handle_stream_complete() raise - def _collect_text_from_chunk(self, chunk: "SendStreamingMessageResponse") -> None: - """Extract text from a streaming chunk and add to collected parts.""" + async def abill_on_disconnect(self) -> None: + if not self.chunks: + return + with anyio.CancelScope(shield=True): + await self._handle_stream_complete() + + def _collect_text_from_chunk(self, payload: Mapping[str, object]) -> None: + """Record the reply text a streaming chunk carries.""" try: - chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} - text: Final = A2ARequestUtils.extract_text_from_response(chunk_dict) - if text: - self.collected_text_parts.append(text) + text: Final = A2ARequestUtils.extract_text_from_stream_event(payload) + if not text: + return + artifact_update: Final = _artifact_update(payload) + if artifact_update is None and text == A2A_COMPLETION_BRIDGE_WORKING_MESSAGE: + return + if artifact_update is None: + self.latest_message_text = text + return + artifact_id, append = artifact_update + artifact_text: Final = self.artifact_texts.get(artifact_id, "") + text if append else text + self.artifact_texts = {**self.artifact_texts, artifact_id: artifact_text} except Exception: verbose_logger.debug("Failed to extract text from A2A streaming chunk") - def _is_completed_chunk(self, chunk: "SendStreamingMessageResponse") -> bool: + def _reply_text(self) -> str: + artifact_text: Final = " ".join(text for text in self.artifact_texts.values() if text) + return artifact_text or self.latest_message_text + + def _is_completed_chunk(self, payload: Mapping[str, object]) -> bool: """Check if chunk indicates stream completion.""" try: - chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} - result: Final = chunk_dict.get("result", {}) + result: Final = payload.get("result", {}) if isinstance(result, dict): status: Final = result.get("status", {}) if isinstance(status, dict): @@ -101,6 +138,8 @@ class A2AStreamingIterator: async def _handle_stream_complete(self) -> None: """Handle logging and token counting when stream completes.""" + if self.final_chunk is None and self.chunks: + self.final_chunk = self.chunks[-1] try: end_time: Final = datetime.now() @@ -109,9 +148,7 @@ class A2AStreamingIterator: input_text: Final = A2ARequestUtils.extract_text_from_message(input_message) prompt_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(input_text) - # Use the last (most complete) text from chunks - output_text: Final = self.collected_text_parts[-1] if self.collected_text_parts else "" - completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(output_text) + completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(self._reply_text()) total_tokens: Final = prompt_tokens + completion_tokens @@ -128,7 +165,7 @@ class A2AStreamingIterator: self.logging_obj.model_call_details["stream"] = False # Calculate cost using A2ACostCalculator - response_cost: Final = A2ACostCalculator.calculate_a2a_cost(self.logging_obj) + response_cost: Final = 0.0 if self.saw_error else A2ACostCalculator.calculate_a2a_cost(self.logging_obj) self.logging_obj.model_call_details["response_cost"] = response_cost # Build result for logging @@ -167,9 +204,7 @@ class A2AStreamingIterator: # Add final chunk result if available if self.final_chunk: try: - chunk_dict: Final = as_str_mapping( - _OPAQUE_PAYLOAD.validate_python(self.final_chunk.model_dump(mode="json", exclude_none=True)) - ) + chunk_dict: Final = as_str_mapping(_OPAQUE_PAYLOAD.validate_python(_chunk_payload(self.final_chunk))) if chunk_dict is not None: result["result"] = chunk_dict.get("result", {}) except Exception: diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index 17d1b86a609..a9ec0de116e 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -6,13 +6,17 @@ import hashlib from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Protocol +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import has_attributes +from litellm.litellm_core_utils.core_helpers import as_str_mapping, has_attributes if TYPE_CHECKING: from a2a.types import SendMessageRequest, SendStreamingMessageRequest +_OBJECT_ITEMS: Final = TypeAdapter(tuple[object, ...]) + class _SupportsModelDump(Protocol): def model_dump(self, *, mode: str) -> object: ... @@ -34,6 +38,9 @@ class A2ARequestUtils: """ if message is None: return "" + # a2a-sdk Part RootModels keep kind/text on part.root, so read pydantic messages as dicts + if has_attributes(message, _SupportsModelDump, "model_dump"): + return A2ARequestUtils.extract_text_from_message(message.model_dump(mode="json")) # Handle both dict and object access if isinstance(message, dict): @@ -63,8 +70,8 @@ class A2ARequestUtils: Returns: Text from response message parts """ - result: Final = response_dict.get("result", {}) - if not isinstance(result, dict): + result: Final = as_str_mapping(response_dict.get("result")) + if result is None: return "" # Direct message format (A2A spec): detect by explicit kind tag only. @@ -72,10 +79,36 @@ class A2ARequestUtils: # type that happens to include a "parts" field. if result.get("kind") == "message": return A2ARequestUtils.extract_text_from_message(result) + if result.get("kind") == "task": + return A2ARequestUtils.extract_text_from_task(result) message: Final = result.get("message", {}) return A2ARequestUtils.extract_text_from_message(message) + @staticmethod + def extract_text_from_task(task: Mapping[str, object]) -> str: + artifacts: Final = task.get("artifacts") + artifact_items: Final = _OBJECT_ITEMS.validate_python(artifacts) if isinstance(artifacts, list) else () + artifact_texts: Final = tuple( + A2ARequestUtils.extract_text_from_message(artifact) for artifact in artifact_items + ) + if any(artifact_texts): + return " ".join(text for text in artifact_texts if text) + status: Final = as_str_mapping(task.get("status")) + return A2ARequestUtils.extract_text_from_message(status.get("message")) if status is not None else "" + + @staticmethod + def extract_text_from_stream_event(event: Mapping[str, object]) -> str: + result: Final = as_str_mapping(event.get("result")) + if result is None: + return "" + if result.get("kind") == "artifact-update": + return A2ARequestUtils.extract_text_from_message(result.get("artifact")) + if result.get("kind") == "status-update": + status: Final = as_str_mapping(result.get("status")) + return A2ARequestUtils.extract_text_from_message(status.get("message")) if status is not None else "" + return A2ARequestUtils.extract_text_from_response(event) + @staticmethod def get_input_message_from_request( request: "SendMessageRequest | SendStreamingMessageRequest", @@ -128,17 +161,8 @@ class A2ARequestUtils: Returns: Tuple of (prompt_tokens, completion_tokens, total_tokens) """ - # Count input tokens. Dump the message to a dict first so extraction hits - # the dict branch — request-side parts are a2a-sdk Part RootModels whose - # kind/text live on part.root, which the object branch cannot read. This - # mirrors how the response side already works (it operates on model_dump). input_message: Final = A2ARequestUtils.get_input_message_from_request(request) - input_text: Final = A2ARequestUtils.extract_text_from_message( - input_message.model_dump(mode="json") - if has_attributes(input_message, _SupportsModelDump, "model_dump") - else input_message - ) - prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text) + prompt_tokens: Final = A2ARequestUtils.count_tokens(A2ARequestUtils.extract_text_from_message(input_message)) # Count output tokens output_text: Final = A2ARequestUtils.extract_text_from_response(response_dict) diff --git a/litellm/constants.py b/litellm/constants.py index fc767a1dad6..c17e776f54b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -902,6 +902,7 @@ PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset( SUBSCRIPTION_BACKED_PROVIDERS: Final = frozenset({"chatgpt"}) A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: Final = frozenset({"databricks_agent"}) +A2A_COMPLETION_BRIDGE_WORKING_MESSAGE: Final = "Processing request..." LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [ "openai", diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 219ce906340..5e1d861e7e9 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -431,6 +431,7 @@ async def _handle_stream_message( proxy_server_request: dict[str, object] | None = None, *, agent_extra_headers: dict[str, str] | None = None, + agent_name: str | None = None, user_api_key_dict: UserAPIKeyAuth | None = None, request_data: dict[str, object] | None = None, proxy_logging_obj: ProxyLogging | None = None, @@ -504,6 +505,7 @@ async def _handle_stream_message( metadata=metadata, proxy_server_request=proxy_server_request, agent_extra_headers=agent_extra_headers, + agent_name=agent_name, ) if ( @@ -932,6 +934,7 @@ async def invoke_agent_a2a( metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), agent_extra_headers=agent_extra_headers, + agent_name=str(agent_name), user_api_key_dict=user_api_key_auth_for_call, request_data=data, proxy_logging_obj=proxy_logging_obj, diff --git a/tests/integration/spend/test_agent_spend.py b/tests/integration/spend/test_agent_spend.py index 58fcd4e9de2..fca848fbb43 100644 --- a/tests/integration/spend/test_agent_spend.py +++ b/tests/integration/spend/test_agent_spend.py @@ -269,9 +269,4 @@ def test_a2a_send_bills_input_and_output_tokens_at_the_agent_per_token_prices(ga def test_a2a_stream_bills_input_and_output_tokens_at_the_agent_per_token_prices(gateway: Gateway) -> None: - pytest.skip( - "BUG: message/stream to an agent priced with input_cost_per_token/output_cost_per_token logs spend 0.0 " - "with prompt_tokens 0 and completion_tokens 0, because A2AStreamingIterator reads text from the request " - "Part RootModels without dumping them and ignores artifact-update text, while message/send bills 1.125" - ) _assert_a2a_calls_are_billed(gateway, _PER_TOKEN, _PER_TOKEN_COST, (("message/stream", "/a2a/{agent}"),)) diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py index 38b896723f7..2eaa04a767e 100644 --- a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py +++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_config.py @@ -148,6 +148,8 @@ async def test_handle_streaming_runs_the_agent_and_ignores_unrelated_keywords(re "kind": "artifact-update", "taskId": ANY, "artifact": {"artifactId": ANY, "parts": [{"kind": "text", "text": "wxo says hi"}]}, + "append": False, + "lastChunk": True, } assert chunks[3]["result"] == { "contextId": ANY, diff --git a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py index 6d685674984..0d4103ae979 100644 --- a/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py +++ b/tests/unit/a2a_protocol/test_a2a_streaming_iterator.py @@ -5,7 +5,9 @@ async_success_handler (cross-thread pydantic mutation segfaults pydantic-core). """ import asyncio +import importlib.util import time +from collections.abc import AsyncIterator, Callable, Mapping from datetime import datetime from typing import Final from types import SimpleNamespace @@ -15,6 +17,7 @@ import pytest import litellm from litellm.a2a_protocol import streaming_iterator as a2a_streaming_iterator_module from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator +from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -176,7 +179,9 @@ async def test_stream_logs_its_first_chunk_time_measured_from_the_call_start(): iterator: Final = A2AStreamingIterator( stream=_stream(), - request=SimpleNamespace(params=SimpleNamespace(message={"role": "user", "parts": [{"kind": "text", "text": "hi"}]})), + request=SimpleNamespace( + params=SimpleNamespace(message={"role": "user", "parts": [{"kind": "text", "text": "hi"}]}) + ), logging_obj=logging_obj, agent_name="test-agent", ) @@ -189,3 +194,133 @@ async def test_stream_logs_its_first_chunk_time_measured_from_the_call_start(): assert first_chunk_at is not None assert recorder.logged == [(call_start, first_chunk_at)] + + +_REPLY: Final = ( + "Paris is the capital of France. It sits on the Seine, holds about two million people, " + "and has been the seat of the French government for most of the last thousand years." +) + + +def _artifact_event(text: str, append: bool | None, artifact_id: str = "a1") -> Mapping[str, object]: + result: Final = { + "kind": "artifact-update", + "taskId": "t1", + "contextId": "c1", + "artifact": {"artifactId": artifact_id, "parts": [{"kind": "text", "text": text}]}, + } + return {"jsonrpc": "2.0", "id": "r1", "result": result if append is None else {**result, "append": append}} + + +def _status_event(state: str, text: str | None = None) -> Mapping[str, object]: + message: Final = {"kind": "message", "role": "agent", "messageId": "m1", "parts": [{"kind": "text", "text": text}]} + status: Final = {"state": state} if text is None else {"state": state, "message": message} + return { + "jsonrpc": "2.0", + "id": "r1", + "result": { + "kind": "status-update", + "taskId": "t1", + "contextId": "c1", + "final": state == "completed", + "status": status, + }, + } + + +_FIRST_ARTIFACT: Final = "Paris is the capital of France and sits on the Seine." +_SECOND_ARTIFACT: Final = "It has been the seat of the French government for most of the last thousand years." + + +async def _snapshot_stream() -> AsyncIterator[Mapping[str, object]]: + for end in range(40, len(_REPLY) + 40, 40): + yield _artifact_event(_REPLY[:end], append=None) + + +async def _watsonx_stream() -> AsyncIterator[Mapping[str, object]]: + from litellm.a2a_protocol.providers.watsonx_orchestrate.transformation import ( + WatsonxOrchestrateTransformation, + ) + + async for event in WatsonxOrchestrateTransformation.fake_streaming_from_text(_REPLY, "r1", delay_ms=0): + yield event + + +async def _two_artifact_stream() -> AsyncIterator[Mapping[str, object]]: + yield _artifact_event(_FIRST_ARTIFACT, append=None, artifact_id="a1") + yield _artifact_event(_SECOND_ARTIFACT, append=None, artifact_id="a2") + yield _status_event("completed") + + +async def _status_message_stream() -> AsyncIterator[Mapping[str, object]]: + yield _status_event("working", "Looking that up") + yield _status_event("working", _REPLY) + yield _status_event("completed") + + +async def _artifact_without_id_stream() -> AsyncIterator[Mapping[str, object]]: + yield { + "jsonrpc": "2.0", + "id": "r1", + "result": { + "kind": "artifact-update", + "taskId": "t1", + "contextId": "c1", + "artifact": {"parts": [{"kind": "text", "text": _REPLY}]}, + }, + } + yield _status_event("completed") + + +async def _url_agent_sdk_stream() -> AsyncIterator[object]: + from a2a.compat.v0_3.types import SendStreamingMessageResponse + + for event in ( + _status_event("working", "Looking that up"), + _artifact_event(_REPLY, append=False), + _status_event("completed"), + ): + yield SendStreamingMessageResponse.model_validate(event) + + +@pytest.mark.parametrize( + ("agent_stream", "reply"), + ( + pytest.param(_snapshot_stream, _REPLY, id="replacing_snapshots"), + pytest.param(_watsonx_stream, _REPLY, id="watsonx_chunks"), + pytest.param(_two_artifact_stream, f"{_FIRST_ARTIFACT} {_SECOND_ARTIFACT}", id="two_artifacts"), + pytest.param(_status_message_stream, _REPLY, id="status_message_reply"), + pytest.param(_artifact_without_id_stream, _REPLY, id="artifact_without_id"), + pytest.param( + _url_agent_sdk_stream, + _REPLY, + id="url_agent_sdk_events", + marks=pytest.mark.skipif(importlib.util.find_spec("a2a") is None, reason="needs the optional a2a-sdk"), + ), + ), +) +@pytest.mark.asyncio +async def test_stream_counts_the_whole_reply_once( + agent_stream: Callable[[], AsyncIterator[object]], reply: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + logging_obj: Final = LitellmLogging( + model="a2a_agent/trip-planner", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="asend_message_streaming", + start_time=time.time(), + litellm_call_id="artifact-count", + function_id="artifact-count", + ) + iterator: Final = A2AStreamingIterator( + stream=agent_stream(), + request=SimpleNamespace(params=SimpleNamespace(message={"role": "user", "parts": []})), + logging_obj=logging_obj, + agent_name="trip-planner", + ) + + _ = [chunk async for chunk in iterator] + + assert logging_obj.model_call_details["usage"].completion_tokens == A2ARequestUtils.count_tokens(reply) diff --git a/tests/unit/a2a_protocol/test_main.py b/tests/unit/a2a_protocol/test_main.py index 75fd9eda75b..c0cabee1127 100644 --- a/tests/unit/a2a_protocol/test_main.py +++ b/tests/unit/a2a_protocol/test_main.py @@ -1,12 +1,16 @@ """Tests for litellm/a2a_protocol/main.py non-streaming send behavior.""" import asyncio +import contextlib import datetime import json -from typing import Final +from collections.abc import AsyncGenerator, Callable, Coroutine, Iterator, Mapping +from typing import Final, TypeAlias import httpx import pytest +import respx +from pydantic import BaseModel, ConfigDict pytest.importorskip("a2a.compat.v0_3.conversions") @@ -27,6 +31,7 @@ from litellm.a2a_protocol.main import ( create_a2a_client, ) from litellm.a2a_protocol.exceptions import A2AError +from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.caching.llm_caching_handler import LLMClientCache from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT from litellm.llms.custom_httpx.http_handler import ( @@ -798,6 +803,403 @@ def test_streaming_logging_obj_keeps_agent_credentials_out_of_logging_params(): assert logging_obj.model_call_details["litellm_params"] == expected +class _LoggedSpend(BaseModel): + model_config = ConfigDict(frozen=True) + + call_type: str + model: str + prompt_tokens: int + completion_tokens: int + response_cost: float + + +class _A2ASpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: Final = asyncio.Event() + self.spend: _LoggedSpend | None = None + self.agent_id: object = None + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + spend: Final = _LoggedSpend.model_validate(kwargs["standard_logging_object"]) + if spend.call_type not in ("asend_message", "asend_message_streaming"): + return + self.spend = spend + self.agent_id = kwargs.get("agent_id") + self.logged.set() + + async def logged_spend(self) -> _LoggedSpend: + await asyncio.wait_for(self.logged.wait(), timeout=10) + assert self.spend is not None + return self.spend + + +@pytest.fixture +def a2a_spend_recorder(monkeypatch: pytest.MonkeyPatch) -> _A2ASpendRecorder: + recorder: Final = _A2ASpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + return recorder + + +@pytest.fixture +def respx_reachable_httpx(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +_AgentPrices: TypeAlias = Mapping[str, float] +_ExpectedCost: TypeAlias = Callable[[int, int], float] + +_BRIDGE_PROMPT: Final = "What is the capital of France?" +_BRIDGE_REPLY: Final = "Paris is the capital of France." +_LONG_AGENT_REPLY: Final = ( + "Paris is the capital of France. It sits on the Seine, holds about two million people, " + "and has been the seat of the French government for most of the last thousand years." +) +_PER_TOKEN_PRICES: Final[_AgentPrices] = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +_AGENTCORE_ARN: Final = "arn:aws:bedrock-agentcore:us-west-2:123456789:runtime/trip_planner" +_AGENTCORE_PARAMS: Final[Mapping[str, object]] = { + "custom_llm_provider": "bedrock", + "model": f"bedrock/agentcore/{_AGENTCORE_ARN}", + "api_key": "test-jwt-token", +} +_JSON_RPC_ERROR: Final[Mapping[str, object]] = { + "jsonrpc": "2.0", + "id": "agentcore-call", + "error": {"code": -32603, "message": "agent crashed"}, +} + + +def _per_token_cost(prompt_tokens: int, completion_tokens: int) -> float: + return prompt_tokens * 0.001 + completion_tokens * 0.002 + + +_BRIDGE_AGENT_PRICES: Final = ( + pytest.param({"cost_per_query": 0.25}, lambda prompt_tokens, completion_tokens: 0.25, id="cost_per_query"), + pytest.param(_PER_TOKEN_PRICES, _per_token_cost, id="per_token"), +) + + +def _bridge_agent_params(prices: _AgentPrices) -> dict[str, object]: + return { + "custom_llm_provider": "langgraph", + "model": "agent", + "mock_response": _BRIDGE_REPLY, + **prices, + } + + +def _bridge_message_params() -> MessageSendParams: + return MessageSendParams( + message={ + "messageId": "m1", + "role": "user", + "parts": [{"kind": "text", "text": _BRIDGE_PROMPT}], + } + ) + + +def _assert_billed_at_agent_prices( + spend: _LoggedSpend, recorder: _A2ASpendRecorder, expected_cost: _ExpectedCost, reply: str +) -> None: + assert spend.prompt_tokens == A2ARequestUtils.count_tokens(_BRIDGE_PROMPT) + assert spend.completion_tokens == A2ARequestUtils.count_tokens(reply) + assert spend.response_cost == pytest.approx(expected_cost(spend.prompt_tokens, spend.completion_tokens)) + assert recorder.agent_id == "agent-1" + + +async def _drain_bridge_stream(litellm_params: Mapping[str, object]) -> None: + from litellm.a2a_protocol.main import asend_message_streaming + + _ = [ + chunk + async for chunk in asend_message_streaming( + request=SendStreamingMessageRequest(id="bridge-stream", params=_bridge_message_params()), + api_base="http://127.0.0.1:9", + litellm_params=dict(litellm_params), + agent_id="agent-1", + agent_name="trip-planner", + ) + ] + + +def _sse(*events: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(f"data: {json.dumps(event)}\n\n" for event in events), + ) + + +def _agentcore_appended_chunks(text: str, chunk_size: int) -> tuple[Mapping[str, object], ...]: + return tuple( + { + "jsonrpc": "2.0", + "id": "agentcore-call", + "result": { + "kind": "artifact-update", + "taskId": "t1", + "contextId": "c1", + "append": start > 0, + "artifact": {"artifactId": "a1", "parts": [{"kind": "text", "text": text[start : start + chunk_size]}]}, + }, + } + for start in range(0, len(text), chunk_size) + ) + + +@pytest.mark.parametrize(("prices", "expected_cost"), _BRIDGE_AGENT_PRICES) +@pytest.mark.asyncio +async def test_completion_bridge_send_bills_the_agents_prices( + prices: _AgentPrices, expected_cost: _ExpectedCost, a2a_spend_recorder: _A2ASpendRecorder +) -> None: + await asend_message( + request=SendMessageRequest(id="bridge-send", params=_bridge_message_params()), + api_base="http://127.0.0.1:9", + litellm_params=_bridge_agent_params(prices), + agent_id="agent-1", + ) + + spend: Final = await a2a_spend_recorder.logged_spend() + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, expected_cost, _BRIDGE_REPLY) + + +@pytest.mark.parametrize(("prices", "expected_cost"), _BRIDGE_AGENT_PRICES) +@pytest.mark.asyncio +async def test_completion_bridge_stream_bills_the_agents_prices( + prices: _AgentPrices, expected_cost: _ExpectedCost, a2a_spend_recorder: _A2ASpendRecorder +) -> None: + await _drain_bridge_stream(_bridge_agent_params(prices)) + + spend: Final = await a2a_spend_recorder.logged_spend() + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, expected_cost, _BRIDGE_REPLY) + assert spend.model == "a2a_agent/trip-planner" + + +@pytest.mark.asyncio +async def test_agentcore_send_bills_a_task_reply_at_the_agents_token_prices( + a2a_spend_recorder: _A2ASpendRecorder, respx_reachable_httpx: None +) -> None: + task_reply: Final = { + "jsonrpc": "2.0", + "id": "bridge-send", + "result": { + "kind": "task", + "id": "t1", + "contextId": "c1", + "status": {"state": "completed"}, + "artifacts": [{"artifactId": "a1", "parts": [{"kind": "text", "text": _LONG_AGENT_REPLY}]}], + }, + } + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock(return_value=httpx.Response(200, json=task_reply)) + await asend_message( + request=SendMessageRequest(id="bridge-send", params=_bridge_message_params()), + litellm_params={**_AGENTCORE_PARAMS, **_PER_TOKEN_PRICES}, + agent_id="agent-1", + ) + + spend: Final = await a2a_spend_recorder.logged_spend() + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, _per_token_cost, _LONG_AGENT_REPLY) + + +@pytest.mark.parametrize( + "extra_fields", + (pytest.param({}, id="result_only"), pytest.param({"error": None}, id="null_error_field")), +) +@pytest.mark.asyncio +async def test_agentcore_stream_bills_every_appended_reply_chunk( + extra_fields: Mapping[str, object], a2a_spend_recorder: _A2ASpendRecorder, respx_reachable_httpx: None +) -> None: + events: Final = tuple({**event, **extra_fields} for event in _agentcore_appended_chunks(_LONG_AGENT_REPLY, 40)) + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock(return_value=_sse(*events)) + await _drain_bridge_stream({**_AGENTCORE_PARAMS, **_PER_TOKEN_PRICES}) + + spend: Final = await a2a_spend_recorder.logged_spend() + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, _per_token_cost, _LONG_AGENT_REPLY) + + +_PYDANTIC_COMPLETED_TASK: Final[Mapping[str, object]] = { + "jsonrpc": "2.0", + "id": "bridge-stream", + "result": { + "kind": "task", + "id": "t1", + "contextId": "c1", + "status": {"state": "completed"}, + "artifacts": [{"artifactId": "a1", "parts": [{"kind": "text", "text": _LONG_AGENT_REPLY}]}], + }, +} + + +@pytest.mark.asyncio +async def test_pydantic_ai_stream_bills_the_whole_chunked_reply( + a2a_spend_recorder: _A2ASpendRecorder, respx_reachable_httpx: None +) -> None: + with respx.mock(assert_all_called=True) as router: + router.post("http://127.0.0.1:9/").mock(return_value=httpx.Response(200, json=_PYDANTIC_COMPLETED_TASK)) + await _drain_bridge_stream({"custom_llm_provider": "pydantic_ai_agents", **_PER_TOKEN_PRICES}) + + spend: Final = await a2a_spend_recorder.logged_spend() + assert len(_LONG_AGENT_REPLY) > 50 + _assert_billed_at_agent_prices(spend, a2a_spend_recorder, _per_token_cost, _LONG_AGENT_REPLY) + + +async def _close_after_first_event(stream: AsyncGenerator[object, None]) -> None: + _ = await stream.__anext__() + await stream.aclose() + + +async def _cancel_reader_after_first_event(stream: AsyncGenerator[object, None]) -> None: + first_event: Final = asyncio.Event() + + async def _read() -> None: + async for _ in stream: + first_event.set() + + reader: Final = asyncio.create_task(_read()) + await first_event.wait() + reader.cancel() + with contextlib.suppress(asyncio.CancelledError): + await reader + + +def _pydantic_bridge_agent_stream() -> AsyncGenerator[object, None]: + from litellm.a2a_protocol.main import asend_message_streaming + + return asend_message_streaming( + request=SendStreamingMessageRequest(id="bridge-stream", params=_bridge_message_params()), + api_base="http://127.0.0.1:9", + litellm_params={"custom_llm_provider": "pydantic_ai_agents", "cost_per_query": 0.25}, + agent_id="agent-1", + agent_name="trip-planner", + ) + + +def _url_agent_stream() -> AsyncGenerator[object, None]: + from litellm.a2a_protocol.main import asend_message_streaming + + return asend_message_streaming( + request=SendStreamingMessageRequest(id="reply", params=_bridge_message_params()), + api_base="http://127.0.0.1:9", + litellm_params={"cost_per_query": 0.25}, + agent_id="agent-1", + ) + + +@pytest.mark.parametrize( + ("open_stream", "billed_model"), + ( + pytest.param(_pydantic_bridge_agent_stream, "a2a_agent/trip-planner", id="bridge_agent"), + pytest.param(_url_agent_stream, f"a2a_agent/{_AGENT_CARD['name']}", id="url_agent"), + ), +) +@pytest.mark.parametrize( + "drop_stream", + ( + pytest.param(_close_after_first_event, id="proxy_closes_the_stream"), + pytest.param(_cancel_reader_after_first_event, id="reader_cancelled_mid_stream"), + ), +) +@pytest.mark.asyncio +async def test_stream_the_client_drops_still_bills_the_agents_price( + open_stream: Callable[[], AsyncGenerator[object, None]], + billed_model: str, + drop_stream: Callable[[AsyncGenerator[object, None]], Coroutine[object, object, None]], + a2a_spend_recorder: _A2ASpendRecorder, + isolated_client_cache: LLMClientCache, + respx_reachable_httpx: None, +) -> None: + _ = await _seed_shared_a2a_client(rpc_stream_events=_LANGGRAPH_STREAM_EVENTS) + with respx.mock(assert_all_called=False) as router: + router.post("http://127.0.0.1:9/").mock(return_value=httpx.Response(200, json=_PYDANTIC_COMPLETED_TASK)) + await drop_stream(open_stream()) + + spend: Final = await a2a_spend_recorder.logged_spend() + assert spend.response_cost == pytest.approx(0.25) + assert spend.model == billed_model + + +class _AgentThatNeverAnswers(_RequestRecorder): + def __init__(self) -> None: + super().__init__() + self.message_sent: Final = asyncio.Event() + + def __call__(self, request: httpx.Request) -> httpx.Response | Coroutine[object, object, httpx.Response]: + if request.method == "GET": + return super().__call__(request) + return self._never_answer() + + async def _never_answer(self) -> httpx.Response: + self.message_sent.set() + reply: Final[asyncio.Future[httpx.Response]] = asyncio.get_running_loop().create_future() + return await reply + + +@pytest.mark.asyncio +async def test_stream_dropped_before_the_agent_answers_bills_nothing( + a2a_spend_recorder: _A2ASpendRecorder, isolated_client_cache: LLMClientCache +) -> None: + agent: Final = _AgentThatNeverAnswers() + _ = await _seed_shared_a2a_client(recorder=agent) + stream: Final = _url_agent_stream() + reader: Final = asyncio.create_task(anext(stream)) + await asyncio.wait_for(agent.message_sent.wait(), timeout=10) + reader.cancel() + with contextlib.suppress(asyncio.CancelledError): + await reader + + with pytest.raises(TimeoutError): + await asyncio.wait_for(a2a_spend_recorder.logged.wait(), timeout=2) + + +@pytest.mark.asyncio +async def test_completion_bridge_stream_with_an_empty_reply_bills_no_reply_tokens( + a2a_spend_recorder: _A2ASpendRecorder, +) -> None: + await _drain_bridge_stream({**_bridge_agent_params(_PER_TOKEN_PRICES), "mock_response": ""}) + + spend: Final = await a2a_spend_recorder.logged_spend() + assert spend.completion_tokens == 0 + assert spend.response_cost == pytest.approx(_per_token_cost(spend.prompt_tokens, 0)) + + +@pytest.mark.parametrize( + "upstream_reply", + ( + pytest.param(httpx.Response(200, json=_JSON_RPC_ERROR), id="send"), + pytest.param(_sse(_JSON_RPC_ERROR), id="stream"), + ), +) +@pytest.mark.asyncio +async def test_agentcore_error_reply_is_not_billed_at_the_agents_price( + upstream_reply: httpx.Response, a2a_spend_recorder: _A2ASpendRecorder, respx_reachable_httpx: None +) -> None: + agent_params: Final = {**_AGENTCORE_PARAMS, "cost_per_query": 0.25} + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock(return_value=upstream_reply) + if upstream_reply.headers["content-type"] == "text/event-stream": + await _drain_bridge_stream(agent_params) + else: + response = await asend_message( + request=SendMessageRequest(id="agentcore-call", params=_bridge_message_params()), + litellm_params=dict(agent_params), + agent_id="agent-1", + ) + assert response.model_dump(mode="json", exclude_none=True)["error"]["message"] == "agent crashed" + + spend: Final = await a2a_spend_recorder.logged_spend() + assert spend.response_cost == 0.0 + + class MockA2AResponse: def __init__(self, text: str): self._payload = { diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 7abe6ad7799..d6d6774d414 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -5,6 +5,7 @@ Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request """ import asyncio +import importlib.util import json import socket import sys @@ -17,8 +18,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from pydantic import BaseModel, ConfigDict import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.agents import AgentCaller @@ -44,6 +47,7 @@ def _sse_upstream_response(*lines: str, first_line_delay: float = 0.0) -> httpx. class CapturedAgentCall: request_id: object agent_extra_headers: dict[str, str] | None + agent_name: object @pytest.mark.asyncio @@ -627,7 +631,11 @@ async def _invoke_message_method( kwargs: Final = downstream.call_args.kwargs request_id: Final = kwargs["request"].__dict__["id"] if is_send else kwargs["request_id"] - return CapturedAgentCall(request_id=request_id, agent_extra_headers=kwargs.get("agent_extra_headers")) + return CapturedAgentCall( + request_id=request_id, + agent_extra_headers=kwargs.get("agent_extra_headers"), + agent_name=kwargs.get("agent_name"), + ) @pytest.mark.asyncio @@ -1684,6 +1692,76 @@ def test_build_message_send_params_proto_fallback_ignores_unknown_fields(): assert result.message.role.value == "user" +@pytest.mark.asyncio +async def test_message_stream_names_the_agent_by_its_card_name(): + agent: Final = _make_agent_mock() + agent.agent_card_params = {"url": "http://backend-agent:10001", "name": "trip-planner"} + + captured: Final = await _invoke_message_method( + "message/stream", + _make_request_mock("message/stream", _HELLO_MESSAGE_PARAMS), + UserAPIKeyAuth(api_key="sk-test", user_id="u1"), + agent=agent, + ) + + assert captured.agent_name == "trip-planner" + + +class _StreamedSpend(BaseModel): + model_config = ConfigDict(frozen=True) + + call_type: str + model: str + response_cost: float + + +class _StreamedSpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: Final = asyncio.Event() + self.spend: _StreamedSpend | None = None + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + spend: Final = _StreamedSpend.model_validate(kwargs["standard_logging_object"]) + if spend.call_type != "asend_message_streaming": + return + self.spend = spend + self.logged.set() + + +@pytest.mark.skipif(importlib.util.find_spec("a2a") is None, reason="needs the optional a2a-sdk") +@pytest.mark.asyncio +async def test_handle_stream_message_bills_a_bridge_agent_under_its_name(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.agent_endpoints.a2a_endpoints import _handle_stream_message + + recorder: Final = _StreamedSpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + + response: Final = await _handle_stream_message( + api_base=None, + request_id="req-1", + params=dict(_HELLO_MESSAGE_PARAMS), + litellm_params={ + "custom_llm_provider": "langgraph", + "model": "agent", + "mock_response": "pong", + "cost_per_query": 0.25, + }, + agent_id="agent-1", + agent_name="trip-planner", + ) + _ = [chunk async for chunk in response.body_iterator] + + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + assert recorder.spend == _StreamedSpend( + call_type="asend_message_streaming", model="a2a_agent/trip-planner", response_cost=0.25 + ) + + @pytest.mark.asyncio async def test_handle_stream_message_rejects_invalid_params_with_32602(): from litellm.proxy.agent_endpoints.a2a_endpoints import _handle_stream_message