Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/memory/memory

1"""

2MEMORY MANAGEMENT

4CRUD endpoints for user/team-scoped memory entries.

6POST /v1/memory - Create a memory entry

7GET /v1/memory - List memory entries visible to the caller

8GET /v1/memory/{key} - Get a single memory entry by key

9PUT /v1/memory/{key} - Upsert (create or update) a memory entry by key

10DELETE /v1/memory/{key} - Delete a memory entry by key

12Scoping:

13- Rows carry both `user_id` and `team_id` (each optional).

14- Visibility: PROXY_ADMIN sees all rows. Non-admin callers see rows whose

15 `user_id` matches their own OR whose `team_id` matches their own.

16- On create, `user_id`/`team_id` default to the caller's identity unless

17 the caller is a PROXY_ADMIN who explicitly supplies a different scope.

18"""

20import json

21from collections.abc import Mapping

22from typing import TYPE_CHECKING, Final

24from fastapi import APIRouter, Depends, HTTPException, Query

25from typing_extensions import ReadOnly, TypedDict

27from litellm._logging import verbose_proxy_logger

28from litellm.proxy._types import (

29 CommonProxyErrors,

30 LitellmUserRoles,

31 UserAPIKeyAuth,

32 user_api_key_has_admin_view,

33)

34from litellm.proxy.auth.user_api_key_auth import user_api_key_auth

35from litellm.repositories.prisma_protocols import TableActions

36from litellm.repositories.table_repositories import MemoryRepository

37from litellm.repositories.team_repository import TeamRepository

38from litellm.types.memory_management import (

39 LiteLLM_MemoryRow,

40 MemoryCreateRequest,

41 MemoryDeleteResponse,

42 MemoryListResponse,

43 MemoryUpdateRequest,

44)

46if TYPE_CHECKING: 46 ↛ 47line 46 didn't jump to line 47 because the condition on line 46 was never true

47 from prisma import models as prisma_models

49 from litellm.proxy.utils import PrismaClient

51router: Final = APIRouter()

54def _memory_table(prisma_client: "PrismaClient") -> TableActions["prisma_models.LiteLLM_MemoryTable"]:

55 return MemoryRepository(prisma_client).table

58def _serialize_metadata_for_prisma(metadata: object) -> str:

59 """

60 Encode a `metadata` payload for the `Json?` column.

62 `metadata` is typed `Optional[Any]`, so callers may send dicts, lists,

63 or JSON scalars (including plain Python strings like `"hello"`).

64 prisma-client-python rejects raw Python values on `Json?` columns

65 (`MissingRequiredValueError` / `DataError`), and Postgres `jsonb`

66 rejects bare-word strings as invalid JSON — so always `json.dumps`,

67 regardless of input type. Roundtrip on read deserializes back to the

68 original Python value.

69 """

70 return json.dumps(metadata)

73def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:

74 return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN

77def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object] | None:

78 """

79 Prisma `where` fragment restricting rows to those the caller can see.

80 Returns None for admins (no restriction).

81 """

82 if user_api_key_has_admin_view(user_api_key_dict): 82 ↛ 84line 82 didn't jump to line 84 because the condition on line 82 was always true

83 return None

84 ors: Final = [

85 {field: value}

86 for field, value in (("user_id", user_api_key_dict.user_id), ("team_id", user_api_key_dict.team_id))

87 if value

88 ]

89 if not ors:

90 # Caller has neither user_id nor team_id — match nothing.

91 return {"memory_id": "__no_match__"}

92 return {"OR": ors}

95class _StartsWith(TypedDict):

96 startsWith: ReadOnly[str]

99class _MemoryKeyWhere(TypedDict):

100 key: ReadOnly[str | _StartsWith]

101

102

103class _MemoryIdWhere(TypedDict):

104 memory_id: ReadOnly[str]

105

106

107class _MemorySearchWhere(TypedDict):

108 OR: ReadOnly[tuple[_MemoryKeyWhere, _MemoryIdWhere]]

109

110

111def _key_filter(search: str | None, key_prefix: str | None, key: str | None) -> Mapping[str, object] | None:

112 """`search` matches a key prefix or an exact memory_id; otherwise `key_prefix` wins over `key`."""

113 if search is not None:

114 search_where: Final[_MemorySearchWhere] = {"OR": ({"key": {"startsWith": search}}, {"memory_id": search})}

115 return search_where

116 if key_prefix is not None:

117 prefix_where: Final[_MemoryKeyWhere] = {"key": {"startsWith": key_prefix}}

118 return prefix_where

119 if key is not None:

120 exact_where: Final[_MemoryKeyWhere] = {"key": key}

121 return exact_where

122 return None

123

124

125def _row_to_model(row: "prisma_models.LiteLLM_MemoryTable") -> LiteLLM_MemoryRow:

126 return LiteLLM_MemoryRow(

127 memory_id=row.memory_id,

128 key=row.key,

129 value=row.value,

130 metadata=getattr(row, "metadata", None),

131 user_id=row.user_id,

132 team_id=row.team_id,

133 created_at=row.created_at,

134 created_by=row.created_by,

135 updated_at=row.updated_at,

136 updated_by=row.updated_by,

137 )

138

139

140def _require_prisma() -> "PrismaClient":

141 from litellm.proxy.proxy_server import prisma_client

142

143 if prisma_client is None: 143 ↛ 144line 143 didn't jump to line 144 because the condition on line 143 was never true

144 raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)

145 return prisma_client

146

147

148def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HTTPException:

149 """

150 Build a 500 HTTPException with a generic, caller-safe `detail` while

151 logging the actual exception server-side. Avoids leaking internal Prisma /

152 DB details (table names, columns, connection metadata) to API callers.

153 """

154 verbose_proxy_logger.exception(log_message, exc)

155 return HTTPException(status_code=500, detail=default_detail)

156

157

158async def _assert_write_access(

159 prisma_client: "PrismaClient", row: "prisma_models.LiteLLM_MemoryTable", user_api_key_dict: UserAPIKeyAuth

160) -> None:

161 """

162 Enforce ownership for mutations (PUT/DELETE).

163

164 The visibility filter uses an OR (`user_id == caller OR team_id == caller`)

165 so team members can READ each other's team-scoped rows. That's intentional

166 for list/get. For writes, broader visibility != broader authority — without

167 this check, any team member could overwrite or delete a teammate's

168 personal row whenever both `user_id` and `team_id` are stamped on it.

169

170 Rules (mirroring how key/team management endpoints gate team-scoped writes):

171 - PROXY_ADMIN: always allowed.

172 - Personal ownership (`row.user_id == caller.user_id`): allowed.

173 - Pure team row (`row.user_id is None`, `row.team_id` set):

174 caller must be a team admin of `row.team_id` (members_with_roles entry

175 with `role == "admin"`), or an org admin for that team's organization.

176 Plain team members can only READ team rows, not modify them — same

177 pattern as `_validate_team_member_add_permissions` etc.

178 - Anything else: 403.

179 """

180 if _is_admin(user_api_key_dict): 180 ↛ 182line 180 didn't jump to line 182 because the condition on line 180 was always true

181 return

182 row_user_id: Final = getattr(row, "user_id", None)

183 row_team_id: Final = getattr(row, "team_id", None)

184

185 # Personal ownership.

186 if row_user_id and row_user_id == user_api_key_dict.user_id:

187 return

188

189 # Pure team row — only team admins (or org admins) may write.

190 if row_user_id is None and row_team_id is not None:

191 if await _is_team_admin_for(prisma_client, user_api_key_dict, row_team_id):

192 return

193

194 raise HTTPException(

195 status_code=403,

196 detail="You do not have permission to modify this memory entry.",

197 )

198

199

200async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool:

201 """

202 True if the caller is a team admin of `team_id`, or an org admin for the

203 team's organization. Mirrors the auth pattern used by team-management

204 endpoints (`_is_user_team_admin` + `_is_user_org_admin_for_team`).

205

206 Imported lazily to avoid a circular import with proxy_server during the

207 memory router's module load.

208 """

209 from litellm.proxy.management_endpoints.common_utils import (

210 _is_user_org_admin_for_team,

211 _is_user_team_admin,

212 )

213

214 try:

215 team_obj: Final = await TeamRepository(prisma_client).find_by_id(team_id, id_field="team_id")

216 except Exception as e:

217 verbose_proxy_logger.exception("Error loading team for write-auth check (team_id=%s): %s", team_id, e)

218 return False

219 if team_obj is None:

220 return False

221

222 if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):

223 return True

224

225 # Org-admin path is best-effort: it pulls from the user cache via

226 # `get_user_object` which depends on the proxy_server module being

227 # initialized. In tests / non-proxy contexts that import path may fail —

228 # treat any error as "not an org admin" rather than crashing the request.

229 try:

230 if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj):

231 return True

232 except Exception as e:

233 verbose_proxy_logger.debug("Org-admin check skipped during write-auth (team_id=%s): %s", team_id, e)

234 return False

235

236

237def _is_unique_violation(exc: Exception) -> bool:

238 """

239 Detect a Prisma unique-constraint violation.

240

241 Prefer the typed error code `P2002` from `PrismaClientKnownRequestError`;

242 fall back to string matching so we stay robust across Prisma versions

243 where the typed class may be unavailable or differently named.

244 """

245 code: Final = getattr(exc, "code", None)

246 if code == "P2002": 246 ↛ 248line 246 didn't jump to line 248 because the condition on line 246 was always true

247 return True

248 msg: Final = str(exc)

249 return "P2002" in msg or "Unique" in msg or "unique" in msg or "UniqueViolation" in msg

250

251

252def _resolve_scope(

253 user_api_key_dict: UserAPIKeyAuth,

254 requested_user_id: str | None,

255 requested_team_id: str | None,

256) -> tuple[str | None, str | None]:

257 """

258 Resolve the (user_id, team_id) to stamp on a new row.

259

260 - PROXY_ADMIN: may override either dimension via the request body.

261 - Everyone else: the requested values must match their own (or be omitted).

262

263 Also rejects identity-less creation: a row with both user_id and team_id

264 NULL is invisible to every non-admin caller (the visibility filter would

265 never match it), so we refuse to create orphan rows unless the caller is

266 a PROXY_ADMIN who is explicitly stamping a global/shared row.

267 """

268 if _is_admin(user_api_key_dict): 268 ↛ 273line 268 didn't jump to line 273 because the condition on line 268 was always true

269 user_id = requested_user_id if requested_user_id is not None else user_api_key_dict.user_id

270 team_id = requested_team_id if requested_team_id is not None else user_api_key_dict.team_id

271 return user_id, team_id

272

273 if requested_user_id is not None and requested_user_id != user_api_key_dict.user_id:

274 raise HTTPException(

275 status_code=403,

276 detail="Only proxy admins may set user_id to a different user.",

277 )

278 if requested_team_id is not None and requested_team_id != user_api_key_dict.team_id:

279 raise HTTPException(

280 status_code=403,

281 detail="Only proxy admins may set team_id to a different team.",

282 )

283 user_id = user_api_key_dict.user_id

284 team_id = user_api_key_dict.team_id

285 if not user_id and not team_id:

286 # Orphan row: no user_id and no team_id means no non-admin can ever

287 # see it again via the visibility filter. Reject up front.

288 raise HTTPException(

289 status_code=400,

290 detail=(

291 "Cannot create a memory entry without a user_id or team_id. "

292 "Authenticate with a key that has a user_id or team_id, or call "

293 "as a proxy admin."

294 ),

295 )

296 return user_id, team_id

297

298

299@router.post(

300 "/v1/memory",

301 tags=["memory management"],

302 dependencies=[Depends(user_api_key_auth)],

303 response_model=LiteLLM_MemoryRow,

304)

305async def create_memory(

306 body: MemoryCreateRequest,

307 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),

308):

309 """Create a new memory entry for the caller (or, for admins, any scope)."""

310 prisma_client: Final = _require_prisma()

311 user_id, team_id = _resolve_scope(user_api_key_dict, body.user_id, body.team_id)

312

313 # `metadata` is a `Json?` column — prisma-client-python rejects raw

314 # Python values, so JSON-encode any non-null payload and omit the field

315 # entirely when None so the column defaults to SQL NULL.

316 create_data: Final[dict[str, object]] = {

317 "key": body.key,

318 "value": body.value,

319 "user_id": user_id,

320 "team_id": team_id,

321 "created_by": user_api_key_dict.user_id,

322 "updated_by": user_api_key_dict.user_id,

323 }

324 if body.metadata is not None:

325 create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata)

326

327 try:

328 row: Final = await _memory_table(prisma_client).create(data=create_data)

329 except Exception as e:

330 # Key is globally unique. Any duplicate → 409.

331 if _is_unique_violation(e): 331 ↛ 336line 331 didn't jump to line 336 because the condition on line 331 was always true

332 raise HTTPException(

333 status_code=409,

334 detail=f"Memory with key '{body.key}' already exists.",

335 )

336 raise _internal_error(

337 "Error creating memory: %s",

338 e,

339 "Internal error creating memory entry.",

340 )

341

342 return _row_to_model(row)

343

344

345@router.get(

346 "/v1/memory",

347 tags=["memory management"],

348 dependencies=[Depends(user_api_key_auth)],

349 response_model=MemoryListResponse,

350)

351async def list_memory(

352 key: str | None = Query(None, description="Filter by exact key match."),

353 key_prefix: str | None = Query(

354 None,

355 description=(

356 "Filter by key prefix (Redis-style namespace scan). "

357 "Mutually exclusive with `key`; if both are provided, `key_prefix` wins."

358 ),

359 ),

360 search: str | None = Query(

361 None,

362 description=(

363 "Match entries whose key starts with this value or whose memory_id equals it. "

364 "Takes precedence over `key_prefix` and `key` when provided."

365 ),

366 ),

367 page: int = Query(1, ge=1),

368 page_size: int = Query(50, ge=1, le=500),

369 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),

370):

371 """List memory entries visible to the caller."""

372 prisma_client: Final = _require_prisma()

373

374 # AND the key filter with the visibility filter via an explicit top-level

375 # "AND": both sides can carry an "OR" key (`search`, non-admin visibility),

376 # so merging them by key would let one clobber the other and leak rows.

377 key_filter: Final = _key_filter(search=search, key_prefix=key_prefix, key=key)

378

379 vis: Final = _visibility_filter(user_api_key_dict)

380 where: Mapping[str, object] | None

381 if vis is None: 381 ↛ 383line 381 didn't jump to line 383 because the condition on line 381 was always true

382 where = key_filter

383 elif key_filter is None:

384 where = vis

385 else:

386 where = {"AND": [key_filter, vis]}

387

388 try:

389 total: Final = await _memory_table(prisma_client).count(where=where)

390 rows: Final = await _memory_table(prisma_client).find_many(

391 where=where,

392 order={"updated_at": "desc"},

393 skip=(page - 1) * page_size,

394 take=page_size,

395 )

396 except Exception as e:

397 raise _internal_error("Error listing memory: %s", e, "Internal error listing memory entries.")

398

399 return MemoryListResponse(memories=[_row_to_model(r) for r in rows], total=total)

400

401

402async def _find_memory_for_caller(

403 prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth

404) -> "prisma_models.LiteLLM_MemoryTable":

405 """Look up a memory row by key, scoped to the caller's visibility."""

406 key_filter: Final[Mapping[str, object]] = {"key": key}

407 vis: Final = _visibility_filter(user_api_key_dict)

408 where: Final[Mapping[str, object]] = key_filter if vis is None else {"AND": [key_filter, vis]}

409 rows = await _memory_table(prisma_client).find_many(where=where, take=1, order={"updated_at": "desc"})

410 if not rows:

411 raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found")

412 return rows[0]

413

414

415async def _find_visible_memory_or_none(

416 prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth

417) -> "prisma_models.LiteLLM_MemoryTable | None":

418 """The caller-visible row for `key`, or None when nothing is visible to them."""

419 try:

420 return await _find_memory_for_caller(prisma_client, key, user_api_key_dict)

421 except HTTPException as e:

422 if e.status_code == 404: 422 ↛ 424line 422 didn't jump to line 424 because the condition on line 422 was always true

423 return None

424 raise

425

426

427@router.get(

428 "/v1/memory/{key:path}",

429 tags=["memory management"],

430 dependencies=[Depends(user_api_key_auth)],

431 response_model=LiteLLM_MemoryRow,

432)

433async def get_memory(

434 key: str,

435 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),

436):

437 """Get a single memory entry by key, scoped to the caller."""

438 prisma_client: Final = _require_prisma()

439 row: Final = await _find_memory_for_caller(prisma_client, key, user_api_key_dict)

440 return _row_to_model(row)

441

442

443@router.put(

444 "/v1/memory/{key:path}",

445 tags=["memory management"],

446 dependencies=[Depends(user_api_key_auth)],

447 response_model=LiteLLM_MemoryRow,

448)

449async def upsert_memory(

450 key: str,

451 body: MemoryUpdateRequest,

452 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),

453):

454 """

455 Upsert a memory entry by key within the caller's scope.

456

457 If no row exists for (key, caller.user_id, caller.team_id), create one.

458 If one exists, update the value/metadata fields that were provided.

459 """

460 prisma_client: Final = _require_prisma()

461

462 # `metadata` is a `Json?` column. prisma-client-python has no

463 # `JsonNull`/`DbNull` sentinel for writing a true SQL NULL

464 # (RobertCraigie/prisma-client-py#714), so an explicit `metadata: null`

465 # is encoded as the JSON literal `null` instead — stored as Postgres

466 # `jsonb 'null'`, which prisma deserializes back to Python `None` on

467 # read. From a caller's perspective `PUT {"metadata": null}` clears

468 # the field (subsequent reads return `metadata: null`), matching the

469 # natural expectation. Callers wanting a strict SQL NULL must use

470 # raw SQL — there is no typed-client path.

471 #

472 # When `metadata` is omitted from the request body entirely (not in

473 # `model_fields_set`), the column is preserved as-is.

474 fields_sent: Final = body.model_fields_set

475 metadata_in_payload: Final = "metadata" in fields_sent

476

477 data: Final[dict[str, object]] = {}

478 if body.value is not None:

479 data["value"] = body.value

480 if metadata_in_payload:

481 data["metadata"] = _serialize_metadata_for_prisma(body.metadata)

482 if not data:

483 raise HTTPException(

484 status_code=400,

485 detail="Request body must include at least one of: value, metadata.",

486 )

487 data["updated_by"] = user_api_key_dict.user_id

488

489 try:

490 existing: Final = await _find_visible_memory_or_none(prisma_client, key, user_api_key_dict)

491 if existing is not None:

492 # Visibility != write authority. Make sure the caller actually

493 # owns this row (their user_id matches, or it's a pure team row in

494 # their team) — otherwise a teammate could overwrite a personal

495 # entry through the OR-based visibility filter.

496 await _assert_write_access(prisma_client, existing, user_api_key_dict)

497 row = await _memory_table(prisma_client).update(

498 where={"memory_id": existing.memory_id},

499 data=data,

500 )

501 else:

502 if body.value is None:

503 raise HTTPException(

504 status_code=400,

505 detail="Cannot create a new memory via PUT without a 'value'.",

506 )

507 # PUT-create must honor admin scope override, matching POST semantics.

508 user_id, team_id = _resolve_scope(user_api_key_dict, body.user_id, body.team_id)

509 # Omit `metadata` when None so the column defaults to SQL NULL;

510 # otherwise JSON-encode for Prisma — same pattern as

511 # `create_memory` above.

512 create_data: Final[dict[str, object]] = {

513 "key": key,

514 "value": body.value,

515 "user_id": user_id,

516 "team_id": team_id,

517 "created_by": user_api_key_dict.user_id,

518 "updated_by": user_api_key_dict.user_id,

519 }

520 if body.metadata is not None: 520 ↛ 521line 520 didn't jump to line 521 because the condition on line 520 was never true

521 create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata)

522 try:

523 row = await _memory_table(prisma_client).create(data=create_data)

524 except Exception as e:

525 # Race: a concurrent PUT/POST created the row after our check.

526 # Re-read and fall back to an update so the PUT stays idempotent

527 # instead of surfacing a 500 on a unique-violation.

528 if not _is_unique_violation(e):

529 raise

530 existing_after_race: Final = await _find_visible_memory_or_none(prisma_client, key, user_api_key_dict)

531 if existing_after_race is None:

532 # Row exists globally but isn't visible to this caller

533 # (owned by someone else). Treat as conflict.

534 raise HTTPException(

535 status_code=409,

536 detail=f"Memory with key '{key}' already exists.",

537 )

538 # Same write-authorization check as the non-race path.

539 await _assert_write_access(prisma_client, existing_after_race, user_api_key_dict)

540 row = await _memory_table(prisma_client).update(

541 where={"memory_id": existing_after_race.memory_id},

542 data=data,

543 )

544 except HTTPException:

545 raise

546 except Exception as e:

547 raise _internal_error("Error upserting memory: %s", e, "Internal error updating memory entry.")

548

549 if row is None: 549 ↛ 550line 549 didn't jump to line 550 because the condition on line 549 was never true

550 raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found")

551 return _row_to_model(row)

552

553

554@router.delete(

555 "/v1/memory/{key:path}",

556 tags=["memory management"],

557 dependencies=[Depends(user_api_key_auth)],

558 response_model=MemoryDeleteResponse,

559)

560async def delete_memory(

561 key: str,

562 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),

563):

564 """Delete a memory entry by key, scoped to the caller."""

565 prisma_client: Final = _require_prisma()

566 row: Final = await _find_memory_for_caller(prisma_client, key, user_api_key_dict)

567 # Visibility != write authority — see the upsert handler for the rationale.

568 await _assert_write_access(prisma_client, row, user_api_key_dict)

569 try:

570 deleted: Final = await _memory_table(prisma_client).delete(where={"memory_id": row.memory_id})

571 except Exception as e:

572 raise _internal_error("Error deleting memory: %s", e, "Internal error deleting memory entry.")

573

574 if deleted is None: 574 ↛ 575line 574 didn't jump to line 575 because the condition on line 574 was never true

575 raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found")

576 return MemoryDeleteResponse(key=key, deleted=True)

Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/memory/memory_endpoints.py: 64%

209 statements