mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
fix(memory): invalidate worker hints and bound pilot validation
This commit is contained in:
parent
0269841595
commit
7297247aab
6 changed files with 378 additions and 5 deletions
|
|
@ -83,6 +83,9 @@ correct, or delete entries in Memory; callers can use the self-service API.
|
|||
- Preparation stores durable facts supported by the conversation, then searches
|
||||
and reads relevant entries. Search is bounded keyword matching in Postgres.
|
||||
There is no vector database, extraction model, scheduler, or nightly process.
|
||||
- On gateway/backend deployments without shared Redis, first-time activation
|
||||
can take up to 30 seconds to reach another process. Policy revocation is
|
||||
checked against the primary database before memory operations.
|
||||
- Stored references are untrusted data. They cannot grant API permissions or
|
||||
change the namespace derived from authentication. Current user corrections
|
||||
take precedence. Replacements require the current revision.
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""An isolated office pilot that preserves upstream gateway credentials."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
|
|
@ -46,6 +47,7 @@ forward_credential: Final = ForwardCredential()
|
|||
class PilotGateway:
|
||||
def __init__(self, app: ASGIApp) -> None:
|
||||
self.app = app
|
||||
self.validation_slots = asyncio.Semaphore(16)
|
||||
self.upstream = get_async_httpx_client(
|
||||
httpxSpecialProvider.PassThroughEndpoint,
|
||||
params={"timeout": 20, "client_alias": "memory-pilot-upstream"},
|
||||
|
|
@ -96,6 +98,15 @@ class PilotGateway:
|
|||
{"error": "Upstream keys can only use inference and their own memories"}, status_code=403
|
||||
)(scope, receive, send)
|
||||
return
|
||||
try:
|
||||
await asyncio.wait_for(self.validation_slots.acquire(), timeout=0.05)
|
||||
except TimeoutError:
|
||||
await JSONResponse(
|
||||
{"error": "Pilot credential validation is busy; retry shortly"},
|
||||
status_code=503,
|
||||
headers={"Retry-After": "1"},
|
||||
)(scope, receive, send)
|
||||
return
|
||||
try:
|
||||
models: Final = await self.upstream.get(
|
||||
_UPSTREAM + "/v1/models", headers={"Authorization": "Bearer " + credential}
|
||||
|
|
@ -103,6 +114,8 @@ class PilotGateway:
|
|||
except httpx.HTTPError:
|
||||
await JSONResponse({"error": "Upstream gateway unavailable"}, status_code=503)(scope, receive, send)
|
||||
return
|
||||
finally:
|
||||
self.validation_slots.release()
|
||||
if models.is_error:
|
||||
await JSONResponse({"error": "Upstream gateway rejected this key"}, status_code=models.status_code)(
|
||||
scope, receive, send
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from fastapi import HTTPException
|
|||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
|
||||
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient
|
||||
from litellm.repositories.table_repositories import MemoryPolicyRepository, MemoryPreferenceRepository
|
||||
from litellm.types.memory_v2 import MemoryPolicy, MemoryScope, MemoryStatus
|
||||
|
|
@ -22,8 +23,16 @@ async def gateway_memory_is_configured(prisma_client: object, cache: DualCache)
|
|||
without shared Redis when an administrator first enables memory.
|
||||
"""
|
||||
cached: Final = await cache.async_get_cache(key=_CONFIGURED_CACHE_KEY)
|
||||
if isinstance(cached, bool):
|
||||
return cached
|
||||
if cached is True:
|
||||
return True
|
||||
if cached is False:
|
||||
if cache.redis_cache is None:
|
||||
return False
|
||||
# A backend mutation evicts Redis, but another worker can still hold
|
||||
# a negative local hint (Redis Cluster may not support pub/sub).
|
||||
shared: Final = await cache.redis_cache.async_get_cache(key=_CONFIGURED_CACHE_KEY)
|
||||
if shared is False:
|
||||
return False
|
||||
rows: Final = await MemoryPolicyRepository(memory_primary_client(prisma_client)).table.find_many(take=1)
|
||||
configured: Final = bool(rows)
|
||||
await cache.async_set_cache(key=_CONFIGURED_CACHE_KEY, value=configured, ttl=30)
|
||||
|
|
@ -33,7 +42,7 @@ async def gateway_memory_is_configured(prisma_client: object, cache: DualCache)
|
|||
async def invalidate_memory_configuration() -> None:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
await user_api_key_cache.async_delete_cache(key=_CONFIGURED_CACHE_KEY)
|
||||
await evict_and_broadcast(cache_keys=(_CONFIGURED_CACHE_KEY,), user_api_key_cache=user_api_key_cache)
|
||||
|
||||
|
||||
def memory_primary_client(prisma_client: object) -> WriterPinnedClient:
|
||||
|
|
|
|||
54
tests/test_litellm/proxy/memory/test_memory_pilot.py
Normal file
54
tests/test_litellm/proxy/memory/test_memory_pilot.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
"""Bound unauthenticated upstream validation without replacing gateway authentication."""
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_validation_is_bounded_and_slots_release_after_failure(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("UPSTREAM_LITELLM_BASE_URL", "https://upstream.example.invalid")
|
||||
filename = Path(__file__).resolve().parents[4] / "deploy" / "memory-pilot" / "pilot.py"
|
||||
spec = importlib.util.spec_from_file_location("memory_pilot_test", filename)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
gateway = module.PilotGateway(Starlette())
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
count = 0
|
||||
|
||||
async def upstream_get(*args: object, **kwargs: object) -> httpx.Response:
|
||||
nonlocal count
|
||||
count += 1
|
||||
if count == 16:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
raise httpx.ConnectError("unavailable")
|
||||
|
||||
upstream = MagicMock(get=AsyncMock(side_effect=upstream_get))
|
||||
gateway.upstream = upstream
|
||||
database = MagicMock()
|
||||
database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
|
||||
with patch.multiple( # test-quality-ok: Inject external database/config; exercise real ASGI admission.
|
||||
"litellm.proxy.proxy_server", prisma_client=database, master_key="local-admin"
|
||||
):
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=gateway), base_url="http://pilot") as client:
|
||||
pending = [
|
||||
asyncio.create_task(client.get("/v1/models", headers={"Authorization": f"Bearer sk-invalid-{i}"}))
|
||||
for i in range(16)
|
||||
]
|
||||
await asyncio.wait_for(entered.wait(), timeout=2)
|
||||
refused = await client.get("/v1/models", headers={"Authorization": "Bearer sk-overload"})
|
||||
assert refused.status_code == 503 and refused.headers["retry-after"] == "1"
|
||||
assert upstream.get.await_count == 16
|
||||
release.set()
|
||||
assert all(response.status_code == 503 for response in await asyncio.gather(*pending))
|
||||
again = await client.get("/v1/models", headers={"Authorization": "Bearer sk-next"})
|
||||
assert again.status_code == 503 and "unavailable" in again.text
|
||||
assert upstream.get.await_count == 17
|
||||
|
|
@ -351,8 +351,9 @@ async def test_gateway_rounds_keep_separate_limiter_contexts_and_original_client
|
|||
)
|
||||
auth = UserAPIKeyAuth(token="a" * 64, user_id="owner", team_id="team", project_id="project", org_id="org")
|
||||
before = get_request_stash()
|
||||
# test-quality-ok: inject only database, cache, and ASGI provider edges; the memory and limiter code runs unchanged.
|
||||
with patch.multiple(proxy_server, app=provider, prisma_client=prisma_edge, user_api_key_cache=DualCache()):
|
||||
with patch.multiple( # test-quality-ok: Inject database/cache/ASGI provider edges; run real memory and limiter code.
|
||||
proxy_server, app=provider, prisma_client=prisma_edge, user_api_key_cache=DualCache()
|
||||
):
|
||||
prepared = await prepare_gateway_memory(original, request, auth, "anthropic_messages")
|
||||
release.set()
|
||||
owners = await asyncio.gather(*deferred)
|
||||
|
|
@ -364,3 +365,42 @@ async def test_gateway_rounds_keep_separate_limiter_contexts_and_original_client
|
|||
assert prepared["tools"] == original["tools"] and prepared["tool_choice"] == original["tool_choice"]
|
||||
assert prepared["stream"] is True and prepared["max_tokens"] == 1
|
||||
assert "8347" in json.dumps(prepared) and "8347" not in json.dumps(original)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_backend_activation_invalidates_a_gateway_negative_hint_without_pubsub(prisma_edge: MagicMock) -> None:
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy.memory.policy import gateway_memory_is_configured, invalidate_memory_configuration
|
||||
|
||||
shared = {}
|
||||
|
||||
async def get(key, **kwargs):
|
||||
return shared.get(key)
|
||||
|
||||
async def set_value(key, value, **kwargs):
|
||||
shared[key] = value
|
||||
|
||||
async def delete(key, **kwargs):
|
||||
shared.pop(key, None)
|
||||
|
||||
redis = MagicMock(
|
||||
async_get_cache=AsyncMock(side_effect=get),
|
||||
async_set_cache=AsyncMock(side_effect=set_value),
|
||||
async_delete_cache=AsyncMock(side_effect=delete),
|
||||
)
|
||||
gateway_cache = DualCache(redis_cache=redis)
|
||||
backend_cache = DualCache(redis_cache=redis)
|
||||
policies = prisma_edge.db.litellm_memorypolicy.find_many
|
||||
policies.return_value = []
|
||||
assert not await gateway_memory_is_configured(prisma_edge, gateway_cache)
|
||||
assert not await gateway_memory_is_configured(prisma_edge, gateway_cache)
|
||||
policies.assert_awaited_once()
|
||||
policies.return_value = [_POLICY]
|
||||
with patch.multiple( # test-quality-ok: Inject external worker caches and Redis; run real invalidation.
|
||||
"litellm.proxy.proxy_server", user_api_key_cache=backend_cache, redis_usage_cache=None
|
||||
):
|
||||
await invalidate_memory_configuration()
|
||||
assert await gateway_memory_is_configured(prisma_edge, gateway_cache)
|
||||
assert policies.await_count == 2
|
||||
|
|
|
|||
254
tests/test_litellm/proxy/memory/test_memory_v2_management.py
Normal file
254
tests/test_litellm/proxy/memory/test_memory_v2_management.py
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
"""Exercise real policy administration and self-service logic at the database edge."""
|
||||
|
||||
from collections.abc import Iterator
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.memory import management
|
||||
from litellm.proxy.memory.policy import MemoryIdentity, memory_digest
|
||||
from litellm.types.memory_v2 import MemoryCapture, MemoryPolicy, MemoryPolicyInput, MemoryPreference
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def database() -> Iterator[MagicMock]:
|
||||
client = MagicMock()
|
||||
for name in (
|
||||
"litellm_memorypolicy",
|
||||
"litellm_memorypreference",
|
||||
"litellm_memorytable",
|
||||
"litellm_teamtable",
|
||||
"litellm_projecttable",
|
||||
"litellm_organizationtable",
|
||||
"litellm_organizationmembership",
|
||||
"litellm_verificationtoken",
|
||||
"litellm_usertable",
|
||||
):
|
||||
table = getattr(client.db, name)
|
||||
table.find_unique = AsyncMock(return_value=None)
|
||||
table.find_first = AsyncMock(return_value=None)
|
||||
table.find_many = AsyncMock(return_value=[])
|
||||
table.upsert = AsyncMock()
|
||||
table.delete = AsyncMock()
|
||||
table.delete_many = AsyncMock(return_value=0)
|
||||
table.create = AsyncMock()
|
||||
client.db.litellm_teamtable.find_unique.return_value = {
|
||||
"team_id": "team",
|
||||
"organization_id": None,
|
||||
"members_with_roles": [{"user_id": "team-admin", "role": "admin"}],
|
||||
}
|
||||
client.db.litellm_projecttable.find_unique.return_value = {"project_id": "project", "team_id": "team"}
|
||||
client.db.litellm_verificationtoken.find_unique.return_value = {
|
||||
"token": "a" * 64,
|
||||
"user_id": "owner",
|
||||
"team_id": "team",
|
||||
"org_id": "explicit-org",
|
||||
}
|
||||
client.db.litellm_usertable.find_unique.return_value = {"user_id": "owner"}
|
||||
client.db.litellm_organizationtable.find_unique.return_value = {
|
||||
"organization_id": "org",
|
||||
"budget_id": "budget",
|
||||
"created_by": "admin",
|
||||
"updated_by": "admin",
|
||||
}
|
||||
with patch.multiple( # test-quality-ok: Replace only the external database and cache; exercise actual authorization and repositories.
|
||||
"litellm.proxy.proxy_server", prisma_client=client, user_api_key_cache=MagicMock(async_delete_cache=AsyncMock())
|
||||
):
|
||||
yield client
|
||||
|
||||
|
||||
def auth(user: str = "owner", role: LitellmUserRoles = LitellmUserRoles.INTERNAL_USER) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(token="a" * 64, user_id=user, user_role=role, team_id="team", org_id="explicit-org")
|
||||
|
||||
|
||||
def policy(**changes: object) -> MemoryPolicy:
|
||||
return MemoryPolicy.model_validate(
|
||||
{
|
||||
"policy_id": memory_digest("gateway", "*"),
|
||||
"target_type": "gateway",
|
||||
"target_id": "*",
|
||||
"activation": "automatic",
|
||||
"scope": "key",
|
||||
"updated_by": "admin",
|
||||
"updated_at": datetime(2026, 9, 12, tzinfo=timezone.utc),
|
||||
**changes,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"target,target_id",
|
||||
[
|
||||
("gateway", "*"),
|
||||
("organization", "org"),
|
||||
("team", "team"),
|
||||
("project", "project"),
|
||||
("key", "a" * 64),
|
||||
("user", "owner"),
|
||||
],
|
||||
)
|
||||
async def test_proxy_admin_can_configure_each_existing_target(database: MagicMock, target: str, target_id: str) -> None:
|
||||
request = MemoryPolicyInput.model_validate(
|
||||
{"target_type": target, "target_id": target_id, "activation": "opt_in", "scope": "key"}
|
||||
)
|
||||
database.db.litellm_memorypolicy.upsert.return_value = policy(**request.model_dump())
|
||||
result = await management.set_policy(request, auth("admin", LitellmUserRoles.PROXY_ADMIN))
|
||||
assert result.activation == "opt_in"
|
||||
written = database.db.litellm_memorypolicy.upsert.call_args.kwargs
|
||||
assert written["where"]["policy_id"] == memory_digest(target, target_id)
|
||||
assert written["data"]["create"]["updated_by"] == "admin"
|
||||
assert written["data"]["update"]["target_id"] == target_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("target,target_id", [("team", "team"), ("project", "project"), ("key", "a" * 64)])
|
||||
async def test_team_admin_can_configure_owned_targets_but_cannot_broaden_to_user_scope(
|
||||
database: MagicMock, target: str, target_id: str
|
||||
) -> None:
|
||||
request = MemoryPolicyInput.model_validate(
|
||||
{"target_type": target, "target_id": target_id, "activation": "automatic", "scope": "team"}
|
||||
)
|
||||
database.db.litellm_memorypolicy.upsert.return_value = policy(**request.model_dump())
|
||||
assert (await management.set_policy(request, auth("team-admin"))).scope == "team"
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await management.set_policy(request.model_copy(update={"scope": "user"}), auth("team-admin"))
|
||||
assert denied.value.status_code == 403
|
||||
database.db.litellm_memorypolicy.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_membership_is_checked_for_the_selected_organization(database: MagicMock) -> None:
|
||||
request = MemoryPolicyInput(
|
||||
target_type="organization", target_id="org", activation="automatic", scope="organization"
|
||||
)
|
||||
database.db.litellm_organizationmembership.find_first.return_value = SimpleNamespace(user_role="org_admin")
|
||||
database.db.litellm_memorypolicy.upsert.return_value = policy(**request.model_dump())
|
||||
assert (await management.set_policy(request, auth("org-admin"))).scope == "organization"
|
||||
assert database.db.litellm_organizationmembership.find_first.call_args.kwargs["where"] == {
|
||||
"organization_id": "org",
|
||||
"user_id": "org-admin",
|
||||
"user_role": "org_admin",
|
||||
}
|
||||
database.db.litellm_organizationmembership.find_first.return_value = None
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await management.set_policy(request, auth("org-admin"))
|
||||
assert denied.value.status_code == 403
|
||||
database.db.litellm_memorypolicy.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"role",
|
||||
[LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY],
|
||||
)
|
||||
async def test_non_administrator_cannot_change_gateway_policy(database: MagicMock, role: LitellmUserRoles) -> None:
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await management.set_policy(
|
||||
MemoryPolicyInput(target_type="gateway", target_id="*", activation="automatic"), auth(role=role)
|
||||
)
|
||||
assert denied.value.status_code == 403
|
||||
database.db.litellm_memorypolicy.upsert.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_policy_listing_requires_owned_target_and_preserves_pagination(database: MagicMock) -> None:
|
||||
database.db.litellm_memorypolicy.find_many.return_value = [policy(target_type="team", target_id="team")]
|
||||
assert len(await management.list_policies("team", "team", 100, auth("team-admin"))) == 1
|
||||
assert database.db.litellm_memorypolicy.find_many.call_args.kwargs == {
|
||||
"where": {"target_type": "team", "target_id": "team"},
|
||||
"take": 100,
|
||||
"skip": 100,
|
||||
"order": {"policy_id": "asc"},
|
||||
}
|
||||
for caller, target, expected in [(auth(), None, 403), (auth("admin", LitellmUserRoles.PROXY_ADMIN), "team", 400)]:
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await management.list_policies(target, None, 0, caller)
|
||||
assert denied.value.status_code == expected
|
||||
assert (
|
||||
len(await management.list_policies(None, None, 0, auth("admin", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY))) == 1
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_policy_delete_rechecks_administrator_and_reports_missing(database: MagicMock) -> None:
|
||||
table = database.db.litellm_memorypolicy
|
||||
table.find_unique.return_value = policy()
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await management.delete_policy("policy", auth())
|
||||
assert denied.value.status_code == 403
|
||||
table.delete.assert_not_awaited()
|
||||
assert (await management.delete_policy("policy", auth("admin", LitellmUserRoles.PROXY_ADMIN))).status_code == 204
|
||||
table.find_unique.return_value = None
|
||||
with pytest.raises(HTTPException) as missing:
|
||||
await management.delete_policy("policy", auth("admin", LitellmUserRoles.PROXY_ADMIN))
|
||||
assert missing.value.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preference_updates_are_bound_to_authenticated_subject(database: MagicMock) -> None:
|
||||
assert not (await management.get_preference(auth())).enabled
|
||||
table = database.db.litellm_memorypreference
|
||||
assert (await management.set_preference(MemoryPreference(enabled=True), auth())).enabled
|
||||
assert table.upsert.call_args.kwargs["where"] == {"subject": memory_digest("user", "owner")}
|
||||
table.find_unique.return_value = SimpleNamespace(enabled=True)
|
||||
assert (await management.get_preference(auth())).enabled
|
||||
assert not (await management.set_preference(MemoryPreference(enabled=False), auth())).enabled
|
||||
assert table.delete_many.call_args.kwargs["where"] == {"subject": memory_digest("user", "owner")}
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await management.set_preference(
|
||||
MemoryPreference(enabled=False), auth(role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY)
|
||||
)
|
||||
assert denied.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dashboard_session_can_select_own_key_but_not_another_users_key(database: MagicMock) -> None:
|
||||
database.db.litellm_memorypolicy.find_many.return_value = [policy()]
|
||||
session = auth().model_copy(update={"token": "session", "team_id": UI_TEAM_ID})
|
||||
own = await management.access_for_key(session, "a" * 64)
|
||||
assert own.active and own.identity.organization_id == "explicit-org"
|
||||
assert own.namespace == MemoryIdentity.from_auth(auth()).namespace("key")
|
||||
assert (await management.get_status(None, auth())).active
|
||||
database.db.litellm_verificationtoken.find_unique.return_value["user_id"] = "someone-else"
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await management.access_for_key(session, "a" * 64)
|
||||
assert denied.value.status_code == 403
|
||||
database.db.litellm_verificationtoken.find_unique.return_value = None
|
||||
with pytest.raises(HTTPException) as missing:
|
||||
await management.access_for_key(auth("admin", LitellmUserRoles.PROXY_ADMIN), "a" * 64)
|
||||
assert missing.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entry_endpoints_apply_namespace_and_delete_after_disable(database: MagicMock) -> None:
|
||||
database.db.litellm_memorypolicy.find_many.return_value = [policy()]
|
||||
namespace = MemoryIdentity.from_auth(auth()).namespace("key")
|
||||
table = database.db.litellm_memorytable
|
||||
assert await management.list_entries("demo", 2, 4, None, auth()) == []
|
||||
assert table.find_many.call_args.kwargs["where"]["namespace"] == namespace
|
||||
now = datetime(2026, 9, 12, tzinfo=timezone.utc)
|
||||
table.create.return_value = SimpleNamespace(
|
||||
memory_id="entry",
|
||||
key=f"memory-v2:{namespace}:demo",
|
||||
namespace=namespace,
|
||||
value="Use port 8347",
|
||||
metadata={"title": "Demo", "evidence": "User said so"},
|
||||
updated_at=now,
|
||||
)
|
||||
saved = await management.capture_entry(
|
||||
MemoryCapture(key="demo", title="Demo", content="Use port 8347", evidence="User said so"), None, auth()
|
||||
)
|
||||
assert saved.content == "Use port 8347" and saved.key == "demo"
|
||||
database.db.litellm_memorypolicy.find_many.return_value = [policy(activation="disabled")]
|
||||
table.delete_many.return_value = 1
|
||||
assert (await management.delete_entry("entry", None, auth())).status_code == 204
|
||||
assert table.delete_many.call_args.kwargs["where"] == {"namespace": namespace, "memory_id": "entry"}
|
||||
table.delete_many.return_value = 0
|
||||
with pytest.raises(HTTPException) as missing:
|
||||
await management.delete_entry("entry", None, auth())
|
||||
assert missing.value.status_code == 404
|
||||
Loading…
Add table
Reference in a new issue