""" Test for LiteLLM A2A module. Run with: pytest tests/agent_tests/test_a2a.py -v -s """ import asyncio import os import sys import json from typing import Optional from uuid import uuid4 import pytest import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import StandardLoggingPayload sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path from a2a.types import MessageSendParams, SendMessageRequest @pytest.mark.asyncio async def test_asend_message_with_client_decorator(): """ Test asend_message standalone function with @client decorator. This tests the LiteLLM logging integration. """ litellm._turn_on_debug() from litellm.a2a_protocol import asend_message, create_a2a_client # Create the A2A client first a2a_client = await create_a2a_client(base_url="http://localhost:10001") # Build the request matching A2A SDK spec send_message_payload = { "message": { "role": "user", "parts": [ { "kind": "text", "text": "Hello from @client decorated function!", } ], "messageId": uuid4().hex, }, } request = SendMessageRequest( id=str(uuid4()), params=MessageSendParams(**send_message_payload), ) # Send message using standalone function with @client decorator response = await asend_message(a2a_client=a2a_client, request=request) # Print response for debugging print("\n=== A2A Response (standalone with @client) ===") print(response.model_dump(mode="json", exclude_none=True)) # Basic assertions assert response is not None class TestA2ALogger(CustomLogger): """Custom logger to capture A2A logging payloads for testing.""" def __init__(self): self.standard_logging_payload: Optional[StandardLoggingPayload] = None self.logged_kwargs: Optional[dict] = None self.log_success_called = False super().__init__() async def async_log_success_event( self, kwargs, response_obj, start_time, end_time ): print("TestA2ALogger: async_log_success_event called") self.log_success_called = True self.logged_kwargs = kwargs self.standard_logging_payload = kwargs.get("standard_logging_object", None) print(f"Captured standard_logging_payload: {self.standard_logging_payload}") @pytest.mark.asyncio async def test_a2a_logging_payload(): """ Test that A2A calls create a standard logging payload. Validates the @client decorator integration with LiteLLM logging. """ # Reset callbacks and set up custom logger litellm.logging_callback_manager._reset_all_callbacks() test_logger = TestA2ALogger() litellm.callbacks = [test_logger] from litellm.a2a_protocol import asend_message, create_a2a_client # Create the A2A client first a2a_client = await create_a2a_client(base_url="http://localhost:10001") # Build the request send_message_payload = { "message": { "role": "user", "parts": [ { "kind": "text", "text": "Hello! Testing logging payload.", } ], "messageId": uuid4().hex, }, } request = SendMessageRequest( id=str(uuid4()), params=MessageSendParams(**send_message_payload), ) # Send message response = await asend_message(a2a_client=a2a_client, request=request) # Give async logging time to complete await asyncio.sleep(1) # Print debug info print("\n=== Logging Validation ===") print(f"log_success_called: {test_logger.log_success_called}") print(f"standard_logging_payload: {test_logger.standard_logging_payload}") print(f"logged kwargs: {json.dumps(test_logger.logged_kwargs, indent=4, default=str)}") # Verify logging was called assert test_logger.log_success_called is True assert test_logger.standard_logging_payload is not None # Verify standard_logging_payload exists slp = test_logger.standard_logging_payload assert slp is not None # Get values from standard logging payload logged_model = slp.get("model") if isinstance(slp, dict) else getattr(slp, "model", None) logged_provider = slp.get("custom_llm_provider") if isinstance(slp, dict) else getattr(slp, "custom_llm_provider", None) call_type = slp.get("call_type") if isinstance(slp, dict) else getattr(slp, "call_type", None) response_cost = slp.get("response_cost") if isinstance(slp, dict) else getattr(slp, "response_cost", None) print(f"\n=== Standard Logging Payload Validation ===") print(f"model: {logged_model}") print(f"custom_llm_provider: {logged_provider}") print(f"call_type: {call_type}") print(f"response_cost: {response_cost}") # Verify model and custom_llm_provider are set correctly assert logged_model is not None, "model should be set" assert "a2a_agent/" in logged_model, f"model should contain 'a2a_agent/', got: {logged_model}" assert logged_provider == "a2a_agent", f"custom_llm_provider should be 'a2a_agent', got: {logged_provider}" # Verify call_type is correct for A2A assert call_type == "asend_message", f"call_type should be 'asend_message', got: {call_type}" # Verify response_cost is set to 0.0 (not None, not an error) # This confirms the A2A cost calculator is working assert response_cost is not None, "response_cost should not be None" assert response_cost == 0.0, f"response_cost should be 0.0 for A2A, got: {response_cost}" @pytest.mark.asyncio async def test_pydantic_ai_non_streaming(): """ Test non-streaming requests to Pydantic AI agents. Pydantic AI agents follow A2A protocol but don't support streaming. This test validates non-streaming requests work correctly. """ litellm._turn_on_debug() from litellm.a2a_protocol import asend_message # Build the request send_message_payload = { "message": { "role": "user", "parts": [ { "kind": "text", "text": "Hello from Pydantic AI test!", } ], "messageId": uuid4().hex, }, } request = SendMessageRequest( id=str(uuid4()), params=MessageSendParams(**send_message_payload), ) # Send message using Pydantic AI provider response = await asend_message( request=request, api_base="http://localhost:9999", litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, ) # Print response for debugging print("\n=== Pydantic AI Non-Streaming Response ===") print(response.model_dump(mode="json", exclude_none=True)) # Basic assertions assert response is not None assert hasattr(response, "result") # Verify result structure result = response.result assert result is not None # Pydantic AI returns a task with history/artifacts, not a direct message # Check for either format result_dict = result if isinstance(result, dict) else result.model_dump(mode="python", exclude_none=True) has_message = "message" in result_dict has_history = "history" in result_dict has_artifacts = "artifacts" in result_dict assert has_message or has_history or has_artifacts, ( f"Result should contain 'message', 'history', or 'artifacts'. Got: {list(result_dict.keys())}" ) # If it's a task response (Pydantic AI style), verify we got agent response if has_history: history = result_dict.get("history", []) agent_messages = [m for m in history if m.get("role") == "agent"] assert len(agent_messages) > 0, "Should have at least one agent message in history" # Verify agent message has text content agent_msg = agent_messages[-1] parts = agent_msg.get("parts", []) text_parts = [p for p in parts if p.get("kind") == "text"] assert len(text_parts) > 0, "Agent message should have text content" print(f"\nAgent response: {text_parts[0].get('text')}") @pytest.mark.asyncio async def test_pydantic_ai_fake_streaming(): """ Test fake streaming for Pydantic AI agents. Pydantic AI agents don't support streaming natively. This test validates that fake streaming works by converting non-streaming responses into streaming chunks. """ litellm._turn_on_debug() from litellm.a2a_protocol import asend_message_streaming # Build the request from a2a.types import SendStreamingMessageRequest send_message_payload = { "message": { "role": "user", "parts": [ { "kind": "text", "text": "Hello from Pydantic AI streaming test!", } ], "messageId": uuid4().hex, }, } request = SendStreamingMessageRequest( id=str(uuid4()), params=MessageSendParams(**send_message_payload), ) # Send streaming message using Pydantic AI provider print("\n=== Pydantic AI Fake Streaming Response ===") chunks_received = 0 task_event_received = False working_event_received = False artifact_event_received = False completed_event_received = False async for chunk in asend_message_streaming( request=request, api_base="http://localhost:9999", litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, ): chunks_received += 1 print(f"\nChunk {chunks_received}:") # Convert chunk to dict for inspection chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else chunk print(json.dumps(chunk_dict, indent=2)) # Check event types result = chunk_dict.get("result", {}) kind = result.get("kind") if kind == "task": task_event_received = True elif kind == "status-update": status = result.get("status", {}) state = status.get("state") if state == "working": working_event_received = True elif state == "completed": completed_event_received = True elif kind == "artifact-update": artifact_event_received = True print(f"\n=== Streaming Summary ===") print(f"Total chunks received: {chunks_received}") print(f"Task event received: {task_event_received}") print(f"Working event received: {working_event_received}") print(f"Artifact event received: {artifact_event_received}") print(f"Completed event received: {completed_event_received}") # Verify we received chunks assert chunks_received > 0, "Should receive at least one chunk" # Verify all required event types were received assert task_event_received, "Should receive task event" assert working_event_received, "Should receive working status event" assert artifact_event_received, "Should receive artifact update event" assert completed_event_received, "Should receive completed status event"