From 9dd5b0c0becffae45742571bfb1e0ba494316ea8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 3 Feb 2026 14:01:15 -0800 Subject: [PATCH] is_a2a_agent_model --- litellm/proxy/agent_endpoints/a2a_routing.py | 13 +++ litellm/proxy/route_llm_request.py | 8 ++ .../proxy/test_route_a2a_models.py | 105 ++++++++++++++++++ 3 files changed, 126 insertions(+) create mode 100644 tests/test_litellm/proxy/test_route_a2a_models.py diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index f48cbbb18a8..591d0b94c06 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/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. diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index c6a93164d49..6147463efde 100644 --- a/litellm/proxy/route_llm_request.py +++ b/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) diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py new file mode 100644 index 00000000000..e8d89c0c7eb --- /dev/null +++ b/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", + )