mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
refactor(memory): remove obsolete policy and end-to-end scaffolding
This commit is contained in:
parent
eb72042a0a
commit
fc8637ad86
24 changed files with 34 additions and 847 deletions
|
|
@ -71,7 +71,6 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/v1/workflows/",
|
||||
"/project/",
|
||||
"/memory/",
|
||||
"/memory/v2/",
|
||||
"/mcp/",
|
||||
# Control plane (see the List Endpoints + Tables standard). Every resource
|
||||
# eventually moves under this prefix, so allowlist it once rather than
|
||||
|
|
|
|||
|
|
@ -3,24 +3,23 @@ ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "namespace" TEXT;
|
|||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_namespace_updated_at_idx"
|
||||
ON "LiteLLM_MemoryTable"("namespace", "updated_at");
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryPolicy" (
|
||||
"policy_id" TEXT NOT NULL,
|
||||
"target_type" TEXT NOT NULL,
|
||||
"target_id" TEXT NOT NULL,
|
||||
"activation" TEXT NOT NULL,
|
||||
"scope" TEXT NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_by" TEXT NOT NULL,
|
||||
CONSTRAINT "LiteLLM_MemoryPolicy_pkey" PRIMARY KEY ("policy_id")
|
||||
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "organization_id" TEXT;
|
||||
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "owner_key_id" TEXT;
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_organization_id_user_id_updated_at_idx" ON "LiteLLM_MemoryTable"("organization_id", "user_id", "updated_at");
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_team_id_updated_at_idx" ON "LiteLLM_MemoryTable"("team_id", "updated_at");
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_owner_key_id_idx" ON "LiteLLM_MemoryTable"("owner_key_id");
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryContinuation" (
|
||||
"id" TEXT NOT NULL,
|
||||
"namespace" TEXT NOT NULL,
|
||||
"key_id" TEXT NOT NULL,
|
||||
"payload" JSONB NOT NULL,
|
||||
"expires_at" TIMESTAMP(3) NOT NULL,
|
||||
CONSTRAINT "LiteLLM_MemoryContinuation_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MemoryPolicy_target_type_target_id_key"
|
||||
ON "LiteLLM_MemoryPolicy"("target_type", "target_id");
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryContinuation_namespace_key_id_expires_at_idx"
|
||||
ON "LiteLLM_MemoryContinuation"("namespace", "key_id", "expires_at");
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryPreference" (
|
||||
"subject" TEXT NOT NULL,
|
||||
"enabled" BOOLEAN NOT NULL DEFAULT false,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
CONSTRAINT "LiteLLM_MemoryPreference_pkey" PRIMARY KEY ("subject")
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryContinuation_expires_at_idx"
|
||||
ON "LiteLLM_MemoryContinuation"("expires_at");
|
||||
|
|
|
|||
|
|
@ -1,14 +0,0 @@
|
|||
CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryContinuation" (
|
||||
"id" TEXT NOT NULL,
|
||||
"namespace" TEXT NOT NULL,
|
||||
"key_id" TEXT NOT NULL,
|
||||
"payload" JSONB NOT NULL,
|
||||
"expires_at" TIMESTAMP(3) NOT NULL,
|
||||
CONSTRAINT "LiteLLM_MemoryContinuation_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryContinuation_namespace_key_id_expires_at_idx"
|
||||
ON "LiteLLM_MemoryContinuation"("namespace", "key_id", "expires_at");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryContinuation_expires_at_idx"
|
||||
ON "LiteLLM_MemoryContinuation"("expires_at");
|
||||
|
|
@ -1 +0,0 @@
|
|||
ALTER TABLE "LiteLLM_MemoryPolicy" ADD COLUMN IF NOT EXISTS "paused" BOOLEAN NOT NULL DEFAULT false;
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "organization_id" TEXT;
|
||||
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "owner_key_id" TEXT;
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_organization_id_user_id_updated_at_idx" ON "LiteLLM_MemoryTable"("organization_id", "user_id", "updated_at");
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_team_id_updated_at_idx" ON "LiteLLM_MemoryTable"("team_id", "updated_at");
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_owner_key_id_idx" ON "LiteLLM_MemoryTable"("owner_key_id");
|
||||
|
|
@ -1451,26 +1451,6 @@ model LiteLLM_MemoryTable {
|
|||
@@index([owner_key_id])
|
||||
}
|
||||
|
||||
model LiteLLM_MemoryPolicy {
|
||||
policy_id String @id
|
||||
target_type String
|
||||
target_id String
|
||||
activation String
|
||||
scope String
|
||||
paused Boolean @default(false)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
|
||||
@@unique([target_type, target_id])
|
||||
}
|
||||
|
||||
model LiteLLM_MemoryPreference {
|
||||
subject String @id
|
||||
enabled Boolean @default(false)
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
}
|
||||
|
||||
model LiteLLM_MemoryContinuation {
|
||||
id String @id
|
||||
namespace String
|
||||
|
|
|
|||
|
|
@ -932,9 +932,6 @@ class LiteLLMRoutes(enum.Enum):
|
|||
# updating this list — the default-allow behavior covers it automatically.
|
||||
admin_viewer_routes = (
|
||||
[
|
||||
"/memory/v2/settings",
|
||||
"/memory/v2/status",
|
||||
"/memory/v2/entries",
|
||||
"/user/list",
|
||||
"/user/available_users",
|
||||
"/user/available_roles",
|
||||
|
|
|
|||
|
|
@ -203,7 +203,7 @@ async def _assert_write_access(
|
|||
)
|
||||
|
||||
|
||||
async def is_memory_team_admin(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool:
|
||||
async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool:
|
||||
"""
|
||||
True if the caller is a team admin of `team_id`, or an org admin for the
|
||||
team's organization. Mirrors the auth pattern used by team-management
|
||||
|
|
@ -591,4 +591,3 @@ async def delete_memory(
|
|||
|
||||
|
||||
_require_prisma = require_memory_prisma
|
||||
_is_team_admin_for = is_memory_team_admin
|
||||
|
|
|
|||
|
|
@ -70,7 +70,6 @@ class MemoryIdentity:
|
|||
key_id: str | None
|
||||
user_id: str | None
|
||||
team_id: str | None
|
||||
project_id: str | None
|
||||
organization_id: str | None
|
||||
read_only: bool
|
||||
role: str | None = None
|
||||
|
|
@ -89,7 +88,6 @@ class MemoryIdentity:
|
|||
key_id=key_id,
|
||||
user_id=auth.user_id,
|
||||
team_id=None if dashboard else auth.team_id,
|
||||
project_id=auth.project_id,
|
||||
organization_id=auth.org_id,
|
||||
read_only=auth.user_role
|
||||
in (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY),
|
||||
|
|
|
|||
|
|
@ -70,10 +70,10 @@ def _before(cursor: tuple[datetime, str] | None) -> Mapping[str, object]:
|
|||
|
||||
|
||||
class MemoryStore:
|
||||
def __init__(self, prisma_client: object, access: MemoryAccess, *, actor: str | None = None) -> None:
|
||||
def __init__(self, prisma_client: object, access: MemoryAccess) -> None:
|
||||
self.prisma_client = memory_primary_client(prisma_client)
|
||||
self.access = access
|
||||
self.actor = actor or access.identity.user_id or access.identity.key_id
|
||||
self.actor = access.identity.user_id or access.identity.key_id
|
||||
self.table = MemoryRepository(self.prisma_client).table
|
||||
|
||||
async def authorize(self, *, write: bool = False, require_active: bool = True) -> MemoryAccess:
|
||||
|
|
|
|||
|
|
@ -1451,26 +1451,6 @@ model LiteLLM_MemoryTable {
|
|||
@@index([owner_key_id])
|
||||
}
|
||||
|
||||
model LiteLLM_MemoryPolicy {
|
||||
policy_id String @id
|
||||
target_type String
|
||||
target_id String
|
||||
activation String
|
||||
scope String
|
||||
paused Boolean @default(false)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
|
||||
@@unique([target_type, target_id])
|
||||
}
|
||||
|
||||
model LiteLLM_MemoryPreference {
|
||||
subject String @id
|
||||
enabled Boolean @default(false)
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
}
|
||||
|
||||
model LiteLLM_MemoryContinuation {
|
||||
id String @id
|
||||
namespace String
|
||||
|
|
|
|||
|
|
@ -98,7 +98,7 @@ class LinkedSpendResetWrites:
|
|||
table: BatchTable
|
||||
|
||||
def queue_spend_zero(self, where: Mapping[str, object]) -> None:
|
||||
self.table.update_many(where=where, data={"spend": 0}) # mutable-ok: Callable type parameter syntax.
|
||||
self.table.update_many(where=where, data={"spend": 0})
|
||||
|
||||
def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None:
|
||||
"""``decrement`` rather than a read-then-set, so spend written between the
|
||||
|
|
@ -116,9 +116,7 @@ class BudgetWindowWrites:
|
|||
def queue_window_advance(self, budget_id: str, budget_reset_at: datetime) -> None:
|
||||
"""``update_many`` so a tier deleted between the read and the commit is a
|
||||
no-op row count instead of a P2025 that aborts the whole chunk."""
|
||||
self.table.update_many(
|
||||
where={"budget_id": budget_id}, data={"budget_reset_at": budget_reset_at}
|
||||
) # mutable-ok: Callable type parameter syntax.
|
||||
self.table.update_many(where={"budget_id": budget_id}, data={"budget_reset_at": budget_reset_at})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
|
|||
|
|
@ -89,12 +89,6 @@ class MemorySearch(BaseModel):
|
|||
offset: int = Field(default=0, ge=0, le=10000)
|
||||
|
||||
|
||||
class MemoryRead(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
memory_id: str = Field(min_length=1, max_length=64)
|
||||
|
||||
|
||||
class MemoryObservation(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
|
|
|
|||
|
|
@ -1451,26 +1451,6 @@ model LiteLLM_MemoryTable {
|
|||
@@index([owner_key_id])
|
||||
}
|
||||
|
||||
model LiteLLM_MemoryPolicy {
|
||||
policy_id String @id
|
||||
target_type String
|
||||
target_id String
|
||||
activation String
|
||||
scope String
|
||||
paused Boolean @default(false)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
|
||||
@@unique([target_type, target_id])
|
||||
}
|
||||
|
||||
model LiteLLM_MemoryPreference {
|
||||
subject String @id
|
||||
enabled Boolean @default(false)
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
}
|
||||
|
||||
model LiteLLM_MemoryContinuation {
|
||||
id String @id
|
||||
namespace String
|
||||
|
|
|
|||
|
|
@ -90,15 +90,6 @@
|
|||
- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"}
|
||||
- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"}
|
||||
- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"}
|
||||
- {id: mgmt.memory_v2.gateway.capture_recall, module: mgmt, tier: P0, surface: api, assertions: [capture_recall], source: "memory/gateway.py", rationale: "Existing clients store and recall across fresh conversations and native streams"}
|
||||
- {id: mgmt.memory_v2.settings.enrollment, module: mgmt, tier: P0, surface: api, assertions: [enrollment], source: "memory/policy.py", rationale: "Admin enrollment follows owners across keys; disabling stops automatic memory and preserves authorized inspection"}
|
||||
- {id: mgmt.memory_v2.entries.isolation, module: mgmt, tier: P0, surface: api, assertions: [isolation], source: "memory/store.py", rationale: "Owner keys share records while unrelated users and V1 remain isolated"}
|
||||
- {id: mgmt.memory_v2.settings.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "memory/management.py", rationale: "Only proxy administrators can enable memory for everyone or selected users"}
|
||||
- {id: mgmt.memory_v2.entries.correction_delete, module: mgmt, tier: P0, surface: api, assertions: [correction_delete], source: "memory/store.py", rationale: "Corrections reject stale revisions and deletion removes facts from future recall"}
|
||||
- {id: mgmt.memory_v2.gateway.client_tools, module: mgmt, tier: P0, surface: api, assertions: [client_tools], source: "memory/gateway.py", rationale: "Application tool calls and their continuations stay under client control"}
|
||||
- {id: mgmt.memory_v2.gateway.billing, module: mgmt, tier: P0, surface: api, assertions: [billing], source: "memory/gateway.py", rationale: "Preparation and final answer each create distinct billed calls under the authenticated key, user, and team"}
|
||||
|
||||
- {id: mgmt.memory_v2.entries.team_permissions, module: mgmt, tier: P0, surface: api, assertions: [team_permissions], source: "memory/policy.py", rationale: "Existing team grants allow reads without writes; log grants do not confer memory access; revocation takes effect"}
|
||||
|
||||
- {id: mgmt.user.jwt.database_roles, module: mgmt, tier: P0, surface: api, assertions: [database_roles], source: "auth/handle_jwt.py", rationale: "User-only JWT subjects retain their seeded database roles and memberships"}
|
||||
- {id: mgmt.key.jwt.viewer_denied, module: mgmt, tier: P0, surface: api, assertions: [viewer_denied], source: "auth/route_checks.py", rationale: "An admin viewer can read a key but cannot update it or change stored state"}
|
||||
|
|
|
|||
|
|
@ -1,109 +0,0 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
from e2e_http import NoBody, Result, unwrap
|
||||
from models import (
|
||||
MemoryCaptureBody,
|
||||
MemoryEntriesData,
|
||||
MemoryEntryData,
|
||||
MemoryEntryParams,
|
||||
MemorySettingsBody,
|
||||
MemoryStatusData,
|
||||
)
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MemoryClient:
|
||||
proxy: ProxyClient
|
||||
|
||||
def settings(self) -> MemorySettingsBody:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/memory/v2/settings",
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=MemorySettingsBody,
|
||||
)
|
||||
)
|
||||
|
||||
def set_settings(self, body: MemorySettingsBody, *, caller: str | None = None) -> Result[MemorySettingsBody]:
|
||||
return self.proxy.transport.put(
|
||||
"/memory/v2/settings",
|
||||
headers=self.proxy.transport.bearer(caller) if caller else self.proxy.transport.master,
|
||||
json=body,
|
||||
response_type=MemorySettingsBody,
|
||||
)
|
||||
|
||||
def read(self, key: str, memory_id: str) -> Result[MemoryEntryData]:
|
||||
return self.proxy.transport.get(
|
||||
f"/memory/v2/entries/{memory_id}",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
params=NoBody(),
|
||||
response_type=MemoryEntryData,
|
||||
)
|
||||
|
||||
def update(self, key: str, memory_id: str, body: MemoryCaptureBody) -> Result[MemoryEntryData]:
|
||||
return self.proxy.transport.put(
|
||||
f"/memory/v2/entries/{memory_id}",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=body,
|
||||
response_type=MemoryEntryData,
|
||||
)
|
||||
|
||||
def status(self, key: str) -> MemoryStatusData:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/memory/v2/status",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
params=NoBody(),
|
||||
response_type=MemoryStatusData,
|
||||
)
|
||||
)
|
||||
|
||||
def entries(self, key: str, params: MemoryEntryParams = MemoryEntryParams()) -> list[MemoryEntryData]:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/memory/v2/entries",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
params=params,
|
||||
response_type=MemoryEntriesData,
|
||||
)
|
||||
).root
|
||||
|
||||
def capture(self, key: str, body: MemoryCaptureBody) -> Result[MemoryEntryData]:
|
||||
return self.proxy.transport.post(
|
||||
"/memory/v2/entries",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=body,
|
||||
response_type=MemoryEntryData,
|
||||
)
|
||||
|
||||
def delete_entry(self, key: str, memory_id: str) -> Result[NoBody]:
|
||||
return self.proxy.transport.delete(
|
||||
f"/memory/v2/entries/{memory_id}",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def cleanup_user_entries(self, user_id: str) -> None:
|
||||
while True:
|
||||
page = unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/memory/v2/entries",
|
||||
headers=self.proxy.transport.master,
|
||||
params=MemoryEntryParams(user_id=user_id),
|
||||
response_type=MemoryEntriesData,
|
||||
)
|
||||
).root
|
||||
if not page:
|
||||
return
|
||||
for entry in page:
|
||||
unwrap(
|
||||
self.proxy.transport.delete(
|
||||
f"/memory/v2/entries/{entry.memory_id}",
|
||||
headers=self.proxy.transport.master,
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
|
|
@ -1,440 +0,0 @@
|
|||
import fcntl
|
||||
import hashlib
|
||||
import os
|
||||
import tempfile
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import NoBody, Success, unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from management_client import ManagementClient
|
||||
from memory_client import MemoryClient
|
||||
from models import (
|
||||
AnthropicMessagesBody,
|
||||
ChatAssistantTurn,
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
ChatTool,
|
||||
ChatToolFunction,
|
||||
ChatToolResultTurn,
|
||||
KeyGenerateBody,
|
||||
LiteLLMParamsBody,
|
||||
MemoryCaptureBody,
|
||||
MemoryEntryParams,
|
||||
MemoryLegacyParams,
|
||||
MemoryLegacyRows,
|
||||
MemoryResponsesBody,
|
||||
MemorySettingsBody,
|
||||
MemoryStreamEvent,
|
||||
MemoryTeamPermissionBody,
|
||||
MemoryWireResponse,
|
||||
TeamNewBody,
|
||||
UserNewBody,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MemorySubjects:
|
||||
owner: str
|
||||
sibling: str
|
||||
outsider: str
|
||||
user_id: str
|
||||
team_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MemoryModels:
|
||||
chat: str
|
||||
messages: str
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_models(client: ManagementClient, resources: ResourceManager) -> MemoryModels:
|
||||
def register(name: str, model: str, credential: str) -> str:
|
||||
alias: Final = f"e2e-memory-{name}-{unique_marker()}"
|
||||
identifier: Final = client.proxy.create_model(
|
||||
alias,
|
||||
LiteLLMParamsBody(
|
||||
model=model,
|
||||
api_key=credential,
|
||||
api_base=os.environ.get("E2E_MEMORY_API_BASE"),
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: client.proxy.delete_model(identifier))
|
||||
return alias
|
||||
|
||||
return MemoryModels(
|
||||
chat=register(
|
||||
"chat", os.environ.get("E2E_MEMORY_CHAT_MODEL", "openai/gpt-5.6-sol"), "os.environ/OPENAI_API_KEY"
|
||||
),
|
||||
messages=register(
|
||||
"messages",
|
||||
os.environ.get("E2E_MEMORY_MESSAGES_MODEL", "anthropic/claude-haiku-4-5"),
|
||||
"os.environ/ANTHROPIC_API_KEY",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory(client: ManagementClient) -> Iterator[MemoryClient]:
|
||||
# Serialize global configuration changes across local pytest workers.
|
||||
with (Path(tempfile.gettempdir()) / "litellm-memory-v2-e2e.lock").open("w") as lock:
|
||||
fcntl.flock(lock, fcntl.LOCK_EX)
|
||||
memory = MemoryClient(client.proxy)
|
||||
original = memory.settings()
|
||||
try:
|
||||
yield memory
|
||||
finally:
|
||||
unwrap(memory.set_settings(original))
|
||||
fcntl.flock(lock, fcntl.LOCK_UN)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def subjects(client: ManagementClient, memory: MemoryClient, resources: ResourceManager) -> MemorySubjects:
|
||||
marker: Final = unique_marker()
|
||||
owner: Final = client.create_user(
|
||||
UserNewBody(
|
||||
user_email=f"memory-owner-{marker}@example.invalid", user_role="internal_user", auto_create_key=False
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: client.delete_user_strict(owner))
|
||||
other: Final = client.create_user(
|
||||
UserNewBody(
|
||||
user_email=f"memory-other-{marker}@example.invalid", user_role="internal_user", auto_create_key=False
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: client.delete_user_strict(other))
|
||||
team: Final = client.create_team(TeamNewBody(team_alias=f"memory-{marker}"))
|
||||
resources.defer(lambda: client.delete_team(team))
|
||||
client.add_team_member(team, owner)
|
||||
client.add_team_member(team, other)
|
||||
|
||||
def create_key(user: str) -> str:
|
||||
key: Final = unwrap(
|
||||
client.generate_key(KeyGenerateBody(user_id=user, team_id=team, models=[], max_parallel_requests=1))
|
||||
).key
|
||||
resources.defer(lambda: client.delete_key_strict(key))
|
||||
return key
|
||||
|
||||
keys: Final = tuple(create_key(user) for user in (owner, owner, other))
|
||||
unwrap(memory.set_settings(MemorySettingsBody(enabled=True, everyone=False, user_ids=[owner, other])))
|
||||
resources.defer(lambda: memory.cleanup_user_entries(owner))
|
||||
resources.defer(lambda: memory.cleanup_user_entries(other))
|
||||
return MemorySubjects(owner=keys[0], sibling=keys[1], outsider=keys[2], user_id=owner, team_id=team)
|
||||
|
||||
|
||||
def _fact(marker: str) -> MemoryCaptureBody:
|
||||
return MemoryCaptureBody(
|
||||
key=f"release-{marker}",
|
||||
title="Demo project codename",
|
||||
content=f"The demo project codename is {marker}",
|
||||
evidence="Synthetic fact supplied by the authenticated test user",
|
||||
)
|
||||
|
||||
|
||||
def _assert_denied(result: object) -> None:
|
||||
assert not isinstance(result, Success), "An unauthorized memory operation succeeded"
|
||||
assert "403" in str(result) or "unauthorized" in str(result).lower() or "404" in str(result), result
|
||||
|
||||
|
||||
class TestMemoryV2:
|
||||
@pytest.mark.covers("mgmt.memory_v2.gateway.capture_recall")
|
||||
@pytest.mark.parametrize("endpoint", ["chat", "responses", "messages"])
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
def test_gateway_stores_and_recalls_without_client_memory_tools(
|
||||
self,
|
||||
client: ManagementClient,
|
||||
memory: MemoryClient,
|
||||
subjects: MemorySubjects,
|
||||
memory_models: MemoryModels,
|
||||
endpoint: str,
|
||||
stream: bool,
|
||||
) -> None:
|
||||
marker: Final = f"copper-{unique_marker()}"
|
||||
seed: Final = unwrap(
|
||||
client.proxy.chat(
|
||||
subjects.owner,
|
||||
ChatBody(
|
||||
model=memory_models.chat,
|
||||
max_tokens=1200,
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
content=f"Remember this durable preference for future conversations: my demo project codename is {marker}. Confirm briefly.",
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
assert seed.choices
|
||||
stored: Final = memory.entries(subjects.owner)
|
||||
assert any(marker in entry.content for entry in stored), stored
|
||||
prompt: Final = "What is my demo project codename? Return the exact word only."
|
||||
model: Final = memory_models.messages if endpoint == "messages" else memory_models.chat
|
||||
body: Final = (
|
||||
AnthropicMessagesBody(
|
||||
model=model, messages=[ChatMessage(role="user", content=prompt)], max_tokens=1200, stream=stream
|
||||
)
|
||||
if endpoint == "messages"
|
||||
else MemoryResponsesBody(model=model, input=prompt, stream=stream)
|
||||
if endpoint == "responses"
|
||||
else ChatBody(
|
||||
model=model, messages=[ChatMessage(role="user", content=prompt)], max_tokens=1200, stream=stream
|
||||
)
|
||||
)
|
||||
path: Final = {"messages": "/v1/messages", "responses": "/v1/responses", "chat": "/v1/chat/completions"}[
|
||||
endpoint
|
||||
]
|
||||
response: Final = client.proxy.transport.send(
|
||||
path, headers=client.proxy.transport.bearer(subjects.owner), json=body, stream=stream
|
||||
)
|
||||
assert response.status_code == 200, response.body
|
||||
assert response.stream_error is None, response.stream_error
|
||||
output: Final = (
|
||||
"".join(MemoryStreamEvent.model_validate_json(event).text for event in response.stream_events)
|
||||
if stream
|
||||
else response.body
|
||||
)
|
||||
assert marker in output, output
|
||||
if stream:
|
||||
assert response.is_streaming
|
||||
assert response.chunks > 1
|
||||
events: Final = tuple(MemoryStreamEvent.model_validate_json(event) for event in response.stream_events)
|
||||
assert not any(event.has_memory_tools for event in events)
|
||||
if endpoint == "responses":
|
||||
assert all(event.response.instructions is None for event in events if event.response)
|
||||
else:
|
||||
public: Final = MemoryWireResponse.model_validate_json(response.body)
|
||||
assert not public.has_memory_tools
|
||||
if endpoint == "responses":
|
||||
assert public.instructions is None
|
||||
assert any(marker in entry.content for entry in memory.entries(subjects.sibling))
|
||||
assert memory.entries(subjects.outsider) == []
|
||||
|
||||
@pytest.mark.covers("mgmt.memory_v2.settings.enrollment")
|
||||
def test_admin_enrollment_follows_user_and_disable_preserves_dashboard(
|
||||
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
|
||||
) -> None:
|
||||
unwrap(memory.set_settings(MemorySettingsBody()))
|
||||
assert not memory.status(subjects.owner).active
|
||||
marker = f"unsaved-{unique_marker()}"
|
||||
unwrap(
|
||||
client.proxy.chat(
|
||||
subjects.owner,
|
||||
ChatBody(
|
||||
model=memory_models.chat,
|
||||
max_tokens=100,
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user", content=f"Remember that my verification word is {marker}. Confirm briefly."
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
assert memory.entries(subjects.owner) == []
|
||||
unwrap(memory.set_settings(MemorySettingsBody(enabled=True, everyone=False, user_ids=[subjects.user_id])))
|
||||
assert memory.status(subjects.owner).active and memory.status(subjects.sibling).active
|
||||
assert not memory.status(subjects.outsider).active
|
||||
saved = unwrap(memory.capture(subjects.owner, _fact(unique_marker())))
|
||||
unwrap(memory.set_settings(MemorySettingsBody()))
|
||||
assert not memory.status(subjects.owner).active
|
||||
assert unwrap(memory.read(subjects.sibling, saved.memory_id)).content == saved.content
|
||||
_assert_denied(memory.capture(subjects.owner, _fact(unique_marker())))
|
||||
|
||||
@pytest.mark.covers("mgmt.memory_v2.entries.isolation")
|
||||
def test_owner_keys_share_but_other_users_and_legacy_api_do_not(
|
||||
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects
|
||||
) -> None:
|
||||
saved = unwrap(memory.capture(subjects.owner, _fact(unique_marker())))
|
||||
assert memory.entries(subjects.sibling)[0].memory_id == saved.memory_id
|
||||
assert memory.entries(subjects.outsider) == []
|
||||
assert memory.entries(subjects.outsider, MemoryEntryParams(user_id=subjects.user_id)) == []
|
||||
_assert_denied(memory.read(subjects.outsider, saved.memory_id))
|
||||
_assert_denied(memory.delete_entry(subjects.outsider, saved.memory_id))
|
||||
legacy = client.proxy.transport.get(
|
||||
"/v1/memory",
|
||||
headers=client.proxy.transport.bearer(subjects.sibling),
|
||||
params=MemoryLegacyParams(),
|
||||
response_type=MemoryLegacyRows,
|
||||
)
|
||||
if isinstance(legacy, Success):
|
||||
assert saved.memory_id not in [row.memory_id for row in legacy.data.memories]
|
||||
assert memory.entries(subjects.owner)[0].content == saved.content
|
||||
|
||||
@pytest.mark.covers("mgmt.memory_v2.settings.admin_only")
|
||||
def test_members_cannot_enable_memory(self, memory: MemoryClient, subjects: MemorySubjects) -> None:
|
||||
before = memory.settings()
|
||||
_assert_denied(memory.set_settings(MemorySettingsBody(enabled=True), caller=subjects.owner))
|
||||
assert memory.settings() == before
|
||||
|
||||
@pytest.mark.covers("mgmt.memory_v2.entries.team_permissions")
|
||||
def test_delegated_team_reads_allow_recall_but_not_edit_and_can_be_revoked(
|
||||
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
|
||||
) -> None:
|
||||
marker = f"team-{unique_marker()}"
|
||||
saved = unwrap(memory.capture(subjects.owner, _fact(marker)))
|
||||
assert memory.entries(subjects.outsider) == []
|
||||
for permissions, visible in ((["/spend/logs"], False), (["/memory/v2/entries"], True), ([], False)):
|
||||
unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/team/permissions_update",
|
||||
headers=client.proxy.transport.master,
|
||||
json=MemoryTeamPermissionBody(team_id=subjects.team_id, team_member_permissions=permissions),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
entries = memory.entries(subjects.outsider, MemoryEntryParams(team_id=subjects.team_id))
|
||||
assert bool(entries) is visible
|
||||
if visible:
|
||||
assert entries[0].memory_id == saved.memory_id and not entries[0].can_edit
|
||||
assert unwrap(memory.read(subjects.outsider, saved.memory_id)).content == saved.content
|
||||
_assert_denied(memory.update(subjects.outsider, saved.memory_id, _fact(unique_marker())))
|
||||
_assert_denied(memory.delete_entry(subjects.outsider, saved.memory_id))
|
||||
recalled = unwrap(
|
||||
client.proxy.chat(
|
||||
subjects.outsider,
|
||||
ChatBody(
|
||||
model=memory_models.chat,
|
||||
max_tokens=1200,
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
content="Search the team's memories for the demo project codename and return it exactly.",
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
assert marker in recalled.model_dump_json()
|
||||
else:
|
||||
_assert_denied(memory.read(subjects.outsider, saved.memory_id))
|
||||
|
||||
@pytest.mark.covers("mgmt.memory_v2.entries.correction_delete")
|
||||
def test_corrections_require_current_revision_and_deleted_memory_is_not_recalled(
|
||||
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
|
||||
) -> None:
|
||||
original: Final = _fact(unique_marker())
|
||||
saved: Final = unwrap(memory.capture(subjects.owner, original))
|
||||
corrected_word: Final = f"corrected-{unique_marker()}"
|
||||
correction: Final = MemoryCaptureBody(
|
||||
key=original.key,
|
||||
title=original.title,
|
||||
content=f"The demo project codename is {corrected_word}",
|
||||
evidence="The user corrected the earlier word",
|
||||
expected_revision=saved.updated_at,
|
||||
)
|
||||
updated: Final = unwrap(memory.capture(subjects.owner, correction))
|
||||
assert updated.memory_id == saved.memory_id
|
||||
stale: Final = memory.capture(
|
||||
subjects.owner, original.model_copy(update={"expected_revision": saved.updated_at})
|
||||
)
|
||||
assert not isinstance(stale, Success)
|
||||
assert "409" in str(stale), stale
|
||||
assert memory.entries(subjects.owner)[0].content == correction.content
|
||||
unwrap(memory.delete_entry(subjects.owner, saved.memory_id))
|
||||
assert memory.entries(subjects.owner) == []
|
||||
response: Final = unwrap(
|
||||
client.proxy.chat(
|
||||
subjects.owner,
|
||||
ChatBody(
|
||||
model=memory_models.chat,
|
||||
max_tokens=200,
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user", content="What is my demo project codename? If it is unknown, say unknown."
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
assert corrected_word not in response.model_dump_json()
|
||||
|
||||
@pytest.mark.covers("mgmt.memory_v2.gateway.client_tools")
|
||||
def test_client_tool_call_and_continuation_remain_owned_by_client(
|
||||
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
|
||||
) -> None:
|
||||
marker: Final = unique_marker()
|
||||
unwrap(memory.capture(subjects.owner, _fact(marker)))
|
||||
tool: Final = ChatTool(
|
||||
function=ChatToolFunction(
|
||||
name="verify_release",
|
||||
description="Verify a release using the user's verification word",
|
||||
parameters={"type": "object", "properties": {"word": {"type": "string"}}, "required": ["word"]},
|
||||
)
|
||||
)
|
||||
prompt: Final = ChatMessage(
|
||||
role="user",
|
||||
content="Use verify_release with my demo project codename. After the tool returns, save its verification result in memory with the tool as your evidence, then report it.",
|
||||
)
|
||||
response: Final = unwrap(
|
||||
client.proxy.chat(
|
||||
subjects.owner,
|
||||
ChatBody(
|
||||
model=memory_models.chat, messages=[prompt], tools=[tool], tool_choice="required", max_tokens=1200
|
||||
),
|
||||
)
|
||||
)
|
||||
message: Final = response.choices[0].message
|
||||
assert message is not None
|
||||
calls: Final = message.tool_calls
|
||||
assert calls and len(calls) == 1
|
||||
call: Final = calls[0]
|
||||
assert call.function.name == "verify_release"
|
||||
assert marker in (call.function.arguments or "")
|
||||
assert call.id
|
||||
result_marker: Final = f"verified-{unique_marker()}"
|
||||
followup: Final = unwrap(
|
||||
client.proxy.chat(
|
||||
subjects.owner,
|
||||
ChatBody(
|
||||
model=memory_models.chat,
|
||||
max_tokens=1200,
|
||||
tools=[tool],
|
||||
messages=[
|
||||
prompt,
|
||||
ChatAssistantTurn(
|
||||
content=message.content, reasoning_content=message.reasoning_content, tool_calls=calls
|
||||
),
|
||||
ChatToolResultTurn(tool_call_id=call.id, content=result_marker),
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
assert result_marker in followup.model_dump_json()
|
||||
assert any(result_marker in entry.content for entry in memory.entries(subjects.owner))
|
||||
|
||||
@pytest.mark.covers("mgmt.memory_v2.gateway.billing")
|
||||
def test_memory_tool_rounds_are_charged_once_to_the_calling_key(
|
||||
self, client: ManagementClient, memory: MemoryClient, subjects: MemorySubjects, memory_models: MemoryModels
|
||||
) -> None:
|
||||
marker: Final = unique_marker()
|
||||
response: Final = unwrap(
|
||||
client.proxy.chat(
|
||||
subjects.owner,
|
||||
ChatBody(
|
||||
model=memory_models.chat,
|
||||
max_tokens=1200,
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user", content=f"Remember that my demo project codename is {marker}. Confirm briefly."
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
)
|
||||
assert response.choices
|
||||
assert any(marker in row.content for row in memory.entries(subjects.owner))
|
||||
rows: Final = client.proxy.poll_logs_for_key(subjects.owner, min_rows=2)
|
||||
assert 2 <= len(rows) <= 8, rows
|
||||
assert len({row.request_id for row in rows}) == len(rows), rows
|
||||
assert all(row.api_key == hashlib.sha256(subjects.owner.encode()).hexdigest() for row in rows), rows
|
||||
assert all(row.user == subjects.user_id and row.team_id == subjects.team_id for row in rows), rows
|
||||
assert all(row.prompt_tokens and row.completion_tokens for row in rows), rows
|
||||
assert all(row.spend is not None and row.spend > 0 for row in rows), rows
|
||||
|
|
@ -69,7 +69,6 @@ class ObjectPermission(BaseModel):
|
|||
|
||||
|
||||
class KeyGenerateBody(BaseModel):
|
||||
max_parallel_requests: int | None = None
|
||||
models: list[str] = []
|
||||
duration: str | None = None
|
||||
max_budget: float | None = None
|
||||
|
|
@ -1327,154 +1326,3 @@ class ReadinessDetailsResponse(ReadinessResponse):
|
|||
|
||||
litellm_version: str | None = None
|
||||
success_callbacks: list[str] = []
|
||||
|
||||
|
||||
class MemorySettingsBody(BaseModel):
|
||||
enabled: bool = False
|
||||
everyone: bool = True
|
||||
user_ids: list[str] = []
|
||||
|
||||
|
||||
class MemoryStatusData(BaseModel):
|
||||
active: bool
|
||||
enabled: bool
|
||||
user_id: str | None = None
|
||||
team_ids: list[str] = []
|
||||
admin_view: bool = False
|
||||
|
||||
|
||||
class MemoryEntryParams(BaseModel):
|
||||
query: str = ""
|
||||
limit: int = 20
|
||||
user_id: str | None = None
|
||||
team_id: str | None = None
|
||||
offset: int = 0
|
||||
before_updated_at: str | None = None
|
||||
before_memory_id: str | None = None
|
||||
|
||||
|
||||
class MemoryTeamPermissionBody(BaseModel):
|
||||
team_id: str
|
||||
team_member_permissions: list[str]
|
||||
|
||||
|
||||
class MemoryCaptureBody(BaseModel):
|
||||
key: str
|
||||
title: str
|
||||
content: str
|
||||
evidence: str
|
||||
expected_revision: str | None = None
|
||||
|
||||
|
||||
class MemoryEntryData(BaseModel):
|
||||
user_id: str | None = None
|
||||
team_id: str | None = None
|
||||
can_edit: bool = False
|
||||
actor: str | None = None
|
||||
memory_id: str
|
||||
key: str
|
||||
title: str
|
||||
content: str
|
||||
evidence: str
|
||||
updated_at: str
|
||||
|
||||
|
||||
class MemoryEntriesData(RootModel[list[MemoryEntryData]]):
|
||||
pass
|
||||
|
||||
|
||||
class MemoryLegacyParams(BaseModel):
|
||||
key_prefix: str = "memory-v2:"
|
||||
page: int = 1
|
||||
page_size: int = 500
|
||||
|
||||
|
||||
class MemoryLegacyRow(BaseModel):
|
||||
memory_id: str
|
||||
key: str
|
||||
user_id: str | None = None
|
||||
|
||||
|
||||
class MemoryLegacyRows(BaseModel):
|
||||
memories: list[MemoryLegacyRow]
|
||||
total: int
|
||||
|
||||
|
||||
class MemoryResponsesBody(BaseModel):
|
||||
model: str
|
||||
input: str
|
||||
stream: bool
|
||||
max_output_tokens: int = 1200
|
||||
store: bool = False
|
||||
cache: dict[str, bool] = {"no-cache": True}
|
||||
|
||||
|
||||
class MemoryWireTool(BaseModel):
|
||||
name: str | None = None
|
||||
function: ToolCallFunction = ToolCallFunction()
|
||||
|
||||
@property
|
||||
def is_gateway_memory(self) -> bool:
|
||||
return (self.name or self.function.name or "").startswith("litellm_memory_")
|
||||
|
||||
|
||||
class MemoryStreamDelta(BaseModel):
|
||||
content: str | None = None
|
||||
text: str | None = None
|
||||
tool_calls: tuple[MemoryWireTool, ...] | None = None
|
||||
|
||||
|
||||
class MemoryStreamChoice(BaseModel):
|
||||
delta: MemoryStreamDelta = MemoryStreamDelta()
|
||||
message: MemoryStreamDelta = MemoryStreamDelta()
|
||||
|
||||
|
||||
class MemoryWireResponse(BaseModel):
|
||||
instructions: str | None = None
|
||||
tools: tuple[MemoryWireTool, ...] = ()
|
||||
output: tuple[MemoryWireTool, ...] = ()
|
||||
content: tuple[MemoryWireTool, ...] = ()
|
||||
choices: tuple[MemoryStreamChoice, ...] = ()
|
||||
|
||||
@property
|
||||
def has_memory_tools(self) -> bool:
|
||||
return any(
|
||||
tool.is_gateway_memory
|
||||
for tool in (
|
||||
*self.tools,
|
||||
*self.output,
|
||||
*self.content,
|
||||
*(
|
||||
tool
|
||||
for choice in self.choices
|
||||
for part in (choice.delta, choice.message)
|
||||
for tool in part.tool_calls or ()
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MemoryStreamEvent(BaseModel):
|
||||
delta: MemoryStreamDelta | str | None = None
|
||||
choices: list[MemoryStreamChoice] = []
|
||||
response: MemoryWireResponse | None = None
|
||||
item: MemoryWireTool = MemoryWireTool()
|
||||
content_block: MemoryWireTool = MemoryWireTool()
|
||||
|
||||
@property
|
||||
def has_memory_tools(self) -> bool:
|
||||
return bool(
|
||||
self.response
|
||||
and self.response.has_memory_tools
|
||||
or self.item.is_gateway_memory
|
||||
or self.content_block.is_gateway_memory
|
||||
or any(tool.is_gateway_memory for choice in self.choices for tool in choice.delta.tool_calls or ())
|
||||
)
|
||||
|
||||
@property
|
||||
def text(self) -> str:
|
||||
if isinstance(self.delta, str):
|
||||
return self.delta
|
||||
if self.delta:
|
||||
return self.delta.text or self.delta.content or ""
|
||||
return "".join(choice.delta.content or "" for choice in self.choices)
|
||||
|
|
|
|||
|
|
@ -305,7 +305,6 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = (
|
|||
"/global",
|
||||
"/config",
|
||||
"/guardrails",
|
||||
"/memory/v2/settings",
|
||||
"/openapi.json",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1182,10 +1182,8 @@ async def test_end_user_jwt_auth(monkeypatch):
|
|||
"llm_router",
|
||||
router,
|
||||
)
|
||||
# No memory policy exists for this authenticated JWT user.
|
||||
memory_db = MagicMock()
|
||||
memory_db.db.litellm_memorypolicy.find_many = AsyncMock(return_value=[])
|
||||
memory_db.db.litellm_memorypreference.find_unique = AsyncMock(return_value=None)
|
||||
memory_db.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", memory_db)
|
||||
setattr(litellm.proxy.proxy_server, "jwt_handler", jwt_handler)
|
||||
from litellm.proxy.proxy_server import cost_tracking
|
||||
|
|
|
|||
|
|
@ -503,10 +503,8 @@ def test_custom_logger_failure_handler(mock_acompletion, client_no_auth):
|
|||
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
# The request has no memory policy; keep the database edge inert.
|
||||
memory_db = MagicMock()
|
||||
memory_db.db.litellm_memorypolicy.find_many = AsyncMock(return_value=[])
|
||||
memory_db.db.litellm_memorypreference.find_unique = AsyncMock(return_value=None)
|
||||
memory_db.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", memory_db)
|
||||
setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)
|
||||
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ from litellm.proxy.memory.store import MemoryStore
|
|||
from litellm.types.memory_v2 import MemoryCapture, MemorySearch, MemorySettings
|
||||
|
||||
_NOW: Final = datetime(2026, 9, 12, tzinfo=timezone.utc)
|
||||
_IDENTITY: Final = MemoryIdentity("a" * 64, "owner", "team", "project", "org", False)
|
||||
_IDENTITY: Final = MemoryIdentity("a" * 64, "owner", "team", "org", False)
|
||||
_SETTINGS: Final = MemorySettings(enabled=True)
|
||||
_CAPTURE: Final = MemoryCapture(key="demo", title="Demo", content="Use port 8347", evidence="User selected this port")
|
||||
|
||||
|
|
@ -162,7 +162,7 @@ async def test_response_deletion_preserves_auth_paths_and_retry_state(prisma_edg
|
|||
"query_string": b"api-version=test",
|
||||
}
|
||||
)
|
||||
identity = MemoryIdentity("a" * 64, "owner", "team", "project", "org", outcome == "readonly")
|
||||
identity = MemoryIdentity("a" * 64, "owner", "team", "org", outcome == "readonly")
|
||||
if outcome in ("upstream_error", "readonly"):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await serve_memory_response(
|
||||
|
|
@ -194,7 +194,7 @@ async def test_flat_enrollment_follows_the_user_across_keys(prisma_edge: MagicMo
|
|||
param_value=MemorySettings(enabled=True, everyone=False, user_ids=("owner",)).model_dump()
|
||||
)
|
||||
assert (await resolve_memory_access(prisma_edge, _IDENTITY)).active
|
||||
sibling = MemoryIdentity("b" * 64, "owner", "team", "project", "org", False)
|
||||
sibling = MemoryIdentity("b" * 64, "owner", "team", "org", False)
|
||||
assert (await resolve_memory_access(prisma_edge, sibling)).active
|
||||
config.return_value = SimpleNamespace(
|
||||
param_value=MemorySettings(enabled=True, everyone=False, user_ids=("other",)).model_dump()
|
||||
|
|
@ -209,7 +209,7 @@ async def test_store_rechecks_access_before_writing(prisma_edge: MagicMock, chan
|
|||
if change == "missing":
|
||||
prisma_edge.db.litellm_config.find_unique.return_value = None
|
||||
elif change == "readonly":
|
||||
identity = MemoryIdentity("a" * 64, "owner", "team", "project", "org", True)
|
||||
identity = MemoryIdentity("a" * 64, "owner", "team", "org", True)
|
||||
else:
|
||||
config = MemorySettings(enabled=change != "disabled", everyone=False, user_ids=("other",))
|
||||
prisma_edge.db.litellm_config.find_unique.return_value = SimpleNamespace(param_value=config.model_dump())
|
||||
|
|
@ -377,7 +377,7 @@ def request() -> Request:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_only_injection_and_forced_no_tools_do_not_request_reflection(prisma_edge: MagicMock) -> None:
|
||||
read_only: Final = MemoryIdentity("a" * 64, "owner", "team", "project", "org", True)
|
||||
read_only: Final = MemoryIdentity("a" * 64, "owner", "team", "org", True)
|
||||
loop: Final = GatewayMemoryLoop(
|
||||
FastAPI(),
|
||||
request(),
|
||||
|
|
@ -468,7 +468,7 @@ async def test_restore_preserves_hidden_memory_tool_results_and_client_cache_mar
|
|||
*object_items(replacement[1]["content"]),
|
||||
*object_items(items[2]["content"]),
|
||||
)
|
||||
sibling: Final = MemoryIdentity("b" * 64, "owner", "team", "project", "org", False)
|
||||
sibling: Final = MemoryIdentity("b" * 64, "owner", "team", "org", False)
|
||||
assert MemoryContinuations(store(prisma_edge, sibling), "anthropic_messages").identifier(
|
||||
anchor
|
||||
) != continuations.identifier(anchor)
|
||||
|
|
@ -1085,7 +1085,7 @@ async def test_continuation_quota_shares_namespace_lock_across_keys_and_allows_r
|
|||
prisma_edge: MagicMock,
|
||||
) -> None:
|
||||
prisma_edge.db.query_raw.return_value = [{"key_count": 255, "bytes": 32 * 1024 * 1024}]
|
||||
other_key = MemoryIdentity("b" * 64, "owner", "team", "project", "org", False)
|
||||
other_key = MemoryIdentity("b" * 64, "owner", "team", "org", False)
|
||||
for identity in (_IDENTITY, other_key):
|
||||
continuations = MemoryContinuations(MemoryStore(prisma_edge, access_for(identity)), "aresponses")
|
||||
await continuations.save_many((("response", MemoryContinuation(replaces=1, response={"text": "é漢字"})),))
|
||||
|
|
|
|||
|
|
@ -4298,7 +4298,7 @@ class TestDisconnectGatherCleanup:
|
|||
await processing_obj.base_process_llm_request(
|
||||
request=self._disconnect_request(),
|
||||
fastapi_response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(),
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
general_settings={"cancel_on_disconnect": True},
|
||||
proxy_config=MagicMock(spec=ProxyConfig),
|
||||
|
|
|
|||
|
|
@ -33,8 +33,6 @@ export interface DeletedKeysResponse {
|
|||
}
|
||||
|
||||
export interface KeyListCallOptions {
|
||||
includeTeamKeys?: boolean;
|
||||
includeCreatedByKeys?: boolean;
|
||||
organizationID?: string | null;
|
||||
teamID?: string | null;
|
||||
projectID?: string | null;
|
||||
|
|
@ -73,8 +71,8 @@ const keyListCall = async (accessToken: string, page: number, pageSize: number,
|
|||
expand: options.expand,
|
||||
status: options.status,
|
||||
return_full_object: "true",
|
||||
include_team_keys: options.includeTeamKeys ?? true,
|
||||
include_created_by_keys: options.includeCreatedByKeys ?? true,
|
||||
include_team_keys: "true",
|
||||
include_created_by_keys: "true",
|
||||
// Opt into substring matching so the admin key-list search box keeps
|
||||
// matching partial user_id/key_alias. /key/list is exact by default.
|
||||
substring_matching: "true",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue