mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
feat(memory): add key_prefix filter + promote Memory to AI GATEWAY nav
Backend:
- GET /v1/memory now accepts `key_prefix` for Redis-style namespace
scans (e.g. `?key_prefix=user:`). When both `key` and `key_prefix`
are passed, `key_prefix` wins.
- Prefix filter sits under the visibility filter in the Prisma where
clause, so it can never leak rows across user/team scopes.
- New tests: prefix match, and cross-scope isolation (another user's
`user:*` rows must not appear in the caller's results).
UI:
- Memory moved from a Tools submenu to a top-level AI GATEWAY item
(alongside Agents, MCP Servers, Skills) — it's an API primitive,
not a tool-management surface.
- Search box now drives prefix search, matching the Redis mental
model ("type the namespace, see everything under it").
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
b0b1ea3740
commit
cc1f0d4cf3
5 changed files with 109 additions and 12 deletions
|
|
@ -159,6 +159,13 @@ async def create_memory(
|
|||
)
|
||||
async def list_memory(
|
||||
key: Optional[str] = Query(None, description="Filter by exact key match."),
|
||||
key_prefix: Optional[str] = Query(
|
||||
None,
|
||||
description=(
|
||||
"Filter by key prefix (Redis-style namespace scan). "
|
||||
"Mutually exclusive with `key`; if both are provided, `key_prefix` wins."
|
||||
),
|
||||
),
|
||||
page: int = Query(1, ge=1),
|
||||
page_size: int = Query(50, ge=1, le=500),
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
|
|
@ -170,7 +177,12 @@ async def list_memory(
|
|||
vis = _visibility_filter(user_api_key_dict)
|
||||
if vis is not None:
|
||||
where.update(vis)
|
||||
if key is not None:
|
||||
# key_prefix takes precedence over exact `key` if both are passed.
|
||||
# Both sit under the visibility filter (ANDed), so prefix scans can
|
||||
# never leak across user/team scopes.
|
||||
if key_prefix is not None:
|
||||
where["key"] = {"startsWith": key_prefix}
|
||||
elif key is not None:
|
||||
where["key"] = key
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -59,6 +59,21 @@ class _InMemoryMemoryTable:
|
|||
if not any(self._matches(row, clause) for clause in v):
|
||||
return False
|
||||
continue
|
||||
# Support Prisma-style filter dicts: {"startsWith": "..."}, etc.
|
||||
if isinstance(v, dict):
|
||||
actual = getattr(row, k, None)
|
||||
if "startsWith" in v:
|
||||
if not isinstance(actual, str) or not actual.startswith(
|
||||
v["startsWith"]
|
||||
):
|
||||
return False
|
||||
continue
|
||||
if "equals" in v:
|
||||
if actual != v["equals"]:
|
||||
return False
|
||||
continue
|
||||
# Unknown filter — fall back to inequality (treat as no match).
|
||||
return False
|
||||
if getattr(row, k, None) != v:
|
||||
return False
|
||||
return True
|
||||
|
|
@ -226,6 +241,64 @@ class TestMemoryEndpoints:
|
|||
assert keys == {"a", "c"}
|
||||
assert body["total"] == 2
|
||||
|
||||
def test_list_memory_key_prefix_filter(self):
|
||||
"""key_prefix should do a prefix match (Redis-style namespace scan)."""
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.extend(
|
||||
[
|
||||
_make_row(
|
||||
memory_id="m1", key="user:profile", user_id="user-a", team_id=None
|
||||
),
|
||||
_make_row(
|
||||
memory_id="m2", key="user:prefs", user_id="user-a", team_id=None
|
||||
),
|
||||
_make_row(
|
||||
memory_id="m3",
|
||||
key="project:context",
|
||||
user_id="user-a",
|
||||
team_id=None,
|
||||
),
|
||||
]
|
||||
)
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.get("/v1/memory?key_prefix=user:")
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
keys = {m["key"] for m in body["memories"]}
|
||||
assert keys == {"user:profile", "user:prefs"}
|
||||
assert body["total"] == 2
|
||||
|
||||
def test_list_memory_key_prefix_does_not_leak_across_scopes(self):
|
||||
"""
|
||||
Even if a prefix would match another user's keys, the visibility
|
||||
filter must still scope results to the caller.
|
||||
"""
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.extend(
|
||||
[
|
||||
# Caller's own "user:*" rows — should be visible.
|
||||
_make_row(
|
||||
memory_id="m1", key="user:profile", user_id="user-a", team_id=None
|
||||
),
|
||||
# Another user's "user:*" rows — must NOT leak.
|
||||
_make_row(
|
||||
memory_id="m2", key="user:secret", user_id="user-b", team_id=None
|
||||
),
|
||||
_make_row(
|
||||
memory_id="m3", key="user:token", user_id="user-b", team_id=None
|
||||
),
|
||||
]
|
||||
)
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.get("/v1/memory?key_prefix=user:")
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
keys = {m["key"] for m in body["memories"]}
|
||||
assert keys == {"user:profile"}
|
||||
assert body["total"] == 1
|
||||
|
||||
def test_list_memory_admin_sees_all(self):
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.extend(
|
||||
|
|
|
|||
|
|
@ -74,7 +74,8 @@ export const MemoryView: React.FC<MemoryViewProps> = ({ accessToken }) => {
|
|||
queryKey: ["memoryList", appliedSearch],
|
||||
queryFn: () => {
|
||||
if (!accessToken) throw new Error("Access token required");
|
||||
// Current API supports exact key match; treat empty search as list-all.
|
||||
// Prefix search matches the Redis-style mental model (namespace scan):
|
||||
// typing "user:" finds "user:profile", "user:prefs", etc.
|
||||
return fetchMemoryList(accessToken, {
|
||||
keyPrefix: appliedSearch || undefined,
|
||||
pageSize: 200,
|
||||
|
|
@ -254,7 +255,7 @@ export const MemoryView: React.FC<MemoryViewProps> = ({ accessToken }) => {
|
|||
<Space>
|
||||
<Input
|
||||
allowClear
|
||||
placeholder="Filter by exact key"
|
||||
placeholder='Filter by key prefix, e.g. "user:"'
|
||||
prefix={<SearchOutlined />}
|
||||
value={searchInput}
|
||||
onChange={(e) => setSearchInput(e.target.value)}
|
||||
|
|
@ -300,7 +301,7 @@ export const MemoryView: React.FC<MemoryViewProps> = ({ accessToken }) => {
|
|||
<Empty
|
||||
description={
|
||||
appliedSearch
|
||||
? `No memories found for key "${appliedSearch}"`
|
||||
? `No memories with keys starting with "${appliedSearch}"`
|
||||
: "No memories stored yet"
|
||||
}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -139,6 +139,12 @@ const menuGroups: MenuGroup[] = [
|
|||
icon: <ApiOutlined />,
|
||||
roles: all_admin_roles,
|
||||
},
|
||||
{
|
||||
key: "memory",
|
||||
page: "memory",
|
||||
label: "Memory",
|
||||
icon: <BookOutlined />,
|
||||
},
|
||||
{
|
||||
key: "guardrails",
|
||||
page: "guardrails",
|
||||
|
|
@ -180,12 +186,6 @@ const menuGroups: MenuGroup[] = [
|
|||
label: "Tool Policies",
|
||||
icon: <SafetyOutlined />,
|
||||
},
|
||||
{
|
||||
key: "memory",
|
||||
page: "memory",
|
||||
label: "Memory",
|
||||
icon: <BookOutlined />,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
|
|
|
|||
|
|
@ -10011,11 +10011,22 @@ export interface MemoryListResponse {
|
|||
|
||||
export const fetchMemoryList = async (
|
||||
accessToken: string,
|
||||
options: { keyPrefix?: string; page?: number; pageSize?: number } = {},
|
||||
options: {
|
||||
key?: string;
|
||||
keyPrefix?: string;
|
||||
page?: number;
|
||||
pageSize?: number;
|
||||
} = {},
|
||||
): Promise<MemoryListResponse> => {
|
||||
const base = proxyBaseUrl ? `${proxyBaseUrl}/v1/memory` : `/v1/memory`;
|
||||
const params = new URLSearchParams();
|
||||
if (options.keyPrefix) params.append("key", options.keyPrefix);
|
||||
// keyPrefix takes precedence — backend also does, but we omit `key`
|
||||
// to keep the URL clean and intent obvious.
|
||||
if (options.keyPrefix) {
|
||||
params.append("key_prefix", options.keyPrefix);
|
||||
} else if (options.key) {
|
||||
params.append("key", options.key);
|
||||
}
|
||||
if (options.page != null) params.append("page", String(options.page));
|
||||
if (options.pageSize != null)
|
||||
params.append("page_size", String(options.pageSize));
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue