fix mypy linting

Author: ishaan-jaff

litellm/a2a/__init__.py litellm/a2a_protocol/__init__.pylitellm/a2a/__init__.py renamed to litellm/a2a_protocol/__init__.py

@@ -8,7 +8,7 @@
 
 Example usage (standalone functions with @client decorator):
     ```python
-    from litellm.a2a import asend_message
+    from litellm.a2a_protocol import asend_message
     from a2a.types import SendMessageRequest, MessageSendParams
     from uuid import uuid4
 
@@ -31,15 +31,15 @@
 
 Example usage (class-based):
     ```python
-    from litellm.a2a import A2AClient
+    from litellm.a2a_protocol import A2AClient
 
     client = A2AClient(base_url="http://localhost:10001")
     response = await client.send_message(request)
     ```
 """
 
-from litellm.a2a.client import A2AClient
-from litellm.a2a.main import (
+from litellm.a2a_protocol.client import A2AClient
+from litellm.a2a_protocol.main import (
     aget_agent_card,
     asend_message,
     asend_message_streaming,

litellm/a2a/client.py litellm/a2a_protocol/client.pylitellm/a2a/client.py renamed to litellm/a2a_protocol/client.py

@@ -26,7 +26,7 @@ class A2AClient:
 
     Example:
         ```python
-        from litellm.a2a import A2AClient
+        from litellm.a2a_protocol import A2AClient
         from a2a.types import SendMessageRequest, MessageSendParams
         from uuid import uuid4
 
@@ -68,7 +68,7 @@ def __init__(
     async def _get_client(self) -> "A2AClientType":
         """Get or create the underlying A2A client."""
         if self._a2a_client is None:
-            from litellm.a2a.main import create_a2a_client
+            from litellm.a2a_protocol.main import create_a2a_client
 
             self._a2a_client = await create_a2a_client(
                 base_url=self.base_url,
@@ -79,7 +79,7 @@ async def _get_client(self) -> "A2AClientType":
 
     async def get_agent_card(self) -> "AgentCard":
         """Fetch the agent card from the server."""
-        from litellm.a2a.main import aget_agent_card
+        from litellm.a2a_protocol.main import aget_agent_card
 
         return await aget_agent_card(
             base_url=self.base_url,
@@ -91,7 +91,7 @@ async def send_message(
         self, request: "SendMessageRequest"
     ) -> LiteLLMSendMessageResponse:
         """Send a message to the A2A agent."""
-        from litellm.a2a.main import asend_message
+        from litellm.a2a_protocol.main import asend_message
 
         a2a_client = await self._get_client()
         return await asend_message(a2a_client=a2a_client, request=request)
@@ -100,7 +100,7 @@ async def send_message_streaming(
         self, request: "SendStreamingMessageRequest"
     ) -> AsyncIterator["SendStreamingMessageResponse"]:
         """Send a streaming message to the A2A agent."""
-        from litellm.a2a.main import asend_message_streaming
+        from litellm.a2a_protocol.main import asend_message_streaming
 
         a2a_client = await self._get_client()
         async for chunk in asend_message_streaming(a2a_client=a2a_client, request=request):

litellm/a2a/cost_calculator.py litellm/a2a_protocol/cost_calculator.pylitellm/a2a/cost_calculator.py renamed to litellm/a2a_protocol/cost_calculator.py

@@ -91,7 +91,7 @@ async def asend_message(
 
     Example:
         ```python
-        from litellm.a2a import asend_message, create_a2a_client
+        from litellm.a2a_protocol import asend_message, create_a2a_client
         from a2a.types import SendMessageRequest, MessageSendParams
         from uuid import uuid4
 
@@ -205,7 +205,7 @@ async def create_a2a_client(
 
     Example:
         ```python
-        from litellm.a2a import create_a2a_client, asend_message
+        from litellm.a2a_protocol import create_a2a_client, asend_message
 
         # Create client once
         client = await create_a2a_client(base_url="http://localhost:10001")

litellm/a2a/main.py litellm/a2a_protocol/main.pylitellm/a2a/main.py renamed to litellm/a2a_protocol/main.py

@@ -939,7 +939,7 @@ def completion_cost(  # noqa: PLR0915
                     CallTypes.asend_message.value,
                     CallTypes.send_message.value,
                 ):
-                    from litellm.a2a.cost_calculator import A2ACostCalculator
+                    from litellm.a2a_protocol.cost_calculator import A2ACostCalculator
 
                     return A2ACostCalculator.calculate_a2a_cost(
                         litellm_logging_obj=litellm_logging_obj

litellm/cost_calculator.py

@@ -138,7 +138,7 @@ async def invoke_agent_a2a(
     - message/send: Send a message and get a response
     - message/stream: Send a message and stream the response
     """
-    from litellm.a2a import asend_message, create_a2a_client
+    from litellm.a2a_protocol import asend_message, create_a2a_client
     from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
     from litellm.proxy.proxy_server import (
         general_settings,

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -32,7 +32,7 @@ async def test_asend_message_with_client_decorator():
     This tests the LiteLLM logging integration.
     """
     litellm._turn_on_debug()
-    from litellm.a2a import asend_message, create_a2a_client
+    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")
@@ -97,7 +97,7 @@ async def test_a2a_logging_payload():
     test_logger = TestA2ALogger()
     litellm.callbacks = [test_logger]
 
-    from litellm.a2a import asend_message, create_a2a_client
+    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")

tests/agent_tests/test_a2a.py

@@ -75,10 +75,10 @@ async def mock_add_litellm_data(data, **kwargs):
         "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
         side_effect=mock_add_litellm_data,
     ) as mock_add_data, patch(
-        "litellm.a2a.create_a2a_client",
+        "litellm.a2a_protocol.create_a2a_client",
         new_callable=AsyncMock,
     ), patch(
-        "litellm.a2a.asend_message",
+        "litellm.a2a_protocol.asend_message",
         new_callable=AsyncMock,
         return_value=mock_response,
     ), patch(