is_a2a_agent_model

litellm/proxy/agent_endpoints/a2a_routing.py

@@ -16,6 +16,19 @@ from litellm.proxy.route_llm_request import (
 )
 
 
+def is_a2a_agent_model(model_name: Any) -> bool:
+    """
+    Check if the model name is for an A2A agent (a2a/ prefix).
+    
+    Args:
+        model_name: The model name to check
+        
+    Returns:
+        True if this is an A2A agent model, False otherwise
+    """
+    return isinstance(model_name, str) and model_name.startswith("a2a/")
+
+
 async def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]:
     """
     Route A2A agent requests directly to litellm with injected API base.

litellm/proxy/route_llm_request.py

@@ -4,6 +4,12 @@ from fastapi import HTTPException, status
 
 import litellm
 
+# Check for A2A agent models (a2a/ prefix)
+from litellm.proxy.agent_endpoints.a2a_routing import (
+    is_a2a_agent_model,
+    route_a2a_agent_request,
+)
+
 if TYPE_CHECKING:
     from litellm.router import Router as _Router
 
@@ -322,6 +328,8 @@ async def route_request(
                 except Exception:
                     # If router fails (e.g., model not found in router), fall back to direct call
                     return getattr(litellm, f"{route_type}")(**data)
+            elif is_a2a_agent_model(data.get("model", "")):
+                return await route_a2a_agent_request(data, route_type)
 
     elif user_model is not None:
         return getattr(litellm, f"{route_type}")(**data)

tests/test_litellm/proxy/test_route_a2a_models.py

@@ -0,0 +1,105 @@
+"""
+Test A2A model routing in proxy.
+
+Maps to: litellm/proxy/agent_endpoints/a2a_routing.py
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../../.."))
+
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from litellm.proxy.agent_endpoints.a2a_routing import route_a2a_agent_request
+from litellm.proxy.route_llm_request import route_request
+
+
+@pytest.mark.asyncio
+async def test_route_a2a_model_bypasses_router():
+    """Test that a2a/ prefixed models bypass router and go directly to litellm with api_base"""
+
+    # Mock data for chat completion with a2a model
+    data = {
+        "model": "a2a/test-agent",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    # Mock router that doesn't have the a2a model
+    mock_router = Mock()
+    mock_router.model_names = ["gpt-4", "gpt-3.5-turbo"]
+    mock_router.deployment_names = []
+    mock_router.has_model_id = Mock(return_value=False)
+    mock_router.model_group_alias = None
+    mock_router.router_general_settings = Mock(pass_through_all_models=False)
+    mock_router.default_deployment = None
+    mock_router.pattern_router = Mock(patterns=[])
+    mock_router.map_team_model = Mock(return_value=None)
+
+    # Mock agent in registry
+    from litellm.types.agents import AgentResponse
+    
+    mock_agent = AgentResponse(
+        agent_id="test-agent-id",
+        agent_name="test-agent",
+        agent_card_params={"url": "http://agent.example.com"},
+        litellm_params=None,
+    )
+    
+    mock_registry = Mock()
+    mock_registry.get_agent_by_name = Mock(return_value=mock_agent)
+
+    # Mock litellm.acompletion to verify it's called
+    mock_acompletion = AsyncMock(return_value={"id": "test-response"})
+
+    with patch("litellm.acompletion", mock_acompletion):
+        with patch(
+            "litellm.proxy.agent_endpoints.a2a_routing.global_agent_registry",
+            mock_registry,
+        ):
+            result = await route_request(
+                data=data,
+                llm_router=mock_router,
+                user_model=None,
+                route_type="acompletion",
+            )
+
+            # Verify litellm.acompletion was called with api_base injected
+            mock_acompletion.assert_called_once()
+            call_kwargs = mock_acompletion.call_args.kwargs
+            assert call_kwargs["model"] == "a2a/test-agent"
+            assert call_kwargs["api_base"] == "http://agent.example.com"
+
+
+@pytest.mark.asyncio  
+async def test_route_non_a2a_model_raises_error_if_not_in_router():
+    """Test that non-a2a models that aren't in router raise an error"""
+
+    # Mock data for chat completion with model not in router
+    data = {
+        "model": "unknown-model",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    # Mock router without the model
+    mock_router = Mock()
+    mock_router.model_names = ["gpt-4", "gpt-3.5-turbo"]
+    mock_router.deployment_names = []
+    mock_router.has_model_id = Mock(return_value=False)
+    mock_router.model_group_alias = None
+    mock_router.router_general_settings = Mock(pass_through_all_models=False)
+    mock_router.default_deployment = None
+    mock_router.pattern_router = Mock(patterns=[])
+    mock_router.map_team_model = Mock(return_value=None)
+
+    # Should raise ProxyModelNotFoundError
+    from litellm.proxy.route_llm_request import ProxyModelNotFoundError
+    
+    with pytest.raises(ProxyModelNotFoundError):
+        await route_request(
+            data=data,
+            llm_router=mock_router,
+            user_model=None,
+            route_type="acompletion",
+        )