mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
fix(a2a)!: bill completion-bridge agents at their own prices on /a2a send and stream (#49150)
* fix(a2a): bill completion-bridge agents at their own prices on /a2a send and stream message/send on an agent served through the completion bridge (custom_llm_provider agents such as LangGraph or databricks_agent) logged a $0, 0-token row, and message/stream logged no row at all, so cost_per_query and per-token prices never reached spend. The bridge send path now records counted usage, the agent's pricing keys, and the agent id on its logging object, the same way URL agents do. The bridge stream is wrapped in the A2A streaming logger, which now reads the bridge's artifact event for the reply text and gets the agent name from the proxy. Streamed calls also count prompt tokens again, since the message text reader now reads pydantic messages as dicts. * fix(a2a): count chunked and Task replies once and skip agent prices on error replies Bridge streams now follow the A2A append flag per artifact, so chunked pydantic-ai, watsonx, and AgentCore replies count the whole reply once. Send replies that are a Task count their artifact text. A JSON-RPC error reply is logged at cost 0 instead of the agent's price. * fix(a2a): keep one reply text per stream and bill every artifact and status reply * fix(a2a): ignore a null error field on streams and run the stream per-token billing test A stream event whose error field is null no longer zeroes the call's cost. The integration test for per-token stream billing now runs instead of skipping, and the URL agent SDK case skips when the optional a2a-sdk is not installed. * fix(a2a): bill streams the client drops and skip the bridge's working placeholder A stream the client closes or cancels mid-reply now logs its spend row at the agent's prices on both the bridge and URL paths, and a stream dropped before the agent sends anything still logs nothing. The bridge's "Processing request..." status no longer counts as reply text, so an empty bridge reply bills no completion tokens, as send does. Adds proxy-level tests that a bridge stream bills under the agent's card name. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
zoveki
GitHub
parent
87d91d78b7
commit
c8cef1a95b
13 changed files with 807 additions and 84 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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/<agent_name> 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(
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
|
|
|
|||
|
|
@ -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}],
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}"),))
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue