mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
fix(a2a): keep config-defined agents registered and accept the documented agents: key
The public A2A guide tells users to declare agents under a top-level `agents:` key, but the proxy only ever read `agent_list:`, so the documented config was silently ignored and GET /v1/agents returned an empty list. Accept `agents` as the documented spelling and keep `agent_list` working for anyone who found it by reading the source. Selection is by key presence, so an explicitly empty `agents: []` is not overridden by leftover legacy entries. Config-defined agents were also dropped on any database-backed gateway: the periodic reload rebuilt the registry from the DB rows plus a module global that was declared and never assigned. The registry now remembers the agents it loaded from config.yaml and replays them on every rebuild. A database row wins a name collision, mirroring how config-declared MCP servers are unioned under the database registry, so name lookups and deregistration keep addressing exactly one agent. Resolves LIT-4978
This commit is contained in:
Yassin Kortam
parent
fa56283806
commit
a4d28e3df1
4 changed files with 331 additions and 18 deletions
|
|
@ -89,6 +89,7 @@ def agents_table(prisma_client: PrismaClient) -> AgentTableClient:
|
|||
class AgentRegistry:
|
||||
def __init__(self):
|
||||
self.agent_list: list[AgentResponse] = []
|
||||
self.config_agents: tuple[AgentConfig, ...] = ()
|
||||
|
||||
def reset_agent_list(self):
|
||||
self.agent_list = []
|
||||
|
|
@ -117,9 +118,21 @@ class AgentRegistry:
|
|||
return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest()
|
||||
|
||||
def load_agents_from_config(self, agent_config: Sequence[AgentConfig] | None = None):
|
||||
"""
|
||||
Register the agents declared in config.yaml and remember them for later rebuilds.
|
||||
|
||||
A config entry is skipped when its ``agent_name`` is already registered, so a
|
||||
database record always wins over a config entry that reuses its name and the
|
||||
registry never holds two agents under one name. Enforcing that here rather than
|
||||
in the caller keeps the guarantee independent of the order the two sources load
|
||||
in. Passing ``None`` leaves the remembered agents untouched; passing an empty
|
||||
sequence clears them.
|
||||
"""
|
||||
if agent_config is None:
|
||||
return None
|
||||
|
||||
self.config_agents = tuple(agent_config)
|
||||
|
||||
for agent_config_item in agent_config:
|
||||
if not isinstance(agent_config_item, dict):
|
||||
raise ValueError("agent_config must be a list of dictionaries")
|
||||
|
|
@ -129,6 +142,9 @@ class AgentRegistry:
|
|||
if not all([agent_name, agent_card_params]):
|
||||
continue
|
||||
|
||||
if any(agent.agent_name == agent_name for agent in self.agent_list):
|
||||
continue
|
||||
|
||||
# create a stable hash id for config item
|
||||
config_hash = self._create_agent_id(agent_config_item)
|
||||
|
||||
|
|
@ -139,26 +155,29 @@ class AgentRegistry:
|
|||
agent_config: Sequence[AgentConfig] | None = None,
|
||||
db_agents: list[dict[str, Any]] | None = None,
|
||||
):
|
||||
"""
|
||||
Rebuild the registry from the DB rows plus the agents declared in config.yaml.
|
||||
|
||||
``agent_config`` defaults to the agents remembered by the last
|
||||
``load_agents_from_config`` call, so a periodic DB reload does not drop
|
||||
config-defined agents.
|
||||
|
||||
The DB rows are registered first so that a config entry reusing one of their
|
||||
names is dropped by ``load_agents_from_config``, mirroring how config-declared
|
||||
MCP servers are unioned under the database registry. Name lookups and
|
||||
deregistration both address a single agent, so the registry must never hold two
|
||||
under one name.
|
||||
"""
|
||||
self.reset_agent_list()
|
||||
|
||||
if agent_config:
|
||||
for agent_config_item in agent_config:
|
||||
if not isinstance(agent_config_item, dict):
|
||||
raise ValueError("agent_config must be a list of dictionaries")
|
||||
|
||||
self.register_agent(
|
||||
agent_config=AgentResponse(
|
||||
agent_id=self._create_agent_id(agent_config_item),
|
||||
**agent_config_item,
|
||||
)
|
||||
) # type: ignore
|
||||
|
||||
if db_agents:
|
||||
for db_agent in db_agents:
|
||||
if not isinstance(db_agent, dict):
|
||||
raise ValueError("db_agents must be a list of dictionaries")
|
||||
|
||||
self.register_agent(agent_config=AgentResponse(**db_agent)) # type: ignore
|
||||
|
||||
self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents)
|
||||
return self.agent_list
|
||||
|
||||
###########################################################
|
||||
|
|
|
|||
|
|
@ -654,8 +654,6 @@ from fastapi.security import OAuth2PasswordBearer
|
|||
from fastapi.security.api_key import APIKeyHeader
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from litellm.types.agents import AgentConfig
|
||||
|
||||
# import enterprise folder
|
||||
enterprise_router = APIRouter()
|
||||
try:
|
||||
|
|
@ -2000,7 +1998,6 @@ config_passthrough_endpoints: Optional[List[Dict[str, Any]]] = None
|
|||
log_file = "api_log.json"
|
||||
worker_config = None
|
||||
master_key: Optional[str] = None
|
||||
config_agents: Optional[List[AgentConfig]] = None
|
||||
otel_logging = False
|
||||
prisma_client: Optional[PrismaClient] = None
|
||||
shared_aiohttp_session: Optional["ClientSession"] = None # Global shared session for connection reuse
|
||||
|
|
@ -5093,8 +5090,8 @@ class ProxyConfig:
|
|||
global_mcp_tool_registry.load_tools_from_config(mcp_tools_config, config_file_path=config_file_path)
|
||||
|
||||
## AGENTS
|
||||
agent_config = config.get("agent_list", None)
|
||||
if agent_config:
|
||||
agent_config = config.get("agents", config.get("agent_list", None))
|
||||
if agent_config is not None:
|
||||
from litellm.proxy.agent_endpoints.agent_registry import (
|
||||
global_agent_registry,
|
||||
)
|
||||
|
|
@ -6824,7 +6821,7 @@ class ProxyConfig:
|
|||
|
||||
try:
|
||||
db_agents = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client)
|
||||
AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents, agent_config=config_agents)
|
||||
AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.py::ProxyConfig:_init_agents_in_db - {}".format(str(e))
|
||||
|
|
|
|||
|
|
@ -112,3 +112,170 @@ async def test_update_agent_in_db_preserves_explicit_static_headers_and_extra_he
|
|||
|
||||
assert update_data["static_headers"] == '{"Authorization": "Bearer xyz"}'
|
||||
assert update_data["extra_headers"] == ["X-Custom-Header"]
|
||||
|
||||
|
||||
def test_load_agents_from_db_and_config_retains_previously_loaded_config_agents():
|
||||
"""
|
||||
A DB reload calls this without an explicit agent_config. It resets the
|
||||
registry, so it must fall back to the agents remembered from config.yaml
|
||||
instead of dropping them.
|
||||
"""
|
||||
registry = AgentRegistry()
|
||||
registry.load_agents_from_config(
|
||||
[
|
||||
{
|
||||
"agent_name": "config-agent",
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
registry.load_agents_from_db_and_config(
|
||||
db_agents=[
|
||||
{
|
||||
"agent_id": "db-id",
|
||||
"agent_name": "db-agent",
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert sorted(agent.agent_name for agent in registry.get_agent_list()) == [
|
||||
"config-agent",
|
||||
"db-agent",
|
||||
]
|
||||
|
||||
|
||||
def test_load_agents_from_db_and_config_skips_incomplete_config_entries():
|
||||
"""Config entries missing agent_card_params are skipped, not registered half-built."""
|
||||
registry = AgentRegistry()
|
||||
registry.load_agents_from_config([{"agent_name": "no-card"}])
|
||||
|
||||
registry.load_agents_from_db_and_config(db_agents=None)
|
||||
|
||||
assert registry.get_agent_list() == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"db_agent_names",
|
||||
[
|
||||
("shared-name", "other-db-agent"),
|
||||
("other-db-agent", "shared-name"),
|
||||
],
|
||||
ids=["colliding-db-row-first", "colliding-db-row-last"],
|
||||
)
|
||||
def test_load_agents_from_db_and_config_lets_db_rows_win_a_name_collision(db_agent_names):
|
||||
"""
|
||||
A config entry that reuses a DB agent's name must not be registered next to it.
|
||||
|
||||
Name lookups and deregistration both address a single agent, so two entries
|
||||
under one name would shadow the DB row for routing and let a single delete
|
||||
drop both. Parametrized over both DB row orders because the registry is a
|
||||
list and a first-match lookup would otherwise pass on ordering luck.
|
||||
"""
|
||||
registry = AgentRegistry()
|
||||
registry.load_agents_from_config(
|
||||
[
|
||||
{
|
||||
"agent_name": "shared-name",
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
},
|
||||
{
|
||||
"agent_name": "config-only",
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
registry.load_agents_from_db_and_config(
|
||||
db_agents=[
|
||||
{
|
||||
"agent_id": f"db-id-{name}",
|
||||
"agent_name": name,
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
}
|
||||
for name in db_agent_names
|
||||
]
|
||||
)
|
||||
|
||||
registered = registry.get_agent_list()
|
||||
assert sorted(agent.agent_name for agent in registered) == [
|
||||
"config-only",
|
||||
"other-db-agent",
|
||||
"shared-name",
|
||||
]
|
||||
|
||||
shared = registry.get_agent_by_name("shared-name")
|
||||
assert shared is not None
|
||||
assert shared.agent_id == "db-id-shared-name", "the DB row must win the shared name, not the config entry"
|
||||
|
||||
registry.deregister_agent("shared-name")
|
||||
assert [agent.agent_name for agent in registry.get_agent_list() if agent.agent_name == "shared-name"] == [], (
|
||||
"deleting the DB agent must not leave a shadowed config duplicate behind"
|
||||
)
|
||||
|
||||
|
||||
def test_load_agents_from_config_skips_a_name_already_held_by_a_db_agent():
|
||||
"""
|
||||
The one-agent-per-name rule is enforced by the loader, not by the call order.
|
||||
|
||||
A config load that lands on a registry already holding DB rows must not append a
|
||||
colliding entry, otherwise the duplicate is reachable any time the two sources are
|
||||
loaded in the other order.
|
||||
"""
|
||||
registry = AgentRegistry()
|
||||
registry.load_agents_from_db_and_config(
|
||||
db_agents=[
|
||||
{
|
||||
"agent_id": "db-id",
|
||||
"agent_name": "shared-name",
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
registry.load_agents_from_config(
|
||||
[
|
||||
{
|
||||
"agent_name": "shared-name",
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert [agent.agent_id for agent in registry.get_agent_list()] == ["db-id"]
|
||||
|
||||
|
||||
def test_load_agents_from_config_registers_one_agent_per_name_within_the_config():
|
||||
"""Two config entries sharing a name collapse to the first, keeping name lookups unambiguous."""
|
||||
registry = AgentRegistry()
|
||||
registry.load_agents_from_config(
|
||||
[
|
||||
{"agent_name": "dupe", "agent_card_params": _sample_agent_card_params()},
|
||||
{"agent_name": "dupe", "agent_card_params": {**_sample_agent_card_params(), "url": "http://second"}},
|
||||
]
|
||||
)
|
||||
|
||||
registered = registry.get_agent_list()
|
||||
assert len(registered) == 1
|
||||
assert registered[0].agent_card_params["url"] == "http://localhost"
|
||||
|
||||
|
||||
def test_load_agents_from_config_with_an_empty_list_clears_the_remembered_agents():
|
||||
"""
|
||||
An explicitly empty config must forget the previously loaded agents.
|
||||
|
||||
Otherwise the next DB rebuild replays agents the operator removed from config.yaml.
|
||||
``None`` still means "no opinion" and leaves them alone.
|
||||
"""
|
||||
registry = AgentRegistry()
|
||||
registry.load_agents_from_config(
|
||||
[{"agent_name": "removed-agent", "agent_card_params": _sample_agent_card_params()}]
|
||||
)
|
||||
assert registry.config_agents != ()
|
||||
|
||||
registry.load_agents_from_config([])
|
||||
assert registry.config_agents == ()
|
||||
|
||||
registry.load_agents_from_db_and_config(db_agents=None)
|
||||
assert registry.get_agent_list() == [], "a removed config agent must not come back on the next rebuild"
|
||||
|
|
|
|||
|
|
@ -2354,3 +2354,133 @@ async def test_ProxyConfig_load_config_redacts_secret_litellm_setting_keeps_plai
|
|||
assert "num_retries=7" in rendered, (
|
||||
f"non-secret num_retries value was over-redacted; expected it visible in {rendered!r}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ProxyConfig agents from config.yaml
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clean_agent_registry():
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
|
||||
original_agents = list(global_agent_registry.agent_list)
|
||||
original_config_agents = getattr(global_agent_registry, "config_agents", ())
|
||||
global_agent_registry.agent_list = []
|
||||
global_agent_registry.config_agents = ()
|
||||
try:
|
||||
yield global_agent_registry
|
||||
finally:
|
||||
global_agent_registry.agent_list = original_agents
|
||||
global_agent_registry.config_agents = original_config_agents
|
||||
|
||||
|
||||
def _config_agent(agent_name: str) -> Dict[str, Any]:
|
||||
return {
|
||||
"agent_name": agent_name,
|
||||
"agent_card_params": {
|
||||
"name": "Config Agent",
|
||||
"url": "http://localhost:10001",
|
||||
"protocolVersion": "1.0",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class _FakeAgentRow:
|
||||
"""Stand-in for a prisma agent record: supports dict() and .object_permission."""
|
||||
|
||||
def __init__(self, agent_id: str, agent_name: str) -> None:
|
||||
self.agent_id = agent_id
|
||||
self.agent_name = agent_name
|
||||
self.object_permission = None
|
||||
self.spend = 0.0
|
||||
|
||||
def __iter__(self):
|
||||
return iter(
|
||||
{
|
||||
"agent_id": self.agent_id,
|
||||
"agent_name": self.agent_name,
|
||||
"agent_card_params": {"name": self.agent_name, "url": "http://db-agent"},
|
||||
"litellm_params": {},
|
||||
}.items()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("config_key", ["agents", "agent_list"])
|
||||
async def test_ProxyConfig__init_non_llm_configs_registers_agents_from_config(clean_agent_registry, config_key):
|
||||
"""The documented ``agents:`` key must register agents, as must the legacy ``agent_list:``."""
|
||||
await ProxyConfig()._init_non_llm_configs(
|
||||
config={config_key: [_config_agent("config-agent")]},
|
||||
config_file_path=None,
|
||||
)
|
||||
|
||||
assert [agent.agent_name for agent in clean_agent_registry.get_agent_list()] == ["config-agent"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_agent_registry):
|
||||
"""A DB reload rebuilds the registry; config-defined agents must survive it alongside DB rows."""
|
||||
await ProxyConfig()._init_non_llm_configs(
|
||||
config={"agents": [_config_agent("config-agent")]},
|
||||
config_file_path=None,
|
||||
)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_agentstable.find_many = AsyncMock(return_value=[_FakeAgentRow("db-id", "db-agent")])
|
||||
|
||||
await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client)
|
||||
|
||||
assert sorted(agent.agent_name for agent in clean_agent_registry.get_agent_list()) == [
|
||||
"config-agent",
|
||||
"db-agent",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"config, expected_agent_names",
|
||||
[
|
||||
({"agents": [], "agent_list": [_config_agent("legacy-agent")]}, []),
|
||||
(
|
||||
{
|
||||
"agents": [_config_agent("documented-agent")],
|
||||
"agent_list": [_config_agent("legacy-agent")],
|
||||
},
|
||||
["documented-agent"],
|
||||
),
|
||||
({"agent_list": [_config_agent("legacy-agent")]}, ["legacy-agent"]),
|
||||
],
|
||||
ids=["empty-agents-wins", "populated-agents-wins", "agent_list-alone-still-works"],
|
||||
)
|
||||
async def test_ProxyConfig__init_non_llm_configs_prefers_agents_key_by_presence(
|
||||
clean_agent_registry, config, expected_agent_names
|
||||
):
|
||||
"""
|
||||
``agents`` outranks the legacy ``agent_list`` whenever the key is present.
|
||||
|
||||
Selecting on truthiness instead would silently register the legacy entries
|
||||
for a config that spells out ``agents: []``.
|
||||
"""
|
||||
await ProxyConfig()._init_non_llm_configs(config=config, config_file_path=None)
|
||||
|
||||
assert [agent.agent_name for agent in clean_agent_registry.get_agent_list()] == expected_agent_names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig__init_non_llm_configs_empty_agents_key_clears_remembered_agents(clean_agent_registry):
|
||||
"""
|
||||
An explicitly empty ``agents:`` must reach the registry, not be skipped as falsy.
|
||||
|
||||
Skipping it leaves the previously remembered agents in place, so the next DB
|
||||
rebuild replays agents the operator deleted from config.yaml.
|
||||
"""
|
||||
clean_agent_registry.load_agents_from_config([_config_agent("stale-agent")])
|
||||
assert clean_agent_registry.config_agents != ()
|
||||
|
||||
await ProxyConfig()._init_non_llm_configs(config={"agents": []}, config_file_path=None)
|
||||
|
||||
assert clean_agent_registry.config_agents == ()
|
||||
clean_agent_registry.load_agents_from_db_and_config(db_agents=None)
|
||||
assert clean_agent_registry.get_agent_list() == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue