From 3bba16c6a6fa830deda9f99d76efd00a42f5e298 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 12 Sep 2026 21:00:50 -0700 Subject: [PATCH] feat(memory): show a recent memory feed and key-aware opt-in toggle --- deploy/memory-pilot/README.md | 17 +- litellm/proxy/memory/management.py | 22 +- litellm/proxy/memory/store.py | 12 +- .../proxy/memory/test_memory_v2_management.py | 84 ++++ .../_components/AutomaticMemoryEntries.tsx | 406 +++++++++++++----- .../memory/_components/MemorySettings.tsx | 64 ++- .../memory/_components/MemoryTargetPicker.tsx | 5 +- .../memory/page.integration.test.tsx | 186 ++++++-- .../src/app/(dashboard)/memory/page.tsx | 31 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 17 +- 10 files changed, 620 insertions(+), 224 deletions(-) diff --git a/deploy/memory-pilot/README.md b/deploy/memory-pilot/README.md index 255ae15c67c..ae46c31c953 100644 --- a/deploy/memory-pilot/README.md +++ b/deploy/memory-pilot/README.md @@ -2,8 +2,8 @@ Run this branch as an isolated forwarding gateway. Colleagues keep their existing upstream LiteLLM key and model name and change only their gateway base URL. They -need no plugin or client-side memory tools. Choosing the pilot URL opts them into -the pilot; returning to the original URL stops using and collecting pilot memory. +need no plugin or client-side memory tools. Memory is off by default. Each person enables it explicitly for their key; +returning to the original URL stops using and collecting pilot memory. Every model call uses that caller's upstream key. The upstream gateway continues to enforce its model permissions, budgets, rate limits, and @@ -49,7 +49,7 @@ pilot must not be connected to an older gateway's production database. | `PORT` | `4000` | 5. After deployment, open `/ui/memory`, sign in as `admin` using the pilot's master - key, and save a policy for **Whole gateway**, **Enabled automatically**, + key, expand **Advanced settings**, and save a policy for **Whole gateway**, **Users choose whether to opt in**, **Private to each virtual key**. This policy persists across restarts. Memory stays disabled until an administrator enables it. @@ -60,18 +60,23 @@ curl --fail-with-body "$PILOT_URL/v2/memory/policies" \ -H "Authorization: Bearer $PILOT_ADMIN_KEY" \ -H 'Content-Type: application/json' \ -X PUT \ - -d '{"target_type":"gateway","target_id":"*","activation":"automatic","scope":"key"}' + -d '{"target_type":"gateway","target_id":"*","activation":"opt_in","scope":"key"}' ``` -Administrators can instead require opt-in, disable a particular registered key, +Administrators can instead enable memory automatically, disable a particular registered key, or disable the whole gateway. Under an opt-in policy, callers set their preference with `PUT /v2/memory/preference` and `{"enabled":true}` using their own key. +In a normal gateway, signed-in users can also select their key in Memory and turn +on the switch. Pilot administrators can do this for a registered key. The switch +shows the actual state; turning it off stops saving and recall but keeps saved +memories visible. Memories appear newest first, with optional details. ## Try it Set an OpenAI-compatible client's base URL to `https://YOUR-SERVICE.onrender.com/v1`. For Claude Code, set `ANTHROPIC_BASE_URL` to `https://YOUR-SERVICE.onrender.com`. -Retain the same gateway key and model setting. +Retain the same gateway key and model setting. First opt in using the preference +API above, or ask the pilot administrator to turn on memory for your key. In one conversation, say “Remember that my demo project is Cobalt Heron and its staging port is 8347.” In a **new conversation**, ask “What is my demo project and diff --git a/litellm/proxy/memory/management.py b/litellm/proxy/memory/management.py index 31e4d52dea4..483f4f8c225 100644 --- a/litellm/proxy/memory/management.py +++ b/litellm/proxy/memory/management.py @@ -167,8 +167,14 @@ async def delete_policy(policy_id: str, auth: UserAPIKeyAuth = _AUTH) -> Respons @router.get("/preference", response_model=MemoryPreference) -async def get_preference(auth: UserAPIKeyAuth = _AUTH) -> MemoryPreference: - subject: Final = MemoryIdentity.from_auth(auth).preference_subject +async def get_preference( + auth: UserAPIKeyAuth = _AUTH, + key_id: Annotated[str | None, Query(pattern=r"^[a-f0-9]{64}$")] = None, +) -> MemoryPreference: + identity: Final = ( + (await access_for_key(auth, key_id)).identity if key_id is not None else MemoryIdentity.from_auth(auth) + ) + subject: Final = identity.preference_subject row: Final = await MemoryPreferenceRepository(memory_primary_client(require_memory_prisma())).table.find_unique( where={ # mutable-ok: Prisma serializes these as native JSON containers. "subject": subject @@ -178,8 +184,14 @@ async def get_preference(auth: UserAPIKeyAuth = _AUTH) -> MemoryPreference: @router.put("/preference", response_model=MemoryPreference) -async def set_preference(preference: MemoryPreference, auth: UserAPIKeyAuth = _AUTH) -> MemoryPreference: - identity: Final = MemoryIdentity.from_auth(auth) +async def set_preference( + preference: MemoryPreference, + auth: UserAPIKeyAuth = _AUTH, + key_id: Annotated[str | None, Query(pattern=r"^[a-f0-9]{64}$")] = None, +) -> MemoryPreference: + identity: Final = ( + (await access_for_key(auth, key_id)).identity if key_id is not None else MemoryIdentity.from_auth(auth) + ) if identity.read_only: raise HTTPException(status_code=403, detail="Read-only users cannot change memory preferences") subject: Final = identity.preference_subject @@ -251,7 +263,7 @@ async def list_entries( prisma: Final = memory_primary_client(require_memory_prisma()) access: Final = await access_for_key(auth, key_id) return await MemoryStore(prisma, access).search( - MemorySearch(query=query, limit=limit, offset=offset), require_active=False + MemorySearch(query=query, limit=limit, offset=offset), require_active=False, recent_first=True ) diff --git a/litellm/proxy/memory/store.py b/litellm/proxy/memory/store.py index 86a45bbad96..6c10c429af9 100644 --- a/litellm/proxy/memory/store.py +++ b/litellm/proxy/memory/store.py @@ -67,10 +67,18 @@ class MemoryStore: raise HTTPException(status_code=403, detail="Memory is not available under the current policy") return current.namespace - async def search(self, search: MemorySearch, *, require_active: bool = True) -> tuple[MemoryEntry, ...]: + async def search( + self, search: MemorySearch, *, require_active: bool = True, recent_first: bool = False + ) -> tuple[MemoryEntry, ...]: entries: Final = await self.entries(require_active=require_active) ranked: Final = await asyncio.to_thread(fuzzy_memories, search.query, entries) - return tuple(entry for entry, _, _ in ranked[search.offset : search.offset + search.limit]) + matched_ids: Final = frozenset(entry.memory_id for entry, _, _ in ranked) + ordered: Final = ( + tuple(entry for entry in entries if entry.memory_id in matched_ids) + if recent_first + else tuple(entry for entry, _, _ in ranked) + ) + return ordered[search.offset : search.offset + search.limit] async def entries(self, *, require_active: bool = True) -> tuple[MemoryEntry, ...]: namespace: Final = await self.authorize_namespace(require_active=require_active) 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 f4b09d23b0f..3b4fc918016 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_management.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_management.py @@ -270,3 +270,87 @@ async def test_entry_endpoints_apply_namespace_and_delete_after_disable(database with pytest.raises(HTTPException) as missing: await management.delete_entry("entry", None, auth()) assert missing.value.status_code == 404 + + +@pytest.mark.parametrize("key_owner", ["owner", None]) +def test_admin_selected_key_preference_controls_actual_request_activation( + database: MagicMock, key_owner: str | None +) -> None: + database.db.litellm_verificationtoken.find_unique.return_value["user_id"] = key_owner + database.db.litellm_memorypolicy.find_many.return_value = [policy(activation="opt_in")] + table = database.db.litellm_memorypreference + app = FastAPI() + app.include_router(management.router) + app.dependency_overrides[user_api_key_auth] = lambda: auth("admin", LitellmUserRoles.PROXY_ADMIN) + params = {"key_id": "a" * 64} + subject = memory_digest("user", key_owner) if key_owner else memory_digest("key", "a" * 64) + with TestClient(app) as client: + assert client.get("/v2/memory/status", params=params).json()["active"] is False + assert client.put("/v2/memory/preference", params=params, json={"enabled": True}).status_code == 200 + assert table.upsert.call_args.kwargs["where"] == {"subject": subject} + table.find_unique.return_value = SimpleNamespace(enabled=True) + assert client.get("/v2/memory/status", params=params).json()["active"] is True + assert client.get("/v2/memory/preference", params=params).json() == {"enabled": True} + assert client.put("/v2/memory/preference", params=params, json={"enabled": False}).status_code == 200 + assert table.delete_many.call_args.kwargs["where"] == {"subject": subject} + table.find_unique.return_value = None + assert client.get("/v2/memory/status", params=params).json()["active"] is False + + +@pytest.mark.parametrize("role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +def test_preference_selection_cannot_bypass_key_ownership_or_readonly( + database: MagicMock, role: LitellmUserRoles +) -> None: + app = FastAPI() + app.include_router(management.router) + app.dependency_overrides[user_api_key_auth] = lambda: auth(role=role).model_copy( + update={"token": "session", "team_id": UI_TEAM_ID} + ) + with TestClient(app) as client: + own = client.put("/v2/memory/preference", params={"key_id": "a" * 64}, json={"enabled": True}) + assert own.status_code == (403 if role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY else 200) + database.db.litellm_memorypreference.upsert.reset_mock() + database.db.litellm_verificationtoken.find_unique.return_value["user_id"] = "another-owner" + assert client.get("/v2/memory/preference", params={"key_id": "a" * 64}).status_code == 403 + assert ( + client.put("/v2/memory/preference", params={"key_id": "a" * 64}, json={"enabled": True}).status_code == 403 + ) + assert ( + client.put("/v2/memory/preference", params={"key_id": "invalid"}, json={"enabled": True}).status_code == 422 + ) + database.db.litellm_memorypreference.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_dashboard_search_keeps_recency_before_pagination_while_agent_search_ranks_relevance( + database: MagicMock, +) -> None: + from litellm.proxy.memory.store import MemoryStore + from litellm.types.memory_v2 import MemorySearch + + database.db.litellm_memorypolicy.find_many.return_value = [policy()] + access = await management.access_for_key(auth(), None) + now = datetime(2026, 9, 12, tzinfo=timezone.utc) + database.db.litellm_memorytable.find_many.return_value = [ + SimpleNamespace( + memory_id=memory_id, + key=f"memory-v2:{access.namespace}:{memory_id}", + namespace=access.namespace, + value=content, + metadata={"title": title, "evidence": "A user instruction"}, + updated_at=now.replace(day=day), + created_at=now.replace(day=day), + created_by="owner", + ) + for memory_id, title, content, day in [ + ("newer", "Demo configuration", "The demo uses port 8123", 12), + ("unrelated", "Theme", "Use dark mode", 11), + ("older", "port", "port", 10), + ] + ] + assert [entry.memory_id for entry in await management.list_entries("port", 1, 0, None, auth())] == ["newer"] + assert [entry.memory_id for entry in await management.list_entries("port", 1, 1, None, auth())] == ["older"] + assert [entry.memory_id for entry in await MemoryStore(database, access).search(MemorySearch(query="port"))] == [ + "older", + "newer", + ] 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 3a1b671eb30..c99788ee782 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/AutomaticMemoryEntries.tsx @@ -1,27 +1,109 @@ "use client"; -import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; +import { useInfiniteQuery, useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; +import { Brain, ChevronDown, Search } 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"; import { fetchClient } from "@/lib/http/api"; import type { components } from "@/lib/http/schema"; import { toast } from "@/lib/toast"; +import { MemoryPreference } from "./MemorySettings"; import { MemoryKeyPicker } from "./MemoryTargetPicker"; type Entry = components["schemas"]["MemoryEntry"]; type Capture = components["schemas"]["MemoryCapture"]; +type Status = components["schemas"]["MemoryStatus"]; -export function AutomaticMemoryEntries({ userId, readOnly }: Readonly<{ userId: string; readOnly: boolean }>) { +function memoryDescription(status?: Status) { + if (!status) return "What your assistants remember across conversations."; + if (status.activation === "automatic") return "Memory is on for this key. Your administrator manages this setting."; + if (status.activation === "disabled" || !status.scope) + return "Memory is off. Your administrator can make it available for this key."; + 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."; +} + +function emptyDescription(query: string, active: boolean) { + if (query) return "Try a different search."; + if (active) return "Use your assistant as usual. What it remembers will appear here."; + return "Turn on memory, then use your assistant as usual. Your memories will appear here."; +} + +function loadMoreLabel(fetching: boolean, failed: boolean) { + if (fetching) return "Loading…"; + 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) { + const [selection, setSelection] = useState(); + const keys = useKeys(1, 1, { userID: proxyAdmin ? undefined : userId, sortBy: "created_at", sortOrder: "desc" }); + const keyId = selection ?? keys.data?.keys[0]?.token ?? ""; + return ( + +
+ + + {keys.error && ( +

+ {keys.error.message} +

+ )} + {keys.isSuccess && keys.data.total_count === 0 && ( +

Create a virtual key to start using memory.

+ )} +
+ + ); +} + +function MemoryDashboard({ + userId, + keyId, + readOnly, + children, +}: Readonly<{ + userId: string; + keyId: string; + readOnly: boolean; + children: React.ReactNode; +}>) { const cache = useQueryClient(); - const [keyId, setKeyId] = useState(""); const [query, setQuery] = useState(""); - const [offset, setOffset] = useState(0); const [editing, setEditing] = useState(null); const [deleting, setDeleting] = useState(null); const status = useQuery({ @@ -30,44 +112,47 @@ export function AutomaticMemoryEntries({ userId, readOnly }: Readonly<{ userId: queryFn: async ({ signal }) => (await fetchClient.GET("/v2/memory/status", { params: { query: { key_id: keyId } }, signal })).data, }); - const entries = useQuery({ - queryKey: ["memoryEntries", userId, keyId, query, offset], + const entriesOptions = { + queryKey: ["memoryEntries", userId, keyId, query], enabled: !!keyId && !!status.data?.scope, - queryFn: async ({ signal }) => + initialPageParam: 0, + queryFn: async ({ signal, pageParam }: { signal: AbortSignal; pageParam: number }) => ( await fetchClient.GET("/v2/memory/entries", { - params: { query: { key_id: keyId, query, offset, limit: 20 } }, + params: { query: { key_id: keyId, query, offset: pageParam, limit: 20 } }, signal, }) - ).data, - }); + ).data ?? [], + getNextPageParam: (lastPage: Entry[], _pages: Entry[][], offset: number) => + lastPage.length === 20 ? offset + 20 : undefined, + }; + const entries = useInfiniteQuery(entriesOptions); const save = useMutation({ - mutationFn: async ({ key, body }: { key: string; body: Capture }) => - fetchClient.POST("/v2/memory/entries", { params: { query: { key_id: key } }, body }), - onSuccess: (_, variables) => { + mutationFn: async (body: Capture) => + fetchClient.POST("/v2/memory/entries", { params: { query: { key_id: keyId } }, body }), + onSuccess: () => { setEditing(null); toast.success("Memory updated"); - return cache.invalidateQueries({ queryKey: ["memoryEntries", userId, variables.key] }); + return cache.invalidateQueries({ queryKey: ["memoryEntries", userId, keyId] }); }, onError: (error: Error) => toast.error(error.message), }); const remove = useMutation({ - mutationFn: async ({ key, memory_id }: { key: string; memory_id: string }) => - fetchClient.DELETE("/v2/memory/entries/{memory_id}", { params: { path: { memory_id }, query: { key_id: key } } }), - onSuccess: (_, variables) => { + mutationFn: async (memory_id: string) => + fetchClient.DELETE("/v2/memory/entries/{memory_id}", { + params: { path: { memory_id }, query: { key_id: keyId } }, + }), + onSuccess: () => { setDeleting(null); toast.success("Memory deleted"); - return cache.invalidateQueries({ queryKey: ["memoryEntries", userId, variables.key] }); + return cache.invalidateQueries({ queryKey: ["memoryEntries", userId, keyId] }); }, onError: (error: Error) => toast.error(error.message), }); const busy = save.isPending || remove.isPending; - const selectKey = (value: string) => { - setKeyId(value); - setOffset(0); - setEditing(null); - setDeleting(null); - }; + const memories = Array.from( + new Map((entries.data?.pages ?? []).flat().map((entry) => [entry.memory_id, entry])).values(), + ); const saveCorrection = (event: FormEvent) => { event.preventDefault(); if (!editing) return; @@ -83,113 +168,198 @@ export function AutomaticMemoryEntries({ userId, readOnly }: Readonly<{ userId: source: editing.source, expected_revision: editing.updated_at, }; - save.mutate({ key: keyId, body }); + save.mutate(body); }; + const description = memoryDescription(status.data); return ( -
-

- Saved gateway memories -

-
- - +
+
+
+

+ Memory +

+

{description}

+
+ {status.data && !status.error && ( + + )} + {keyId && status.isPending && }
- {status.data && ( -

- {status.data.active ? "Automatic memory is active" : "Automatic memory is off"} - {status.data.scope ? ` · ${status.data.scope} scope` : " · No applicable policy"} -

- )} - {(status.error || entries.error) && ( +
+ {children} + {status.data?.scope && ( +
+ + setQuery(event.target.value)} + /> +
+ )} +
+ {status.error && (

- {status.error?.message || entries.error?.message} + Could not load memory status: {status.error.message}

)} {status.data?.scope && ( - { - setQuery(event.target.value); - setOffset(0); - }} - /> - )} - {entries.isFetching &&

Loading memories...

} - {entries.data?.length === 0 &&

No memories match this search

} -
    - {(entries.data ?? []).map((entry) => ( -
  • -

    {entry.title}

    -

    {entry.content}

    -

    Evidence: {entry.evidence}

    - {!readOnly && ( -
    - - -
    - )} -
    • - ))} -
- {editing && ( -
- -