mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
refactor(agents): resolve attached access groups from the agent registry instead of the DB on the request path
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
18c31e6fc6
commit
1206fa802b
4 changed files with 14 additions and 166 deletions
|
|
@ -1,35 +1,18 @@
|
|||
import asyncio
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Protocol, TypeAlias
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import LiteLLM_AccessGroupTable
|
||||
from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl
|
||||
|
||||
|
||||
class _AgentAccessGroupsRecord(Protocol):
|
||||
@property
|
||||
def access_group_ids(self) -> Sequence[str] | None: ...
|
||||
|
||||
|
||||
class _AgentIdWhere(TypedDict):
|
||||
agent_id: ReadOnly[str]
|
||||
|
||||
|
||||
AccessGroupIds: TypeAlias = tuple[str, ...]
|
||||
AccessGroupIdsLoader: TypeAlias = Callable[[str], Awaitable[AccessGroupIds]] # mutable-ok: Callable params
|
||||
AgentRecordFinder: TypeAlias = Callable[[str], Awaitable[_AgentAccessGroupsRecord | None]] # mutable-ok: Callable
|
||||
LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None
|
||||
AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LoadedAccessGroup]] # mutable-ok: Callable parameter syntax
|
||||
|
||||
_CACHED_IDS: Final = TypeAdapter(list[str])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AgentAccessGroupCeiling:
|
||||
|
|
@ -44,19 +27,6 @@ class AgentAccessGroupCeiling:
|
|||
CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] # mutable-ok: Callable params
|
||||
|
||||
|
||||
def agent_access_group_ids_cache_key(agent_id: str) -> str:
|
||||
return f"agent_access_group_ids:{agent_id}"
|
||||
|
||||
|
||||
def _cached_access_group_ids(cached: object) -> AccessGroupIds | None:
|
||||
if cached is None:
|
||||
return None
|
||||
try:
|
||||
return tuple(_CACHED_IDS.validate_python(cached))
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds:
|
||||
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
|
||||
|
||||
|
|
@ -64,51 +34,6 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds:
|
|||
return tuple(agent.access_group_ids or ()) if agent is not None else ()
|
||||
|
||||
|
||||
async def load_agent_access_group_ids(
|
||||
agent_id: str,
|
||||
cache: DualCache,
|
||||
find_agent: AgentRecordFinder,
|
||||
fallback: AccessGroupIdsLoader,
|
||||
) -> AccessGroupIds:
|
||||
"""The agent row's groups, cached for the management-object TTL and evicted on every agent write."""
|
||||
cache_key: Final = agent_access_group_ids_cache_key(agent_id)
|
||||
cached: Final = _cached_access_group_ids(await cache.async_get_cache(key=cache_key))
|
||||
if cached is not None:
|
||||
return cached
|
||||
try:
|
||||
record: Final = await find_agent(agent_id)
|
||||
except Exception as e: # noqa: BLE001 # prisma raises many error types; the registry snapshot answers instead
|
||||
verbose_proxy_logger.warning("Failed to read access groups for agent %r, using registry: %s", agent_id, e)
|
||||
return await fallback(agent_id)
|
||||
access_group_ids: Final = tuple(record.access_group_ids or ()) if record is not None else ()
|
||||
await cache.async_set_cache(key=cache_key, value=access_group_ids, ttl=get_management_object_ttl(cache))
|
||||
return access_group_ids
|
||||
|
||||
|
||||
async def _load_agent_access_group_ids(agent_id: str) -> AccessGroupIds:
|
||||
from litellm.proxy.agent_endpoints.agent_registry import agents_table
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
return await _registry_access_group_ids(agent_id)
|
||||
db: Final = prisma_client
|
||||
|
||||
async def find_agent(row_agent_id: str) -> _AgentAccessGroupsRecord | None:
|
||||
return await agents_table(db).find_unique(where=_AgentIdWhere(agent_id=row_agent_id))
|
||||
|
||||
return await load_agent_access_group_ids(agent_id, user_api_key_cache, find_agent, _registry_access_group_ids)
|
||||
|
||||
|
||||
async def evict_agent_access_group_ids(agent_ids: Sequence[str]) -> None:
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
await evict_and_broadcast(
|
||||
cache_keys=tuple(agent_access_group_ids_cache_key(agent_id) for agent_id in agent_ids),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
|
||||
async def _load_access_group(access_group_id: str) -> LoadedAccessGroup:
|
||||
from litellm.proxy.auth.auth_checks import get_access_object
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
|
@ -132,7 +57,7 @@ async def _load_access_group(access_group_id: str) -> LoadedAccessGroup:
|
|||
|
||||
async def resolve_agent_access_group_ceiling(
|
||||
agent_id: str,
|
||||
load_access_group_ids: AccessGroupIdsLoader = _load_agent_access_group_ids,
|
||||
load_access_group_ids: AccessGroupIdsLoader = _registry_access_group_ids,
|
||||
load_access_group: AccessGroupLoader = _load_access_group,
|
||||
) -> AgentAccessGroupCeiling | None:
|
||||
"""``None`` when the agent has no access groups attached, so nothing is capped."""
|
||||
|
|
|
|||
|
|
@ -44,7 +44,6 @@ from litellm.proxy.agent_endpoints.agent_search import (
|
|||
global_agent_search_index,
|
||||
search_agents,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import evict_agent_access_group_ids
|
||||
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
|
||||
|
|
@ -697,7 +696,6 @@ async def update_agent(
|
|||
prisma_client=prisma_client,
|
||||
updated_by=updated_by,
|
||||
)
|
||||
await evict_agent_access_group_ids((agent_id,))
|
||||
|
||||
# deregister in memory
|
||||
AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name"))
|
||||
|
|
@ -801,7 +799,6 @@ async def patch_agent(
|
|||
prisma_client=prisma_client,
|
||||
updated_by=updated_by,
|
||||
)
|
||||
await evict_agent_access_group_ids((agent_id,))
|
||||
|
||||
# deregister in memory
|
||||
AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name"))
|
||||
|
|
@ -864,7 +861,6 @@ async def delete_agent(
|
|||
raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found in DB.")
|
||||
|
||||
await AGENT_REGISTRY.delete_agent_from_db(agent_id=agent_id, prisma_client=prisma_client)
|
||||
await evict_agent_access_group_ids((agent_id,))
|
||||
|
||||
AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name"))
|
||||
|
||||
|
|
|
|||
|
|
@ -16,7 +16,6 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import evict_agent_access_group_ids
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
_cache_access_object,
|
||||
_cache_key_object,
|
||||
|
|
@ -783,7 +782,6 @@ async def delete_access_group(
|
|||
|
||||
await invalidate_access_group_cache(access_group_id)
|
||||
_detach_access_group_from_agent_registry(detached_agent_ids, access_group_id)
|
||||
await evict_agent_access_group_ids(detached_agent_ids)
|
||||
await _patch_team_caches_remove_access_group(
|
||||
affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,16 +1,11 @@
|
|||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.models.access_group import LiteLLM_AccessGroupTable
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
|
||||
AgentAccessGroupCeiling,
|
||||
agent_access_group_ids_cache_key,
|
||||
load_agent_access_group_ids,
|
||||
resolve_agent_access_group_ceiling,
|
||||
)
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
|
@ -110,86 +105,20 @@ async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted():
|
|||
assert ceiling.agent_ids == frozenset()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _AgentRow:
|
||||
access_group_ids: Sequence[str] | None
|
||||
|
||||
|
||||
class _FakeAgentTable:
|
||||
def __init__(self, rows: dict[str, _AgentRow], failing: bool = False) -> None:
|
||||
self._rows: Final = rows
|
||||
self._failing: Final = failing
|
||||
self.reads = 0
|
||||
|
||||
async def find_agent(self, agent_id: str) -> _AgentRow | None:
|
||||
self.reads += 1
|
||||
if self._failing:
|
||||
raise RuntimeError("db down")
|
||||
return self._rows.get(agent_id)
|
||||
|
||||
|
||||
async def _registry_snapshot(agent_id: str) -> tuple[str, ...]:
|
||||
return ("registry-group",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_row_is_read_once_then_served_from_cache():
|
||||
cache: Final = DualCache()
|
||||
table: Final = _FakeAgentTable({"agent-1": _AgentRow(["g1", "g2"])})
|
||||
async def test_default_agent_loader_reads_the_attached_groups_from_the_registry():
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
|
||||
first: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot)
|
||||
second: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot)
|
||||
_, load_group = _loaders(None, {"g1": _group("g1", models=("gpt-5",))})
|
||||
global_agent_registry.register_agent(_agent(["g1"]))
|
||||
try:
|
||||
ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_access_group=load_group)
|
||||
finally:
|
||||
global_agent_registry.deregister_agent("agent")
|
||||
|
||||
assert (first, second, table.reads) == (("g1", "g2"), ("g1", "g2"), 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_with_no_row_or_no_groups_caches_an_empty_answer():
|
||||
cache: Final = DualCache()
|
||||
table: Final = _FakeAgentTable({"bare": _AgentRow(None)})
|
||||
|
||||
bare: Final = await load_agent_access_group_ids("bare", cache, table.find_agent, _registry_snapshot)
|
||||
missing: Final = await load_agent_access_group_ids("missing", cache, table.find_agent, _registry_snapshot)
|
||||
again: Final = await load_agent_access_group_ids("missing", cache, table.find_agent, _registry_snapshot)
|
||||
|
||||
assert (bare, missing, again, table.reads) == ((), (), (), 2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_evicted_cache_entry_picks_up_the_patched_row():
|
||||
cache: Final = DualCache()
|
||||
rows: Final = {"agent-1": _AgentRow(["g1"])}
|
||||
table: Final = _FakeAgentTable(rows)
|
||||
await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot)
|
||||
|
||||
rows["agent-1"] = _AgentRow(["g2"])
|
||||
stale: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot)
|
||||
await cache.async_delete_cache(key=agent_access_group_ids_cache_key("agent-1"))
|
||||
fresh: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot)
|
||||
|
||||
assert (stale, fresh) == (("g1",), ("g2",))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unreadable_row_falls_back_to_the_registry_without_caching():
|
||||
cache: Final = DualCache()
|
||||
table: Final = _FakeAgentTable({}, failing=True)
|
||||
|
||||
answer: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot)
|
||||
|
||||
assert answer == ("registry-group",)
|
||||
assert await cache.async_get_cache(key=agent_access_group_ids_cache_key("agent-1")) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_garbage_in_the_cache_is_treated_as_a_miss():
|
||||
cache: Final = DualCache()
|
||||
await cache.async_set_cache(key=agent_access_group_ids_cache_key("agent-1"), value={"not": "a list"})
|
||||
table: Final = _FakeAgentTable({"agent-1": _AgentRow(["g1"])})
|
||||
|
||||
answer: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot)
|
||||
|
||||
assert (answer, table.reads) == (("g1",), 1)
|
||||
assert ceiling == AgentAccessGroupCeiling(
|
||||
access_group_ids=("g1",), models=frozenset({"gpt-5"}), mcp_server_ids=frozenset(), agent_ids=frozenset()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue