mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
feat(proxy): add /v1/memory CRUD endpoints with user/team scoping
New LiteLLM_MemoryTable stores user/team-scoped key/value entries with
optional JSON metadata. Value is a String (LLM-readable text) and metadata
is an optional Json? envelope, matching the Letta + mem0 hybrid model so
future structured fields can be added without a schema migration.
Endpoints:
POST /v1/memory - create
GET /v1/memory - list (caller-scoped; admins see all)
GET /v1/memory/{key} - fetch one
PUT /v1/memory/{key} - upsert
DELETE /v1/memory/{key} - delete
Non-admin callers cannot set a user_id/team_id other than their own.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
26fcbc93e5
commit
c1954017cc
10 changed files with 771 additions and 0 deletions
|
|
@ -0,0 +1,24 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryTable" (
|
||||
"memory_id" TEXT NOT NULL,
|
||||
"key" TEXT NOT NULL,
|
||||
"value" TEXT NOT NULL,
|
||||
"metadata" JSONB,
|
||||
"user_id" TEXT,
|
||||
"team_id" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"created_by" TEXT,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_by" TEXT,
|
||||
|
||||
CONSTRAINT "LiteLLM_MemoryTable_pkey" PRIMARY KEY ("memory_id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_key_user_id_team_id_key" ON "LiteLLM_MemoryTable"("key", "user_id", "team_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_user_id_idx" ON "LiteLLM_MemoryTable"("user_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_team_id_idx" ON "LiteLLM_MemoryTable"("team_id");
|
||||
|
|
@ -1223,3 +1223,23 @@ model LiteLLM_ClaudeCodePluginTable {
|
|||
|
||||
@@map("LiteLLM_ClaudeCodePluginTable")
|
||||
}
|
||||
|
||||
// User/team-scoped memory store. Keyed by (key, user_id, team_id).
|
||||
// `value` is a string (typically markdown/text meant for LLM context);
|
||||
// `metadata` is an optional JSON envelope for structured tags without schema changes.
|
||||
model LiteLLM_MemoryTable {
|
||||
memory_id String @id @default(uuid())
|
||||
key String
|
||||
value String
|
||||
metadata Json?
|
||||
user_id String?
|
||||
team_id String?
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String?
|
||||
|
||||
@@unique([key, user_id, team_id])
|
||||
@@index([user_id])
|
||||
@@index([team_id])
|
||||
}
|
||||
|
|
|
|||
0
litellm/proxy/memory/__init__.py
Normal file
0
litellm/proxy/memory/__init__.py
Normal file
314
litellm/proxy/memory/memory_endpoints.py
Normal file
314
litellm/proxy/memory/memory_endpoints.py
Normal file
|
|
@ -0,0 +1,314 @@
|
|||
"""
|
||||
MEMORY MANAGEMENT
|
||||
|
||||
CRUD endpoints for user/team-scoped memory entries.
|
||||
|
||||
POST /v1/memory - Create a memory entry
|
||||
GET /v1/memory - List memory entries visible to the caller
|
||||
GET /v1/memory/{key} - Get a single memory entry by key
|
||||
PUT /v1/memory/{key} - Upsert (create or update) a memory entry by key
|
||||
DELETE /v1/memory/{key} - Delete a memory entry by key
|
||||
|
||||
Scoping:
|
||||
- Rows carry both `user_id` and `team_id` (each optional).
|
||||
- Visibility: PROXY_ADMIN sees all rows. Non-admin callers see rows whose
|
||||
`user_id` matches their own OR whose `team_id` matches their own.
|
||||
- On create, `user_id`/`team_id` default to the caller's identity unless
|
||||
the caller is a PROXY_ADMIN who explicitly supplies a different scope.
|
||||
"""
|
||||
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import (
|
||||
CommonProxyErrors,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.types.memory_management import (
|
||||
LiteLLM_MemoryRow,
|
||||
MemoryCreateRequest,
|
||||
MemoryDeleteResponse,
|
||||
MemoryListResponse,
|
||||
MemoryUpdateRequest,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
|
||||
|
||||
def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Optional[dict]:
|
||||
"""
|
||||
Prisma `where` fragment restricting rows to those the caller can see.
|
||||
Returns None for admins (no restriction).
|
||||
"""
|
||||
if _is_admin(user_api_key_dict):
|
||||
return None
|
||||
ors: List[dict] = []
|
||||
if user_api_key_dict.user_id:
|
||||
ors.append({"user_id": user_api_key_dict.user_id})
|
||||
if user_api_key_dict.team_id:
|
||||
ors.append({"team_id": user_api_key_dict.team_id})
|
||||
if not ors:
|
||||
# Caller has neither user_id nor team_id — match nothing.
|
||||
return {"memory_id": "__no_match__"}
|
||||
return {"OR": ors}
|
||||
|
||||
|
||||
def _row_to_model(row: Any) -> LiteLLM_MemoryRow:
|
||||
return LiteLLM_MemoryRow(
|
||||
memory_id=row.memory_id,
|
||||
key=row.key,
|
||||
value=row.value,
|
||||
metadata=getattr(row, "metadata", None),
|
||||
user_id=row.user_id,
|
||||
team_id=row.team_id,
|
||||
created_at=row.created_at,
|
||||
created_by=row.created_by,
|
||||
updated_at=row.updated_at,
|
||||
updated_by=row.updated_by,
|
||||
)
|
||||
|
||||
|
||||
def _require_prisma():
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
return prisma_client
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/memory",
|
||||
tags=["memory management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=LiteLLM_MemoryRow,
|
||||
)
|
||||
async def create_memory(
|
||||
body: MemoryCreateRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Create a new memory entry for the caller (or, for admins, any scope)."""
|
||||
prisma_client = _require_prisma()
|
||||
|
||||
if _is_admin(user_api_key_dict):
|
||||
user_id = (
|
||||
body.user_id if body.user_id is not None else user_api_key_dict.user_id
|
||||
)
|
||||
team_id = (
|
||||
body.team_id if body.team_id is not None else user_api_key_dict.team_id
|
||||
)
|
||||
else:
|
||||
# Non-admins cannot set a scope other than their own.
|
||||
if body.user_id is not None and body.user_id != user_api_key_dict.user_id:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Only proxy admins may set user_id to a different user.",
|
||||
)
|
||||
if body.team_id is not None and body.team_id != user_api_key_dict.team_id:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Only proxy admins may set team_id to a different team.",
|
||||
)
|
||||
user_id = user_api_key_dict.user_id
|
||||
team_id = user_api_key_dict.team_id
|
||||
|
||||
try:
|
||||
row = await prisma_client.db.litellm_memorytable.create(
|
||||
data={
|
||||
"key": body.key,
|
||||
"value": body.value,
|
||||
"metadata": body.metadata,
|
||||
"user_id": user_id,
|
||||
"team_id": team_id,
|
||||
"created_by": user_api_key_dict.user_id,
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
# Unique constraint (key, user_id, team_id) → 409.
|
||||
msg = str(e)
|
||||
if "Unique" in msg or "unique" in msg or "UniqueViolation" in msg:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=f"Memory with key '{body.key}' already exists for this scope.",
|
||||
)
|
||||
verbose_proxy_logger.exception("Error creating memory: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return _row_to_model(row)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v1/memory",
|
||||
tags=["memory management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MemoryListResponse,
|
||||
)
|
||||
async def list_memory(
|
||||
key: Optional[str] = Query(None, description="Filter by exact key match."),
|
||||
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),
|
||||
):
|
||||
"""List memory entries visible to the caller."""
|
||||
prisma_client = _require_prisma()
|
||||
|
||||
where: dict = {}
|
||||
vis = _visibility_filter(user_api_key_dict)
|
||||
if vis is not None:
|
||||
where.update(vis)
|
||||
if key is not None:
|
||||
where["key"] = key
|
||||
|
||||
try:
|
||||
total = await prisma_client.db.litellm_memorytable.count(where=where)
|
||||
rows = await prisma_client.db.litellm_memorytable.find_many(
|
||||
where=where,
|
||||
order={"updated_at": "desc"},
|
||||
skip=(page - 1) * page_size,
|
||||
take=page_size,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error listing memory: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return MemoryListResponse(memories=[_row_to_model(r) for r in rows], total=total)
|
||||
|
||||
|
||||
async def _find_memory_for_caller(
|
||||
prisma_client: Any, key: str, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> Any:
|
||||
"""Look up a memory row by key, scoped to the caller's visibility."""
|
||||
where: dict = {"key": key}
|
||||
vis = _visibility_filter(user_api_key_dict)
|
||||
if vis is not None:
|
||||
where.update(vis)
|
||||
rows = await prisma_client.db.litellm_memorytable.find_many(
|
||||
where=where, take=1, order={"updated_at": "desc"}
|
||||
)
|
||||
if not rows:
|
||||
raise HTTPException(
|
||||
status_code=404, detail=f"Memory with key '{key}' not found"
|
||||
)
|
||||
return rows[0]
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v1/memory/{key:path}",
|
||||
tags=["memory management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=LiteLLM_MemoryRow,
|
||||
)
|
||||
async def get_memory(
|
||||
key: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Get a single memory entry by key, scoped to the caller."""
|
||||
prisma_client = _require_prisma()
|
||||
row = await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
|
||||
return _row_to_model(row)
|
||||
|
||||
|
||||
@router.put(
|
||||
"/v1/memory/{key:path}",
|
||||
tags=["memory management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=LiteLLM_MemoryRow,
|
||||
)
|
||||
async def upsert_memory(
|
||||
key: str,
|
||||
body: MemoryUpdateRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""
|
||||
Upsert a memory entry by key within the caller's scope.
|
||||
|
||||
If no row exists for (key, caller.user_id, caller.team_id), create one.
|
||||
If one exists, update the value/metadata fields that were provided.
|
||||
"""
|
||||
prisma_client = _require_prisma()
|
||||
|
||||
data: dict = {}
|
||||
if body.value is not None:
|
||||
data["value"] = body.value
|
||||
if body.metadata is not None:
|
||||
data["metadata"] = body.metadata
|
||||
if not data:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Request body must include at least one of: value, metadata.",
|
||||
)
|
||||
data["updated_by"] = user_api_key_dict.user_id
|
||||
|
||||
try:
|
||||
existing = None
|
||||
try:
|
||||
existing = await _find_memory_for_caller(
|
||||
prisma_client, key, user_api_key_dict
|
||||
)
|
||||
except HTTPException as e:
|
||||
if e.status_code != 404:
|
||||
raise
|
||||
|
||||
if existing is not None:
|
||||
row = await prisma_client.db.litellm_memorytable.update(
|
||||
where={"memory_id": existing.memory_id},
|
||||
data=data,
|
||||
)
|
||||
else:
|
||||
if body.value is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Cannot create a new memory via PUT without a 'value'.",
|
||||
)
|
||||
row = await prisma_client.db.litellm_memorytable.create(
|
||||
data={
|
||||
"key": key,
|
||||
"value": body.value,
|
||||
"metadata": body.metadata,
|
||||
"user_id": user_api_key_dict.user_id,
|
||||
"team_id": user_api_key_dict.team_id,
|
||||
"created_by": user_api_key_dict.user_id,
|
||||
"updated_by": user_api_key_dict.user_id,
|
||||
}
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error upserting memory: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return _row_to_model(row)
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/v1/memory/{key:path}",
|
||||
tags=["memory management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MemoryDeleteResponse,
|
||||
)
|
||||
async def delete_memory(
|
||||
key: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Delete a memory entry by key, scoped to the caller."""
|
||||
prisma_client = _require_prisma()
|
||||
row = await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
|
||||
try:
|
||||
await prisma_client.db.litellm_memorytable.delete(
|
||||
where={"memory_id": row.memory_id}
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error deleting memory: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return MemoryDeleteResponse(key=key, deleted=True)
|
||||
|
|
@ -428,6 +428,7 @@ from litellm.proxy.management_endpoints.team_endpoints import (
|
|||
from litellm.proxy.management_endpoints.tool_management_endpoints import (
|
||||
router as tool_management_router,
|
||||
)
|
||||
from litellm.proxy.memory.memory_endpoints import router as memory_router
|
||||
from litellm.proxy.management_endpoints.ui_sso import (
|
||||
get_disabled_non_admin_personal_key_creation,
|
||||
)
|
||||
|
|
@ -14061,6 +14062,7 @@ app.include_router(model_management_router)
|
|||
app.include_router(model_access_group_management_router)
|
||||
app.include_router(tag_management_router)
|
||||
app.include_router(tool_management_router)
|
||||
app.include_router(memory_router)
|
||||
app.include_router(cost_tracking_settings_router)
|
||||
app.include_router(router_settings_router)
|
||||
app.include_router(fallback_management_router)
|
||||
|
|
|
|||
|
|
@ -1223,3 +1223,23 @@ model LiteLLM_ClaudeCodePluginTable {
|
|||
|
||||
@@map("LiteLLM_ClaudeCodePluginTable")
|
||||
}
|
||||
|
||||
// User/team-scoped memory store. Keyed by (key, user_id, team_id).
|
||||
// `value` is a string (typically markdown/text meant for LLM context);
|
||||
// `metadata` is an optional JSON envelope for structured tags without schema changes.
|
||||
model LiteLLM_MemoryTable {
|
||||
memory_id String @id @default(uuid())
|
||||
key String
|
||||
value String
|
||||
metadata Json?
|
||||
user_id String?
|
||||
team_id String?
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String?
|
||||
|
||||
@@unique([key, user_id, team_id])
|
||||
@@index([user_id])
|
||||
@@index([team_id])
|
||||
}
|
||||
|
|
|
|||
55
litellm/types/memory_management.py
Normal file
55
litellm/types/memory_management.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
"""
|
||||
Pydantic models for Memory management endpoints.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class LiteLLM_MemoryRow(BaseModel):
|
||||
memory_id: str
|
||||
key: str
|
||||
value: str
|
||||
metadata: Optional[Any] = None
|
||||
user_id: Optional[str] = None
|
||||
team_id: Optional[str] = None
|
||||
created_at: Optional[datetime] = None
|
||||
created_by: Optional[str] = None
|
||||
updated_at: Optional[datetime] = None
|
||||
updated_by: Optional[str] = None
|
||||
|
||||
|
||||
class MemoryCreateRequest(BaseModel):
|
||||
key: str = Field(..., description="Memory key (acts as the namespace in the URL).")
|
||||
value: str = Field(
|
||||
..., description="Memory content. Typically markdown/text for LLM context."
|
||||
)
|
||||
metadata: Optional[Any] = Field(
|
||||
default=None,
|
||||
description="Optional JSON metadata (tags, structured fields).",
|
||||
)
|
||||
user_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Scope to this user. Defaults to the caller's user_id.",
|
||||
)
|
||||
team_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Scope to this team. Defaults to the caller's team_id.",
|
||||
)
|
||||
|
||||
|
||||
class MemoryUpdateRequest(BaseModel):
|
||||
value: Optional[str] = None
|
||||
metadata: Optional[Any] = None
|
||||
|
||||
|
||||
class MemoryListResponse(BaseModel):
|
||||
memories: List[LiteLLM_MemoryRow]
|
||||
total: int
|
||||
|
||||
|
||||
class MemoryDeleteResponse(BaseModel):
|
||||
key: str
|
||||
deleted: bool
|
||||
|
|
@ -1223,3 +1223,23 @@ model LiteLLM_ClaudeCodePluginTable {
|
|||
|
||||
@@map("LiteLLM_ClaudeCodePluginTable")
|
||||
}
|
||||
|
||||
// User/team-scoped memory store. Keyed by (key, user_id, team_id).
|
||||
// `value` is a string (typically markdown/text meant for LLM context);
|
||||
// `metadata` is an optional JSON envelope for structured tags without schema changes.
|
||||
model LiteLLM_MemoryTable {
|
||||
memory_id String @id @default(uuid())
|
||||
key String
|
||||
value String
|
||||
metadata Json?
|
||||
user_id String?
|
||||
team_id String?
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String?
|
||||
|
||||
@@unique([key, user_id, team_id])
|
||||
@@index([user_id])
|
||||
@@index([team_id])
|
||||
}
|
||||
|
|
|
|||
0
tests/test_litellm/proxy/memory/__init__.py
Normal file
0
tests/test_litellm/proxy/memory/__init__.py
Normal file
316
tests/test_litellm/proxy/memory/test_memory_endpoints.py
Normal file
316
tests/test_litellm/proxy/memory/test_memory_endpoints.py
Normal file
|
|
@ -0,0 +1,316 @@
|
|||
"""
|
||||
Unit tests for /v1/memory CRUD endpoints.
|
||||
|
||||
Uses FastAPI TestClient with an in-memory fake of the Prisma memory table.
|
||||
Auth is overridden so we can simulate different callers (admin vs. scoped).
|
||||
We patch the endpoint module's `_require_prisma` helper so we never need the
|
||||
real proxy_server import chain (which pulls heavy optional deps).
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.memory.memory_endpoints import router
|
||||
|
||||
|
||||
def _make_row(
|
||||
memory_id: str = "mem-1",
|
||||
key: str = "notes",
|
||||
value: str = "hello",
|
||||
user_id: Optional[str] = "user-a",
|
||||
team_id: Optional[str] = "team-a",
|
||||
metadata: Optional[Any] = None,
|
||||
) -> MagicMock:
|
||||
"""Build a Prisma-like row object."""
|
||||
now = datetime.now(timezone.utc)
|
||||
row = MagicMock()
|
||||
row.memory_id = memory_id
|
||||
row.key = key
|
||||
row.value = value
|
||||
row.metadata = metadata
|
||||
row.user_id = user_id
|
||||
row.team_id = team_id
|
||||
row.created_at = now
|
||||
row.created_by = user_id
|
||||
row.updated_at = now
|
||||
row.updated_by = user_id
|
||||
return row
|
||||
|
||||
|
||||
class _InMemoryMemoryTable:
|
||||
"""Tiny fake of prisma_client.db.litellm_memorytable used by the endpoints."""
|
||||
|
||||
def __init__(self):
|
||||
self.rows: List[MagicMock] = []
|
||||
self._counter = 0
|
||||
|
||||
def _matches(self, row: MagicMock, where: Dict[str, Any]) -> bool:
|
||||
for k, v in where.items():
|
||||
if k == "OR":
|
||||
if not any(self._matches(row, clause) for clause in v):
|
||||
return False
|
||||
continue
|
||||
if getattr(row, k, None) != v:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _filter(self, where: Optional[Dict[str, Any]]) -> List[MagicMock]:
|
||||
if not where:
|
||||
return list(self.rows)
|
||||
return [r for r in self.rows if self._matches(r, where)]
|
||||
|
||||
async def create(self, data: Dict[str, Any]) -> MagicMock:
|
||||
for r in self.rows:
|
||||
if (
|
||||
r.key == data["key"]
|
||||
and r.user_id == data.get("user_id")
|
||||
and r.team_id == data.get("team_id")
|
||||
):
|
||||
raise Exception("UniqueViolation: duplicate key")
|
||||
self._counter += 1
|
||||
row = _make_row(
|
||||
memory_id=f"mem-{self._counter}",
|
||||
key=data["key"],
|
||||
value=data["value"],
|
||||
user_id=data.get("user_id"),
|
||||
team_id=data.get("team_id"),
|
||||
metadata=data.get("metadata"),
|
||||
)
|
||||
row.created_by = data.get("created_by")
|
||||
row.updated_by = data.get("updated_by")
|
||||
self.rows.append(row)
|
||||
return row
|
||||
|
||||
async def count(self, where: Optional[Dict[str, Any]] = None) -> int:
|
||||
return len(self._filter(where))
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
where: Optional[Dict[str, Any]] = None,
|
||||
order: Optional[Dict[str, str]] = None,
|
||||
skip: int = 0,
|
||||
take: Optional[int] = None,
|
||||
) -> List[MagicMock]:
|
||||
_ = order
|
||||
out = self._filter(where)
|
||||
if skip:
|
||||
out = out[skip:]
|
||||
if take is not None:
|
||||
out = out[:take]
|
||||
return out
|
||||
|
||||
async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> MagicMock:
|
||||
for r in self.rows:
|
||||
if r.memory_id == where["memory_id"]:
|
||||
for k, v in data.items():
|
||||
setattr(r, k, v)
|
||||
return r
|
||||
raise Exception("Not found")
|
||||
|
||||
async def delete(self, where: Dict[str, Any]) -> MagicMock:
|
||||
for i, r in enumerate(self.rows):
|
||||
if r.memory_id == where["memory_id"]:
|
||||
return self.rows.pop(i)
|
||||
raise Exception("Not found")
|
||||
|
||||
|
||||
def _make_prisma() -> MagicMock:
|
||||
client = MagicMock()
|
||||
client.db = MagicMock()
|
||||
client.db.litellm_memorytable = _InMemoryMemoryTable()
|
||||
return client
|
||||
|
||||
|
||||
def _make_client(auth: UserAPIKeyAuth) -> TestClient:
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: auth
|
||||
return TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
|
||||
def _user_auth(user_id: str = "user-a", team_id: str = "team-a") -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(api_key="sk-test", user_id=user_id, team_id=team_id)
|
||||
|
||||
|
||||
def _admin_auth() -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="sk-admin",
|
||||
user_id="admin",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
|
||||
def _patch_prisma(prisma: Any):
|
||||
"""Patch the endpoint module's _require_prisma to return our fake."""
|
||||
return patch(
|
||||
"litellm.proxy.memory.memory_endpoints._require_prisma",
|
||||
return_value=prisma,
|
||||
)
|
||||
|
||||
|
||||
class TestMemoryEndpoints:
|
||||
def setup_method(self):
|
||||
self.prisma = _make_prisma()
|
||||
|
||||
def test_create_memory_defaults_scope_to_caller(self):
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.post("/v1/memory", json={"key": "notes", "value": "hello"})
|
||||
assert resp.status_code == 200, resp.text
|
||||
body = resp.json()
|
||||
assert body["key"] == "notes"
|
||||
assert body["value"] == "hello"
|
||||
assert body["user_id"] == "user-a"
|
||||
assert body["team_id"] == "team-a"
|
||||
|
||||
def test_create_memory_duplicate_key_returns_409(self):
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
r1 = client.post("/v1/memory", json={"key": "k", "value": "v1"})
|
||||
assert r1.status_code == 200
|
||||
r2 = client.post("/v1/memory", json={"key": "k", "value": "v2"})
|
||||
assert r2.status_code == 409
|
||||
|
||||
def test_non_admin_cannot_set_foreign_user_id(self):
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.post(
|
||||
"/v1/memory",
|
||||
json={"key": "notes", "value": "x", "user_id": "user-b"},
|
||||
)
|
||||
assert resp.status_code == 403
|
||||
|
||||
def test_admin_can_set_any_scope(self):
|
||||
client = _make_client(_admin_auth())
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.post(
|
||||
"/v1/memory",
|
||||
json={
|
||||
"key": "notes",
|
||||
"value": "x",
|
||||
"user_id": "some-user",
|
||||
"team_id": "some-team",
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["user_id"] == "some-user"
|
||||
assert body["team_id"] == "some-team"
|
||||
|
||||
def test_list_memory_scoped_to_caller(self):
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.extend(
|
||||
[
|
||||
_make_row(memory_id="m1", key="a", user_id="user-a", team_id=None),
|
||||
_make_row(memory_id="m2", key="b", user_id="user-b", team_id=None),
|
||||
_make_row(memory_id="m3", key="c", user_id=None, team_id="team-a"),
|
||||
]
|
||||
)
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.get("/v1/memory")
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
keys = {m["key"] for m in body["memories"]}
|
||||
assert keys == {"a", "c"}
|
||||
assert body["total"] == 2
|
||||
|
||||
def test_list_memory_admin_sees_all(self):
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.extend(
|
||||
[
|
||||
_make_row(memory_id="m1", key="a", user_id="user-a", team_id=None),
|
||||
_make_row(memory_id="m2", key="b", user_id="user-b", team_id=None),
|
||||
]
|
||||
)
|
||||
client = _make_client(_admin_auth())
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.get("/v1/memory")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["total"] == 2
|
||||
|
||||
def test_get_memory_by_key(self):
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.append(
|
||||
_make_row(memory_id="m1", key="notes", value="hi", user_id="user-a")
|
||||
)
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.get("/v1/memory/notes")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["value"] == "hi"
|
||||
|
||||
def test_get_memory_not_visible_returns_404(self):
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.append(
|
||||
_make_row(memory_id="m1", key="notes", 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/notes")
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_put_memory_creates_when_missing(self):
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.put("/v1/memory/notes", json={"value": "new"})
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["value"] == "new"
|
||||
assert len(self.prisma.db.litellm_memorytable.rows) == 1
|
||||
|
||||
def test_put_memory_updates_existing(self):
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.append(
|
||||
_make_row(
|
||||
memory_id="m1",
|
||||
key="notes",
|
||||
value="old",
|
||||
user_id="user-a",
|
||||
team_id="team-a",
|
||||
)
|
||||
)
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.put("/v1/memory/notes", json={"value": "new"})
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["value"] == "new"
|
||||
assert len(table.rows) == 1
|
||||
|
||||
def test_put_memory_empty_body_returns_400(self):
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.put("/v1/memory/notes", json={})
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_delete_memory(self):
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.append(
|
||||
_make_row(memory_id="m1", key="notes", user_id="user-a", team_id="team-a")
|
||||
)
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.delete("/v1/memory/notes")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json() == {"key": "notes", "deleted": True}
|
||||
assert table.rows == []
|
||||
|
||||
def test_delete_memory_not_visible_returns_404(self):
|
||||
table = self.prisma.db.litellm_memorytable
|
||||
table.rows.append(
|
||||
_make_row(memory_id="m1", key="notes", user_id="user-b", team_id=None)
|
||||
)
|
||||
client = _make_client(_user_auth("user-a", "team-a"))
|
||||
with _patch_prisma(self.prisma):
|
||||
resp = client.delete("/v1/memory/notes")
|
||||
assert resp.status_code == 404
|
||||
Loading…
Add table
Reference in a new issue