From c1954017cc229ffdd6c7576f2cd9194ca0745d9a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 21 Apr 2026 20:42:12 -0700 Subject: [PATCH] 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) --- .../migration.sql | 24 ++ .../litellm_proxy_extras/schema.prisma | 20 ++ litellm/proxy/memory/__init__.py | 0 litellm/proxy/memory/memory_endpoints.py | 314 +++++++++++++++++ litellm/proxy/proxy_server.py | 2 + litellm/proxy/schema.prisma | 20 ++ litellm/types/memory_management.py | 55 +++ schema.prisma | 20 ++ tests/test_litellm/proxy/memory/__init__.py | 0 .../proxy/memory/test_memory_endpoints.py | 316 ++++++++++++++++++ 10 files changed, 771 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260421120000_add_memory_table/migration.sql create mode 100644 litellm/proxy/memory/__init__.py create mode 100644 litellm/proxy/memory/memory_endpoints.py create mode 100644 litellm/types/memory_management.py create mode 100644 tests/test_litellm/proxy/memory/__init__.py create mode 100644 tests/test_litellm/proxy/memory/test_memory_endpoints.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260421120000_add_memory_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260421120000_add_memory_table/migration.sql new file mode 100644 index 00000000000..76db89203c3 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260421120000_add_memory_table/migration.sql @@ -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"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 08aa5645251..344cc16d76e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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]) +} diff --git a/litellm/proxy/memory/__init__.py b/litellm/proxy/memory/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py new file mode 100644 index 00000000000..b56b570b2a4 --- /dev/null +++ b/litellm/proxy/memory/memory_endpoints.py @@ -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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d8354a798b1..94f96b57b1b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 08aa5645251..344cc16d76e 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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]) +} diff --git a/litellm/types/memory_management.py b/litellm/types/memory_management.py new file mode 100644 index 00000000000..f400f71313c --- /dev/null +++ b/litellm/types/memory_management.py @@ -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 diff --git a/schema.prisma b/schema.prisma index 08aa5645251..344cc16d76e 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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]) +} diff --git a/tests/test_litellm/proxy/memory/__init__.py b/tests/test_litellm/proxy/memory/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/test_litellm/proxy/memory/test_memory_endpoints.py new file mode 100644 index 00000000000..4bbea02e32b --- /dev/null +++ b/tests/test_litellm/proxy/memory/test_memory_endpoints.py @@ -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