From f4fc27f2e2b2f876d1e59fece9397e40307b8271 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 14 Sep 2026 13:23:31 -0700 Subject: [PATCH] feat(memory): add user activation and contributor table --- litellm/proxy/memory/management.py | 30 ++- litellm/proxy/memory/memory_endpoints.py | 4 +- litellm/proxy/memory/policy.py | 1 + litellm/proxy/memory/store.py | 7 +- litellm/types/memory_v2.py | 5 +- .../proxy/memory/test_memory_endpoints.py | 24 ++- .../proxy/memory/test_memory_v2_management.py | 78 ++++++- .../_components/AutomaticMemoryEntries.tsx | 197 ++++++------------ .../memory/_components/MemoryEntriesTable.tsx | 180 ++++++++++++++++ .../memory/_components/MemorySettings.tsx | 31 +-- .../memory/page.integration.test.tsx | 68 +++++- .../src/app/(dashboard)/memory/page.tsx | 52 +++-- ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 +- 13 files changed, 497 insertions(+), 190 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryEntriesTable.tsx diff --git a/litellm/proxy/memory/management.py b/litellm/proxy/memory/management.py index d0974f1f543..6f1ce72671a 100644 --- a/litellm/proxy/memory/management.py +++ b/litellm/proxy/memory/management.py @@ -1,3 +1,4 @@ +from types import MappingProxyType from typing import Annotated, Final from fastapi import APIRouter, Depends, HTTPException, Query, Response @@ -225,7 +226,15 @@ async def get_status( key_id: str | None = Query(None, pattern=r"^[a-f0-9]{64}$"), auth: UserAPIKeyAuth = _AUTH, ) -> MemoryStatus: - return (await access_for_key(auth, key_id)).status + access: Final = await access_for_key(auth, key_id) + user: Final = ( + await UserRepository(memory_primary_client(require_memory_prisma())).find_by_id(access.identity.user_id) + if access.identity.user_id + else None + ) + return access.status.model_copy( + update=MappingProxyType({"user_name": user.user_alias or user.user_email or user.user_id if user else None}) + ) async def access_for_key(auth: UserAPIKeyAuth, key_id: str | None) -> MemoryAccess: @@ -267,12 +276,25 @@ async def list_entries( raise HTTPException(status_code=422, detail="Provide both memory cursor fields") prisma: Final = memory_primary_client(require_memory_prisma()) access: Final = await access_for_key(auth, key_id) - return await MemoryStore(prisma, access).search( + entries: Final = await MemoryStore(prisma, access).search( MemorySearch(query=query, limit=limit, offset=offset), require_active=False, recent_first=True, before=(before_updated_at, before_memory_id) if before_updated_at and before_memory_id else None, ) + actors: Final = tuple(frozenset(entry.actor for entry in entries if entry.actor)) + if not actors: + return entries + users: Final = await UserRepository(prisma).table.find_many( + where={"user_id": {"in": list(actors)}}, # mutable-ok: Prisma serializes the bounded page's contributor IDs. + take=len(actors), + ) + names: Final = MappingProxyType( + {user.user_id: user.user_alias or user.user_email or user.user_id for user in users} + ) + return tuple( + entry.model_copy(update=MappingProxyType({"actor_name": names.get(entry.actor or "")})) for entry in entries + ) @router.post("/entries", response_model=MemoryEntry) @@ -283,7 +305,9 @@ async def capture_entry( ) -> MemoryEntry: prisma: Final = memory_primary_client(require_memory_prisma()) access: Final = await access_for_key(auth, key_id) - return await MemoryStore(prisma, access).capture(capture) + return await MemoryStore(prisma, access, actor=auth.user_id or MemoryIdentity.from_auth(auth).key_id).capture( + capture + ) @router.delete("/entries/{memory_id}", status_code=204) diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index ea4648dfcb2..69bd9c41fc7 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -82,10 +82,10 @@ class _LegacyMemoryVisibility(TypedDict): def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object] | None: """ Prisma `where` fragment restricting rows to those the caller can see. - Returns None for admins (no restriction). + Administrators can access every V1 row, independently of user/team ownership. """ if user_api_key_has_admin_view(user_api_key_dict): - return None + return {"namespace": None} # mutable-ok: Prisma requires a native JSON filter for V1 visibility. ors: Final = [ {field: value} for field, value in (("user_id", user_api_key_dict.user_id), ("team_id", user_api_key_dict.team_id)) diff --git a/litellm/proxy/memory/policy.py b/litellm/proxy/memory/policy.py index 12d4a7f3895..31a79f06bc1 100644 --- a/litellm/proxy/memory/policy.py +++ b/litellm/proxy/memory/policy.py @@ -172,6 +172,7 @@ class MemoryAccess: scope=self.policy.scope if self.policy and self.namespace else None, opted_in=self.opted_in, policy_id=self.policy.policy_id if self.policy else None, + user_id=self.identity.user_id, ) diff --git a/litellm/proxy/memory/store.py b/litellm/proxy/memory/store.py index 2295e591c1b..e0f897e8fff 100644 --- a/litellm/proxy/memory/store.py +++ b/litellm/proxy/memory/store.py @@ -50,9 +50,10 @@ def memory_entry(row: "LiteLLM_MemoryTable") -> MemoryEntry: class MemoryStore: - def __init__(self, prisma_client: object, access: MemoryAccess) -> None: + def __init__(self, prisma_client: object, access: MemoryAccess, *, actor: str | None = None) -> None: self.prisma_client = memory_primary_client(prisma_client) self.access = access + self.actor = actor or access.identity.user_id or access.identity.key_id self.table = MemoryRepository(self.prisma_client).table async def authorize_namespace(self, *, write: bool = False, require_active: bool = True) -> str: @@ -159,7 +160,7 @@ class MemoryStore: data: Final = { # mutable-ok: Prisma query and write JSON. "value": content, "metadata": json.dumps(metadata), - "updated_by": self.access.identity.user_id or self.access.identity.key_id, + "updated_by": self.actor, } existing: Final = await table.find_unique( where={ # mutable-ok: Prisma query and write JSON. @@ -215,7 +216,7 @@ class MemoryStore: "namespace": namespace, "user_id": self.access.identity.user_id, "team_id": self.access.identity.team_id, - "created_by": self.access.identity.user_id or self.access.identity.key_id, + "created_by": self.actor, } ) return memory_entry(created) diff --git a/litellm/types/memory_v2.py b/litellm/types/memory_v2.py index c36242be2bc..c8ed56f0f97 100644 --- a/litellm/types/memory_v2.py +++ b/litellm/types/memory_v2.py @@ -29,7 +29,7 @@ class MemoryPolicyInput(BaseModel): target_type: MemoryTarget target_id: str = Field(min_length=1, max_length=256) activation: MemoryActivation - scope: MemoryScope = "key" + scope: MemoryScope = "user" @model_validator(mode="after") def validate_target(self) -> Self: @@ -62,6 +62,8 @@ class MemoryStatus(BaseModel): scope: MemoryScope | None opted_in: bool policy_id: str | None + user_id: str | None = None + user_name: str | None = None class MemoryCapture(BaseModel): @@ -90,6 +92,7 @@ class MemoryEntry(BaseModel): updated_at: datetime created_at: datetime | None = None actor: str | None = None + actor_name: str | None = None when_to_use: str = "" scope: str = "" kind: MemoryKind = "context" diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/test_litellm/proxy/memory/test_memory_endpoints.py index 4a4a3551d21..ee1550b2904 100644 --- a/tests/test_litellm/proxy/memory/test_memory_endpoints.py +++ b/tests/test_litellm/proxy/memory/test_memory_endpoints.py @@ -1088,14 +1088,30 @@ class TestMemoryEndpoints: assert resp.status_code == 404, resp.text assert resp.json()["detail"] == "Memory with key 'notes' not found" - def test_visibility_filter_unscoped_for_admin_viewer(self): + def test_visibility_filter_excludes_v2_for_admin_viewer(self): """ - proxy_admin_viewer reads with the same unscoped filter as proxy_admin; - every other role stays row-restricted. + Administrative reads include every V1 owner but exclude V2 namespaces. """ - assert _visibility_filter(_admin_viewer_auth()) is None + assert _visibility_filter(_admin_viewer_auth()) == {"namespace": None} assert _visibility_filter(_user_auth("user-a", "team-a")) is not None + def test_v1_admin_crud_does_not_read_or_change_v2_memories(self): + table = self.prisma.db.litellm_memorytable + automatic = _make_row(memory_id="automatic", key="memory-v2:private:note", value="Private context") + automatic.namespace = "private" + table.rows.extend([_make_row(key="original-v1", value="Existing integration"), automatic]) + client = _make_client(_admin_auth()) + with _patch_prisma(self.prisma): + listed = client.get("/v1/memory") + assert listed.status_code == 200 + assert [row["key"] for row in listed.json()["memories"]] == ["original-v1"] + assert client.get("/v1/memory/memory-v2:private:note").status_code == 404 + assert client.delete("/v1/memory/memory-v2:private:note").status_code == 404 + assert client.put("/v1/memory/memory-v2:private:note", json={"value": "Overwrite"}).status_code == 409 + updated = client.put("/v1/memory/original-v1", json={"value": "Still works", "metadata": None}) + assert updated.status_code == 200 and updated.json()["value"] == "Still works" + assert automatic.value == "Private context" + def test_list_memory_admin_viewer_sees_all(self): """Read parity end-to-end: the viewer's own user_id/team_id must not filter the list.""" table = self.prisma.db.litellm_memorytable 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 73b908dd92d..654793fbc65 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_management.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_management.py @@ -12,7 +12,7 @@ from fastapi.testclient import TestClient from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.memory import management -from litellm.proxy.memory.policy import MemoryIdentity, memory_digest +from litellm.proxy.memory.policy import MemoryIdentity, memory_digest, resolve_memory_access from litellm.types.memory_v2 import MemoryCapture, MemoryPolicy, MemoryPolicyInput, MemoryPreference @@ -96,6 +96,82 @@ def policy(**changes: object) -> MemoryPolicy: ) +@pytest.mark.asyncio +async def test_user_preference_and_memories_follow_owner_across_keys(database: MagicMock) -> None: + database.db.litellm_memorypolicy.find_many.return_value = [policy(scope="user", activation="opt_in")] + first = MemoryIdentity.from_auth(auth()) + second = MemoryIdentity.from_auth(auth().model_copy(update={"token": "b" * 64})) + other = MemoryIdentity.from_auth(auth("other").model_copy(update={"token": "c" * 64})) + other_org = MemoryIdentity.from_auth(auth().model_copy(update={"org_id": "other-org"})) + assert not (await resolve_memory_access(database, first)).active + await management.set_preference(MemoryPreference(enabled=True), auth()) + written = database.db.litellm_memorypreference.upsert.call_args.kwargs["data"]["create"] + + async def preference(*, where: dict[str, str]) -> SimpleNamespace | None: + return SimpleNamespace(enabled=written["enabled"]) if where["subject"] == written["subject"] else None + + database.db.litellm_memorypreference.find_unique.side_effect = preference + one = await resolve_memory_access(database, first) + two = await resolve_memory_access(database, second) + assert one.active and two.active and one.namespace == two.namespace + assert not (await resolve_memory_access(database, other)).active + assert one.namespace != (await resolve_memory_access(database, other_org)).namespace + database.db.litellm_memorypolicy.find_many.return_value = [policy(scope="user", activation="disabled")] + assert not (await resolve_memory_access(database, second)).active + + +@pytest.mark.asyncio +async def test_dashboard_resolves_saved_contributor_names_and_identifies_preference_owner(database: MagicMock) -> None: + database.db.litellm_memorypolicy.find_many.return_value = [policy(scope="team")] + database.db.litellm_usertable.find_unique.return_value = {"user_id": "owner", "user_alias": "Alex Rivera"} + status = await management.get_status(None, auth()) + assert status.user_id == "owner" and status.user_name == "Alex Rivera" + now = datetime(2026, 9, 14, tzinfo=timezone.utc) + database.db.litellm_memorytable.find_many.return_value = [ + SimpleNamespace( + memory_id="entry", + key="memory-v2:namespace:demo", + value="Use port 8123", + metadata={}, + updated_at=now, + created_at=now, + created_by="contributor", + ) + ] + database.db.litellm_usertable.find_many.return_value = [ + SimpleNamespace(user_id="contributor", user_alias="Jamie Davis", user_email=None) + ] + entries = await management.list_entries("", 20, 0, None, auth()) + assert entries[0].actor == "contributor" and entries[0].actor_name == "Jamie Davis" + assert database.db.litellm_usertable.find_many.call_args.kwargs["where"] == {"user_id": {"in": ["contributor"]}} + database.db.litellm_usertable.find_many.return_value = [] + legacy = await management.list_entries("", 20, 0, None, auth()) + assert legacy[0].actor == "contributor" and legacy[0].actor_name is None + + +@pytest.mark.asyncio +async def test_admin_manual_capture_uses_actual_author_instead_of_selected_key_owner(database: MagicMock) -> None: + database.db.litellm_memorypolicy.find_many.return_value = [policy(scope="user")] + now = datetime(2026, 9, 14, tzinfo=timezone.utc) + database.db.litellm_memorytable.create.return_value = SimpleNamespace( + memory_id="entry", + key="memory-v2:namespace:demo", + value="Use port 8123", + metadata={}, + updated_at=now, + created_at=now, + created_by="admin", + ) + saved = await management.capture_entry( + MemoryCapture(key="demo", title="Demo", content="Use port 8123", evidence="Admin correction"), + "a" * 64, + auth("admin", LitellmUserRoles.PROXY_ADMIN), + ) + data = database.db.litellm_memorytable.create.call_args.kwargs["data"] + assert data["created_by"] == data["updated_by"] == saved.actor == "admin" + assert data["user_id"] == "owner" + + @pytest.mark.asyncio @pytest.mark.parametrize( "target,target_id", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx index e17f8d7c80a..fc363a83a3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx @@ -1,15 +1,13 @@ "use client"; import { useInfiniteQuery, useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; -import { Brain, ChevronDown, Search } from "lucide-react"; +import { Brain } from "lucide-react"; import { type FormEvent, useState } from "react"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import { Button } from "@/components/ui/button"; -import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; -import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Skeleton } from "@/components/ui/skeleton"; import { Textarea } from "@/components/ui/textarea"; @@ -19,6 +17,7 @@ import { toast } from "@/lib/toast"; import { MemoryPreference } from "./MemorySettings"; import { MemoryKeyPicker } from "./MemoryTargetPicker"; +import { MemoryEntriesTable } from "./MemoryEntriesTable"; type Entry = components["schemas"]["MemoryEntry"]; type Capture = components["schemas"]["MemoryCapture"]; @@ -28,10 +27,10 @@ type Status = components["schemas"]["MemoryStatus"]; function memoryDescription(status?: Status) { if (!status) return "What your assistants remember across conversations."; if (status.activation === "disabled" || !status.scope) - return "Memory is off. Your administrator can make it available for this key."; + return "Memory is off. Your administrator can make it available."; if (status.activation === "automatic") { - if (status.active) return "Memory is on for this key. Your administrator manages this setting."; - return "Memory is off. Your administrator can make it available for this key."; + if (status.active) return "Memory is on. Your administrator manages this setting."; + return "Memory is off. Your administrator can make it available."; } if (status.active) return "Your assistants can save and recall memories. You can turn this off at any time."; return "Your assistants won't save or recall memories. Turn it on whenever you're ready."; @@ -48,21 +47,6 @@ function loadMoreLabel(fetching: boolean, failed: boolean) { return failed ? "Try again" : "Load more memories"; } -function memoryDetails(entry: Entry) { - const details = { - Evidence: entry.evidence, - "When to use": entry.when_to_use, - Source: entry.source, - Kind: entry.kind, - Certainty: entry.certainty, - Context: entry.scope, - "Saved by": entry.actor, - Created: entry.created_at ? new Date(entry.created_at).toLocaleString() : undefined, - Updated: new Date(entry.updated_at).toLocaleString(), - "Memory ID": entry.memory_id, - }; - return Object.entries(details).filter(([, value]) => !!value); -} type DashboardProps = Readonly<{ userId: string; readOnly: boolean; proxyAdmin: boolean }>; export function AutomaticMemoryEntries({ userId, readOnly, proxyAdmin }: DashboardProps) { @@ -80,7 +64,7 @@ export function AutomaticMemoryEntries({ userId, readOnly, proxyAdmin }: Dashboa
(null); const [deleting, setDeleting] = useState(null); - const status = useQuery({ + const statusOptions = { queryKey: ["memoryStatus", userId, keyId], enabled: !!keyId, - queryFn: async ({ signal }) => + refetchOnMount: true, + queryFn: async ({ signal }: { signal: AbortSignal }) => (await fetchClient.GET("/v2/memory/status", { params: { query: { key_id: keyId } }, signal })).data, - }); + }; + const status = useQuery(statusOptions); const entriesOptions = { queryKey: ["memoryEntries", userId, keyId, query], enabled: !!keyId && !!status.data?.scope, + refetchOnMount: true, initialPageParam: null as Cursor, queryFn: async ({ signal, pageParam }: { signal: AbortSignal; pageParam: Cursor }) => ( @@ -148,7 +135,7 @@ function MemoryDashboard({ onSuccess: () => { setEditing(null); toast.success("Memory updated"); - return cache.invalidateQueries({ queryKey: ["memoryEntries", userId, keyId] }); + return cache.invalidateQueries({ queryKey: ["memoryEntries", userId] }); }, onError: (error: Error) => toast.error(error.message), }); @@ -160,7 +147,7 @@ function MemoryDashboard({ onSuccess: () => { setDeleting(null); toast.success("Memory deleted"); - return cache.invalidateQueries({ queryKey: ["memoryEntries", userId, keyId] }); + return cache.invalidateQueries({ queryKey: ["memoryEntries", userId] }); }, onError: (error: Error) => toast.error(error.message), }); @@ -202,17 +189,12 @@ function MemoryDashboard({
{children} - {status.data?.scope && ( -
- - setQuery(event.target.value)} - /> -
+ {status.data?.user_id && ( +

+ Memory for{" "} + {status.data.user_name ?? status.data.user_id} + {status.data.scope === "user" && ". Shared across this user's keys in this organization."} +

)}
{status.error && ( @@ -221,106 +203,55 @@ function MemoryDashboard({

)} {status.data?.scope && ( -
-
-

Saved memories

- - {status.data.scope !== "key" && - status.data.scope !== "user" && - `Shared with your ${status.data.scope} · `} - Newest first - -
- {!status.data.active && memories.length > 0 && ( -

Saved memories stay here while memory is off.

- )} - {entries.isPending && ( -
- {[0, 1, 2].map((row) => ( - - ))} -
- )} - {entries.isSuccess && memories.length === 0 && ( -
- -

{query ? "No matching memories" : "No memories yet"}

-

{emptyDescription(query, status.data.active)}

-
- )} -
    - {memories.map((entry) => ( -
  • -
    -

    {entry.title}

    - +
    +

    Saved memories

    + Could not load memories: {entries.error.message}

    + ) : ( +
    + +

    {query ? "No matching memories" : "No memories yet"}

    +

    {emptyDescription(query, status.data.active)}

    -

    {entry.content}

    - - } - aria-label={`Details for ${entry.title}`} + ) + } + footer={ +
    + + {memories.length} memories + {!status.data.active && memories.length > 0 && " · Saved memories stay here while memory is off"} + + {(entries.hasNextPage || entries.isError) && ( + - -
    - )} - -     -
    • - ))} -
- {entries.error && ( + {loadMoreLabel(entries.isFetching, entries.isError)} + + )} +
+ } + /> + {entries.isFetchNextPageError && (

- Could not load memories: {entries.error.message} + Could not load more memories: {entries.error.message}

)} - {(entries.hasNextPage || entries.isError) && ( -
- -
- )} )} !!value); +} + +export function MemoryEntriesTable({ + entries, + query, + onQueryChange, + loading, + empty, + footer, + readOnly, + canEdit, + busy, + onEdit, + onDelete, +}: Readonly<{ + entries: Entry[]; + query: string; + onQueryChange: (query: string) => void; + loading: boolean; + empty: React.ReactNode; + footer: React.ReactNode; + readOnly: boolean; + canEdit: boolean; + busy: boolean; + onEdit: (entry: Entry) => void; + onDelete: (entry: Entry) => void; +}>) { + const [selectedId, setSelectedId] = useState(null); + const selected = entries.find((entry) => entry.memory_id === selectedId); + const columns: ColumnDef[] = [ + { + id: "memory", + header: "Memory", + size: 600, + cell: ({ row: { original: entry } }) => ( +
+ +

{entry.content}

+
+ ), + }, + { + id: "contributor", + header: "Contributed by", + size: 180, + cell: ({ row }) => {contributor(row.original)}, + }, + { + id: "saved", + header: "Updated", + size: 130, + cell: ({ row }) => ( + + ), + }, + { + id: "details", + header: () => Details, + size: 36, + cell: () =>