mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
Merge ce4a25721f into 570f41daad
This commit is contained in:
commit
4bfc5ce78d
3 changed files with 182 additions and 10 deletions
|
|
@ -969,6 +969,22 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr
|
|||
return getattr(llm_router, f"{route_type}")(**data)
|
||||
|
||||
elif data["model"] not in router_model_names:
|
||||
a2a_model_name: Final[object] = JSON_OBJECT_ADAPTER.validate_python(data).get("model", "")
|
||||
if isinstance(a2a_model_name, str) and _is_a2a_agent_model(a2a_model_name):
|
||||
from litellm.proxy.common_utils.registry_read_through import (
|
||||
get_agent_with_read_through,
|
||||
)
|
||||
|
||||
agent_name: Final = a2a_model_name[4:]
|
||||
if await get_agent_with_read_through(agent_name) is not None:
|
||||
from litellm.proxy.agent_endpoints.a2a_routing import (
|
||||
route_a2a_agent_request,
|
||||
)
|
||||
|
||||
result: Final = await route_a2a_agent_request(data, route_type, user_api_key_dict=user_api_key_dict)
|
||||
if result is not None:
|
||||
return result
|
||||
|
||||
# Check wildcards before checking deployment_names
|
||||
# Priority: 1. Exact model_name match, 2. Wildcard match, 3. deployment_names match
|
||||
if llm_router.settings().router_general_settings.pass_through_all_models:
|
||||
|
|
@ -1028,16 +1044,12 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr
|
|||
except Exception:
|
||||
# If router fails (e.g., model not found in router), fall back to direct call
|
||||
return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs())
|
||||
elif _is_a2a_agent_model(data.get("model", "")):
|
||||
from litellm.proxy.agent_endpoints.a2a_routing import (
|
||||
route_a2a_agent_request,
|
||||
elif isinstance(a2a_model_name, str) and _is_a2a_agent_model(a2a_model_name):
|
||||
raise ProxyModelNotFoundError(
|
||||
route=ROUTE_ENDPOINT_MAPPING.get(route_type, route_type),
|
||||
model_name=a2a_model_name,
|
||||
retryable_with_model_read_through=False,
|
||||
)
|
||||
|
||||
result: Final = await route_a2a_agent_request(data, route_type, user_api_key_dict=user_api_key_dict)
|
||||
if result is not None:
|
||||
return result
|
||||
# Fall through to raise exception below if result is None
|
||||
|
||||
elif user_model is not None or route_type == "allm_passthrough_route":
|
||||
return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs())
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from typing import Final
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.client import Gateway, object_value, string_value
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
|
@ -164,3 +164,105 @@ def test_a2a_completion_through_proxy(
|
|||
|
||||
requests: Final = wire.drain()
|
||||
assert len(tuple(item for item in requests if item.method == "POST")) == 1
|
||||
|
||||
|
||||
def test_registered_a2a_agent_routes_before_wildcard_model(gateway: Gateway) -> None:
|
||||
agent_name: Final = "translation-a2a-" + uuid.uuid4().hex
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if request.method == "GET":
|
||||
assert request.target in ("/.well-known/agent-card.json", "/.well-known/agent.json")
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"protocolVersion": "0.3",
|
||||
"name": agent_name,
|
||||
"description": "Scripted A2A peer",
|
||||
"version": "1.0.0",
|
||||
"url": wire.url + "/",
|
||||
"capabilities": {"streaming": True},
|
||||
"defaultInputModes": ["text"],
|
||||
"defaultOutputModes": ["text"],
|
||||
"skills": [],
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
assert request.method == "POST" and request.target == "/"
|
||||
outbound: Final = json.loads(request.body)
|
||||
request_id: Final = outbound["id"]
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"kind": "message",
|
||||
"role": "agent",
|
||||
"messageId": agent_name,
|
||||
"parts": [{"kind": "text", "text": "A2A wildcard response"}],
|
||||
},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
with wire_server(upstream) as wire, gateway.scenario() as scenario:
|
||||
wildcard_prefix: Final = "integration-" + uuid.uuid4().hex
|
||||
wildcard_created: Final = gateway.request(
|
||||
"POST",
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": f"{wildcard_prefix}/*",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "integration-provider-key",
|
||||
"api_base": gateway.upstream_url,
|
||||
},
|
||||
},
|
||||
)
|
||||
assert wildcard_created.status_code == 200, wildcard_created.text
|
||||
wildcard_id: Final = string_value(object_value(wildcard_created.json()["model_info"])["id"])
|
||||
scenario.cleanups.callback(scenario.delete_model, wildcard_id)
|
||||
|
||||
missing_agent: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": f"a2a/{agent_name}-missing",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
assert missing_agent.status_code == 400, missing_agent.text
|
||||
|
||||
card: Final = {
|
||||
"protocolVersion": "0.3",
|
||||
"name": agent_name,
|
||||
"description": "Scripted A2A peer",
|
||||
"version": "1.0.0",
|
||||
"url": wire.url + "/",
|
||||
"capabilities": {"streaming": True},
|
||||
"defaultInputModes": ["text"],
|
||||
"defaultOutputModes": ["text"],
|
||||
"skills": [],
|
||||
}
|
||||
created: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/agents",
|
||||
{"agent_name": agent_name, "agent_card_params": card},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
agent_id: Final = string_value(created.json()["agent_id"])
|
||||
scenario.cleanups.callback(
|
||||
lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}")
|
||||
)
|
||||
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": f"a2a/{agent_name}",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["choices"][0]["message"]["content"] == "A2A wildcard response"
|
||||
|
|
|
|||
|
|
@ -2152,6 +2152,64 @@ async def test_route_request_routing_group_name_passes_model_gate():
|
|||
spy.assert_called_once_with(**data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("agent_registered", [True, False])
|
||||
async def test_route_request_a2a_agent_precedes_pattern_route(
|
||||
agent_registered: bool, fresh_agent_read_through: object, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
import uuid
|
||||
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
agent_name: Final = "pattern-route-agent-" + uuid.uuid4().hex
|
||||
model_name: Final = f"a2a/{agent_name}"
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
|
||||
}
|
||||
]
|
||||
)
|
||||
pattern_route: Final = AsyncMock(return_value="pattern response")
|
||||
monkeypatch.setattr(router, "acompletion", pattern_route)
|
||||
agent: Final = AgentResponse(
|
||||
agent_id=f"{agent_name}-id",
|
||||
agent_name=agent_name,
|
||||
agent_card_params={"name": agent_name, "url": "http://agent.example"},
|
||||
)
|
||||
agents_find_unique: Final = AsyncMock(side_effect=[None, agent] if agent_registered else [None, None])
|
||||
fake_prisma, _ = _fake_prisma_client_with_models([])
|
||||
fake_prisma.db.litellm_agentstable = SimpleNamespace(find_unique=agents_find_unique)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
|
||||
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
|
||||
data: Final = {"model": model_name, "messages": [{"role": "user", "content": "hi"}]}
|
||||
admin_auth: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
if agent_registered:
|
||||
agent_route: Final = AsyncMock(return_value="agent response")
|
||||
monkeypatch.setattr(litellm, "acompletion", agent_route)
|
||||
agent_response: Final[object] = await (await route_request(data, router, None, "acompletion", admin_auth))
|
||||
assert agent_response == "agent response"
|
||||
agent_route.assert_awaited_once()
|
||||
agent_call: Final = agent_route.await_args
|
||||
assert agent_call is not None
|
||||
assert agent_call.kwargs["api_base"] == "http://agent.example"
|
||||
pattern_route.assert_not_awaited()
|
||||
return
|
||||
|
||||
pattern_response: Final[object] = await (await route_request(data, router, None, "acompletion"))
|
||||
assert pattern_response == "pattern response"
|
||||
pattern_route.assert_awaited_once_with(**data)
|
||||
assert agents_find_unique.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_route_request_a2a_agent_miss_does_not_consume_model_read_through(fresh_agent_read_through, monkeypatch):
|
||||
from types import SimpleNamespace
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue