diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 5cb238904f..21b2ae81ad 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -7,7 +7,9 @@ Provides standalone functions with @client decorator for LiteLLM logging integra import asyncio from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union +import litellm from litellm._logging import verbose_logger +from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -38,6 +40,29 @@ except ImportError: pass +def _set_usage_on_logging_obj( + kwargs: Dict[str, Any], + prompt_tokens: int, + completion_tokens: int, +) -> None: + """ + Set usage on litellm_logging_obj for standard logging payload. + + Args: + kwargs: The kwargs dict containing litellm_logging_obj + prompt_tokens: Number of input tokens + completion_tokens: Number of output tokens + """ + litellm_logging_obj = kwargs.get("litellm_logging_obj") + if litellm_logging_obj is not None: + usage = litellm.Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + litellm_logging_obj.model_call_details["usage"] = usage + + def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: """ Extract agent info and set model/custom_llm_provider for cost tracking. @@ -123,6 +148,20 @@ async def asend_message( # Wrap in LiteLLM response type for _hidden_params support response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response) + # Calculate token usage from request and response + response_dict = a2a_response.model_dump(mode="json", exclude_none=True) + prompt_tokens, completion_tokens, _ = A2ARequestUtils.calculate_usage_from_request_response( + request=request, + response_dict=response_dict, + ) + + # Set usage on logging obj for standard logging payload + _set_usage_on_logging_obj( + kwargs=kwargs, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + return response diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py new file mode 100644 index 0000000000..d7bd27246a --- /dev/null +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -0,0 +1,165 @@ +""" +A2A Streaming Iterator with token tracking and logging support. +""" + +import asyncio +from datetime import datetime +from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional + +import litellm +from litellm._logging import verbose_logger +from litellm.a2a_protocol.utils import A2ARequestUtils +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.thread_pool_executor import executor + +if TYPE_CHECKING: + from a2a.types import SendStreamingMessageRequest, SendStreamingMessageResponse + + +class A2AStreamingIterator: + """ + Async iterator for A2A streaming responses with token tracking. + + Collects chunks, extracts text, and logs usage on completion. + """ + + def __init__( + self, + stream: AsyncIterator["SendStreamingMessageResponse"], + request: "SendStreamingMessageRequest", + logging_obj: LiteLLMLoggingObj, + agent_name: str = "unknown", + ): + self.stream = stream + self.request = request + self.logging_obj = logging_obj + self.agent_name = agent_name + self.start_time = datetime.now() + + # Collect chunks for token counting + self.chunks: List[Any] = [] + self.collected_text_parts: List[str] = [] + self.final_chunk: Optional[Any] = None + + def __aiter__(self): + return self + + async def __anext__(self) -> "SendStreamingMessageResponse": + try: + chunk = await self.stream.__anext__() + + # Store chunk + self.chunks.append(chunk) + + # Extract text from chunk for token counting + self._collect_text_from_chunk(chunk) + + # Check if this is the final chunk (completed status) + if self._is_completed_chunk(chunk): + 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: Any) -> None: + """Extract text from a streaming chunk and add to collected parts.""" + try: + chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} + text = A2ARequestUtils.extract_text_from_response(chunk_dict) + if text: + self.collected_text_parts.append(text) + except Exception: + verbose_logger.debug("Failed to extract text from A2A streaming chunk") + + def _is_completed_chunk(self, chunk: Any) -> bool: + """Check if chunk indicates stream completion.""" + try: + chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} + result = chunk_dict.get("result", {}) + if isinstance(result, dict): + status = result.get("status", {}) + if isinstance(status, dict): + return status.get("state") == "completed" + except Exception: + pass + return False + + async def _handle_stream_complete(self) -> None: + """Handle logging and token counting when stream completes.""" + try: + end_time = datetime.now() + + # Calculate tokens from collected text + input_message = A2ARequestUtils.get_input_message_from_request(self.request) + input_text = A2ARequestUtils.extract_text_from_message(input_message) + prompt_tokens = A2ARequestUtils.count_tokens(input_text) + + # Use the last (most complete) text from chunks + output_text = self.collected_text_parts[-1] if self.collected_text_parts else "" + completion_tokens = A2ARequestUtils.count_tokens(output_text) + + total_tokens = prompt_tokens + completion_tokens + + # Create usage object + usage = litellm.Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + ) + + # Set usage on logging obj + self.logging_obj.model_call_details["usage"] = usage + + # Build result for logging + result = self._build_logging_result(usage) + + # Call success handlers + asyncio.create_task( + self.logging_obj.async_success_handler( + result=result, + start_time=self.start_time, + end_time=end_time, + cache_hit=None, + ) + ) + + executor.submit( + self.logging_obj.success_handler, + result=result, + cache_hit=None, + start_time=self.start_time, + end_time=end_time, + ) + + verbose_logger.info( + f"A2A streaming completed: prompt_tokens={prompt_tokens}, " + f"completion_tokens={completion_tokens}, total_tokens={total_tokens}" + ) + + except Exception as e: + verbose_logger.debug(f"Error in A2A streaming completion handler: {e}") + + def _build_logging_result(self, usage: litellm.Usage) -> Dict[str, Any]: + """Build a result dict for logging.""" + result: Dict[str, Any] = { + "id": getattr(self.request, "id", "unknown"), + "jsonrpc": "2.0", + "usage": usage.model_dump() if hasattr(usage, "model_dump") else dict(usage), + } + + # Add final chunk result if available + if self.final_chunk: + try: + chunk_dict = self.final_chunk.model_dump(mode="json", exclude_none=True) + result["result"] = chunk_dict.get("result", {}) + except Exception: + pass + + return result + diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py new file mode 100644 index 0000000000..1cdbde9775 --- /dev/null +++ b/litellm/a2a_protocol/utils.py @@ -0,0 +1,138 @@ +""" +Utility functions for A2A protocol. +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union + +import litellm +from litellm._logging import verbose_logger + +if TYPE_CHECKING: + from a2a.types import SendMessageRequest, SendStreamingMessageRequest + + +class A2ARequestUtils: + """Utility class for A2A request/response processing.""" + + @staticmethod + def extract_text_from_message(message: Any) -> str: + """ + Extract text content from A2A message parts. + + Args: + message: A2A message dict or object with 'parts' containing text parts + + Returns: + Concatenated text from all text parts + """ + if message is None: + return "" + + # Handle both dict and object access + if isinstance(message, dict): + parts = message.get("parts", []) + else: + parts = getattr(message, "parts", []) or [] + + text_parts: List[str] = [] + for part in parts: + if isinstance(part, dict): + if part.get("kind") == "text": + text_parts.append(part.get("text", "")) + else: + if getattr(part, "kind", None) == "text": + text_parts.append(getattr(part, "text", "")) + + return " ".join(text_parts) + + @staticmethod + def extract_text_from_response(response_dict: Dict[str, Any]) -> str: + """ + Extract text content from A2A response result. + + Args: + response_dict: A2A response dict with 'result' containing message + + Returns: + Text from response message parts + """ + result = response_dict.get("result", {}) + if not isinstance(result, dict): + return "" + + message = result.get("message", {}) + return A2ARequestUtils.extract_text_from_message(message) + + @staticmethod + def get_input_message_from_request( + request: "Union[SendMessageRequest, SendStreamingMessageRequest]", + ) -> Any: + """ + Extract the input message from an A2A request. + + Args: + request: The A2A SendMessageRequest or SendStreamingMessageRequest + + Returns: + The message object/dict or None + """ + params = getattr(request, "params", None) + if params is None: + return None + return getattr(params, "message", None) + + @staticmethod + def count_tokens(text: str) -> int: + """ + Count tokens in text using litellm.token_counter. + + Args: + text: Text to count tokens for + + Returns: + Token count, or 0 if counting fails + """ + if not text: + return 0 + try: + return litellm.token_counter(text=text) + except Exception: + verbose_logger.debug("Failed to count tokens") + return 0 + + @staticmethod + def calculate_usage_from_request_response( + request: "Union[SendMessageRequest, SendStreamingMessageRequest]", + response_dict: Dict[str, Any], + ) -> Tuple[int, int, int]: + """ + Calculate token usage from A2A request and response. + + Args: + request: The A2A SendMessageRequest or SendStreamingMessageRequest + response_dict: The A2A response as a dict + + Returns: + Tuple of (prompt_tokens, completion_tokens, total_tokens) + """ + # Count input tokens + input_message = A2ARequestUtils.get_input_message_from_request(request) + input_text = A2ARequestUtils.extract_text_from_message(input_message) + prompt_tokens = A2ARequestUtils.count_tokens(input_text) + + # Count output tokens + output_text = A2ARequestUtils.extract_text_from_response(response_dict) + completion_tokens = A2ARequestUtils.count_tokens(output_text) + + total_tokens = prompt_tokens + completion_tokens + + return prompt_tokens, completion_tokens, total_tokens + + +# Backwards compatibility aliases +def extract_text_from_a2a_message(message: Any) -> str: + return A2ARequestUtils.extract_text_from_message(message) + + +def extract_text_from_a2a_response(response_dict: Dict[str, Any]) -> str: + return A2ARequestUtils.extract_text_from_response(response_dict) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 2b52d58e29..6e845c56b5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1651,6 +1651,11 @@ class Logging(LiteLLMLoggingBaseClass): result = self._handle_non_streaming_google_genai_generate_content_response_logging( result=result ) + elif ( + self.call_type == CallTypes.asend_message.value + or self.call_type == CallTypes.send_message.value + ): + result = self._handle_a2a_response_logging(result=result) logging_result = self.normalize_logging_result(result=result) @@ -3243,6 +3248,29 @@ class Logging(LiteLLMLoggingBaseClass): ) return result + def _handle_a2a_response_logging(self, result: Any) -> Any: + """ + Handles logging for A2A (Agent-to-Agent) responses. + + Adds usage from model_call_details to the result if available. + Uses Pydantic's model_copy to avoid modifying the original response. + + Args: + result: The LiteLLMSendMessageResponse from the A2A call + + Returns: + The response object with usage added if available + """ + # Get usage from model_call_details (set by asend_message) + usage = self.model_call_details.get("usage") + if usage is None: + return result + + # Deep copy result and add usage + result_copy = result.model_copy(deep=True) + result_copy.usage = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) + return result_copy + def _get_masked_values( sensitive_object: dict, @@ -3806,7 +3834,10 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 from litellm.integrations.opentelemetry import ( OpenTelemetryConfig, ) - from litellm.integrations.weave.weave_otel import WeaveOtelLogger, get_weave_otel_config + from litellm.integrations.weave.weave_otel import ( + WeaveOtelLogger, + get_weave_otel_config, + ) weave_otel_config = get_weave_otel_config() diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 2eb26dc622..7d8251ebff 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -221,6 +221,9 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase): result: Optional[Dict[str, Any]] = None error: Optional[Dict[str, Any]] = None + # LiteLLM usage tracking + usage: Optional[Dict[str, Any]] = None + model_config = {"extra": "allow"} # LiteLLM private attributes for logging/cost tracking