route_a2a_agent_request

This commit is contained in:
Ishaan Jaffer 2026-02-03 14:04:32 -08:00
parent 9dd5b0c0be
commit 3c73526169
2 changed files with 17 additions and 25 deletions

View file

@ -9,24 +9,6 @@ from typing import Any, Optional
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.route_llm_request import (
ROUTE_ENDPOINT_MAPPING,
ProxyModelNotFoundError,
)
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]:
@ -35,6 +17,13 @@ async def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]:
Returns None if not an A2A request (allows normal routing to continue).
"""
# Import here to avoid circular imports
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.route_llm_request import (
ROUTE_ENDPOINT_MAPPING,
ProxyModelNotFoundError,
)
model_name = data.get("model", "")
# Check if this is an A2A agent request

View file

@ -4,12 +4,6 @@ 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
@ -18,6 +12,11 @@ else:
LitellmRouter = Any
def _is_a2a_agent_model(model_name: Any) -> bool:
"""Check if the model name is for an A2A agent (a2a/ prefix)."""
return isinstance(model_name, str) and model_name.startswith("a2a/")
ROUTE_ENDPOINT_MAPPING = {
"acompletion": "/chat/completions",
"atext_completion": "/completions",
@ -328,7 +327,11 @@ 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", "")):
elif _is_a2a_agent_model(data.get("model", "")):
from litellm.proxy.agent_endpoints.a2a_routing import (
route_a2a_agent_request,
)
return await route_a2a_agent_request(data, route_type)
elif user_model is not None: