From 8b53d48e34eef703654cc69e3bfbf378555664b0 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 14 Sep 2026 17:30:11 -0700 Subject: [PATCH] fix(memory): show enrolled users and verify live team recall --- litellm/proxy/memory/management.py | 39 +++++++++++++++---- litellm/proxy/memory/store.py | 11 +++--- litellm/types/memory_v2.py | 5 +++ tests/e2e/management/test_memory_v2_e2e.py | 21 +++++++++- .../proxy/memory/test_memory_v2_management.py | 7 +++- .../memory/_components/MemorySettings.tsx | 7 ++-- .../memory/page.integration.test.tsx | 14 ++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 26 ++++++++++++- 8 files changed, 107 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/memory/management.py b/litellm/proxy/memory/management.py index 49015bfb324..696ab15710e 100644 --- a/litellm/proxy/memory/management.py +++ b/litellm/proxy/memory/management.py @@ -19,7 +19,15 @@ from litellm.proxy.memory.store import MemoryStore from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository -from litellm.types.memory_v2 import MemoryCapture, MemoryEntry, MemoryQuery, MemorySearch, MemorySettings, MemoryStatus +from litellm.types.memory_v2 import ( + MemoryCapture, + MemoryEntry, + MemoryQuery, + MemorySearch, + MemorySettings, + MemorySettingsView, + MemoryStatus, +) _AUTH: Final = Depends(user_api_key_auth) router: Final = APIRouter(prefix="/v2/memory", tags=["memory management"]) # mutable-ok: FastAPI requires native tags. @@ -30,14 +38,31 @@ def require_memory_admin(auth: UserAPIKeyAuth, *, write: bool = False) -> None: raise HTTPException(status_code=403, detail="Only proxy administrators can configure gateway memory") -@router.get("/settings", response_model=MemorySettings) -async def get_settings(auth: UserAPIKeyAuth = _AUTH) -> MemorySettings: +async def settings_view(settings: MemorySettings) -> MemorySettingsView: + users: Final = ( + await UserRepository(memory_primary_client(require_memory_prisma())).table.find_many( + where={"user_id": {"in": list(settings.user_ids)}}, # mutable-ok: Prisma requires native JSON. + take=len(settings.user_ids), + ) + if settings.user_ids + else () + ) + return MemorySettingsView( + **settings.model_dump(), + user_names=MappingProxyType( + {user.user_id: user.user_alias or user.user_email or user.user_id for user in users} + ), + ) + + +@router.get("/settings", response_model=MemorySettingsView) +async def get_settings(auth: UserAPIKeyAuth = _AUTH) -> MemorySettingsView: require_memory_admin(auth) - return await memory_settings(require_memory_prisma()) + return await settings_view(await memory_settings(require_memory_prisma())) -@router.put("/settings", response_model=MemorySettings) -async def set_settings(settings: MemorySettings, auth: UserAPIKeyAuth = _AUTH) -> MemorySettings: +@router.put("/settings", response_model=MemorySettingsView) +async def set_settings(settings: MemorySettings, auth: UserAPIKeyAuth = _AUTH) -> MemorySettingsView: require_memory_admin(auth, write=True) prisma: Final = memory_primary_client(require_memory_prisma()) selected: Final = tuple(sorted(frozenset(settings.user_ids))) if not settings.everyone else () @@ -53,7 +78,7 @@ async def set_settings(settings: MemorySettings, auth: UserAPIKeyAuth = _AUTH) - saved: Final = settings.model_copy(update=MappingProxyType({"user_ids": selected})) await ConfigRepository(prisma).set_param(MEMORY_CONFIG_PARAM, saved.model_dump(mode="json")) await invalidate_memory_configuration() - return saved + return await settings_view(saved) async def memory_store(auth: UserAPIKeyAuth) -> MemoryStore: diff --git a/litellm/proxy/memory/store.py b/litellm/proxy/memory/store.py index db8159ffb93..a915802434f 100644 --- a/litellm/proxy/memory/store.py +++ b/litellm/proxy/memory/store.py @@ -187,14 +187,13 @@ class MemoryStore: {"user_id": user_id} if user_id else {}, # mutable-ok: Prisma requires native JSON. _before(before), ) - result: Final[tuple[MemoryEntry, ...]] if not search.query.strip(): - result = await self._page(where, limit=search.limit, offset=search.offset) - else: - ranked, _ = await self._ranked(search.query, where, search.offset + search.limit, recent_first=recent_first) - result = tuple(entry for entry, _, _ in ranked[search.offset :]) + page: Final = await self._page(where, limit=search.limit, offset=search.offset) + await self.authorize(require_active=require_active) + return page + ranked, _ = await self._ranked(search.query, where, search.offset + search.limit, recent_first=recent_first) await self.authorize(require_active=require_active) - return result + return tuple(entry for entry, _, _ in ranked[search.offset :]) async def read(self, memory_id: str, *, require_active: bool = True) -> MemoryEntry: access: Final = await self.authorize(require_active=require_active) diff --git a/litellm/types/memory_v2.py b/litellm/types/memory_v2.py index 03c17760b8e..fc76a6947f8 100644 --- a/litellm/types/memory_v2.py +++ b/litellm/types/memory_v2.py @@ -1,4 +1,5 @@ import re +from collections.abc import Mapping from datetime import datetime from typing import Annotated, Literal, TypeAlias @@ -27,6 +28,10 @@ class MemorySettings(BaseModel): user_ids: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=10000) +class MemorySettingsView(MemorySettings): + user_names: Mapping[str, str] + + class MemoryStatus(BaseModel): model_config = ConfigDict(frozen=True) diff --git a/tests/e2e/management/test_memory_v2_e2e.py b/tests/e2e/management/test_memory_v2_e2e.py index 1316cabe591..6e5c2c1228e 100644 --- a/tests/e2e/management/test_memory_v2_e2e.py +++ b/tests/e2e/management/test_memory_v2_e2e.py @@ -276,9 +276,10 @@ class TestMemoryV2: @pytest.mark.covers("mgmt.memory_v2.entries.team_permissions") def test_delegated_team_reads_allow_recall_but_not_edit_and_can_be_revoked( - self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects + self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels ) -> None: - saved = unwrap(memory.capture(subjects.owner, _fact(unique_marker()))) + marker = f"team-{unique_marker()}" + saved = unwrap(memory.capture(subjects.owner, _fact(marker))) assert memory.entries(subjects.outsider) == [] for permissions, visible in ((["/spend/logs"], False), (["/v2/memory/entries"], True), ([], False)): unwrap( @@ -296,6 +297,22 @@ class TestMemoryV2: assert unwrap(memory.read(subjects.outsider, saved.memory_id)).content == saved.content _assert_denied(memory.update(subjects.outsider, saved.memory_id, _fact(unique_marker()))) _assert_denied(memory.delete_entry(subjects.outsider, saved.memory_id)) + recalled = unwrap( + client.proxy.chat( + subjects.outsider, + ChatBody( + model=memory_models.chat, + max_tokens=1200, + messages=[ + ChatMessage( + role="user", + content="Search the team's memories for the demo project codename and return it exactly.", + ) + ], + ), + ) + ) + assert marker in recalled.model_dump_json() else: _assert_denied(memory.read(subjects.outsider, saved.memory_id)) diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_management.py b/tests/test_litellm/proxy/memory/test_memory_v2_management.py index 72126d12643..cedd9f8cdbd 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_management.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_management.py @@ -108,13 +108,16 @@ def row(**changes: object) -> LiteLLM_MemoryTable: @pytest.mark.asyncio async def test_default_off_and_proxy_admin_can_enable_selected_users(database: MagicMock) -> None: admin = auth("admin", LitellmUserRoles.PROXY_ADMIN) - assert (await management.get_settings(admin)) == MemorySettings() + assert not (await management.get_settings(admin)).enabled assert not (await management.get_status(auth())).active - database.db.litellm_usertable.find_many.return_value = [SimpleNamespace(user_id="owner")] + database.db.litellm_usertable.find_many.return_value = [ + SimpleNamespace(user_id="owner", user_alias="Alex", user_email="alex@example.test") + ] saved = await management.set_settings( MemorySettings(enabled=True, everyone=False, user_ids=("owner", "owner")), admin ) assert saved.user_ids == ("owner",) + assert saved.user_names == {"owner": "Alex"} written = database.db.litellm_config.upsert.call_args.kwargs["data"]["update"]["param_value"] database.db.litellm_config.find_unique.return_value = SimpleNamespace(param_value=written) assert (await management.get_status(auth())).active diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemorySettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemorySettings.tsx index c58e75cd367..ff837c0b8ee 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemorySettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemorySettings.tsx @@ -28,7 +28,8 @@ export function MemoryAdministration({ }); const current = draft ?? settings.data; const save = useMutation({ - mutationFn: (body: Settings) => fetchClient.PUT("/v2/memory/settings", { body }), + mutationFn: ({ enabled, everyone, user_ids }: Settings) => + fetchClient.PUT("/v2/memory/settings", { body: { enabled, everyone, user_ids } }), onSuccess: async ({ data }) => { cache.setQueryData(["memorySettings", userId], data); setDraft(null); @@ -98,12 +99,12 @@ export function MemoryAdministration({