mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
fix(memory): isolate request accounting and support split deployments
This commit is contained in:
moe-berri
parent
42c8bb387c
commit
0269841595
19 changed files with 537 additions and 44 deletions
|
|
@ -71,6 +71,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/v1/workflows/",
|
||||
"/project/",
|
||||
"/memory/",
|
||||
"/v2/memory/",
|
||||
"/mcp/",
|
||||
# Control plane (see the List Endpoints + Tables standard). Every resource
|
||||
# eventually moves under this prefix, so allowlist it once rather than
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ set -euo pipefail
|
|||
# writable homes to install the repository's pinned toolchain during uv sync.
|
||||
export RUSTUP_HOME="$PWD/.memory-pilot-rustup"
|
||||
export CARGO_HOME="$PWD/.memory-pilot-cargo"
|
||||
export PRISMA_BINARY_CACHE_DIR="$PWD/.memory-pilot-prisma"
|
||||
python -m pip install uv==0.11.7
|
||||
export UV_PROJECT_ENVIRONMENT="$PWD/.memory-pilot-venv"
|
||||
uv sync --frozen --extra proxy --extra extra_proxy --no-default-groups
|
||||
|
|
|
|||
|
|
@ -7,15 +7,19 @@ from contextvars import ContextVar
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException, Request
|
||||
from starlette.exceptions import HTTPException
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy._types import UI_TEAM_ID, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
|
||||
from litellm.proxy.auth.user_api_key_auth import _get_bearer_token_or_received_api_key
|
||||
from litellm.repositories.verification_token_repository import VerificationTokenRepository
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.utils import CallTypesLiteral
|
||||
|
||||
_UPSTREAM: Final = os.environ["UPSTREAM_LITELLM_BASE_URL"].rstrip("/")
|
||||
|
|
@ -42,25 +46,28 @@ forward_credential: Final = ForwardCredential()
|
|||
class PilotGateway:
|
||||
def __init__(self, app: ASGIApp) -> None:
|
||||
self.app = app
|
||||
self.upstream = httpx.AsyncClient(timeout=20)
|
||||
self.upstream = get_async_httpx_client(
|
||||
httpxSpecialProvider.PassThroughEndpoint,
|
||||
params={"timeout": 20, "client_alias": "memory-pilot-upstream"},
|
||||
)
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] == "lifespan":
|
||||
try:
|
||||
await self.app(scope, receive, send)
|
||||
finally:
|
||||
await self.upstream.aclose()
|
||||
await self.upstream.close()
|
||||
return
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
request: Final = Request(scope, receive)
|
||||
credential: Final = (
|
||||
credential: Final = _get_bearer_token_or_received_api_key(
|
||||
request.headers.get("x-litellm-api-key")
|
||||
or request.headers.get("authorization")
|
||||
or request.headers.get("x-api-key")
|
||||
or ""
|
||||
).removeprefix("Bearer ")
|
||||
)
|
||||
from litellm.proxy.proxy_server import master_key, prisma_client
|
||||
|
||||
if not credential or master_key and secrets.compare_digest(credential, master_key):
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
export PATH="$PWD/.memory-pilot-venv/bin:$PATH"
|
||||
export PRISMA_BINARY_CACHE_DIR="$PWD/.memory-pilot-prisma"
|
||||
export PRISMA_CLI_PATH="$PRISMA_BINARY_CACHE_DIR/node_modules/.bin/prisma"
|
||||
prisma migrate deploy --schema litellm-proxy-extras/litellm_proxy_extras/schema.prisma
|
||||
export WORKER_CONFIG="$PWD/deploy/memory-pilot/proxy_config.yaml"
|
||||
export PYTHONPATH="$PWD/deploy/memory-pilot${PYTHONPATH:+:$PYTHONPATH}"
|
||||
|
|
|
|||
|
|
@ -74,6 +74,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/containers",
|
||||
"/v1/evals",
|
||||
"/v1/memory",
|
||||
"/v2/memory/entries",
|
||||
"/queue/chat/",
|
||||
# Google data plane (v1beta is the Google AI Studio version)
|
||||
"/v1beta/",
|
||||
|
|
@ -123,6 +124,8 @@ GATEWAY_EXACT_PATHS: frozenset[str] = frozenset(
|
|||
"/redoc",
|
||||
"/test",
|
||||
"/debug/memory/summary",
|
||||
"/v2/memory/status",
|
||||
"/v2/memory/preference",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN "namespace" TEXT;
|
||||
ALTER TABLE "LiteLLM_MemoryTable" ADD COLUMN IF NOT EXISTS "namespace" TEXT;
|
||||
|
||||
CREATE INDEX "LiteLLM_MemoryTable_namespace_updated_at_idx"
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MemoryTable_namespace_updated_at_idx"
|
||||
ON "LiteLLM_MemoryTable"("namespace", "updated_at");
|
||||
|
||||
CREATE TABLE "LiteLLM_MemoryPolicy" (
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_MemoryPolicy" (
|
||||
"policy_id" TEXT NOT NULL,
|
||||
"target_type" TEXT NOT NULL,
|
||||
"target_id" TEXT NOT NULL,
|
||||
|
|
@ -15,10 +15,10 @@ CREATE TABLE "LiteLLM_MemoryPolicy" (
|
|||
CONSTRAINT "LiteLLM_MemoryPolicy_pkey" PRIMARY KEY ("policy_id")
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX "LiteLLM_MemoryPolicy_target_type_target_id_key"
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MemoryPolicy_target_type_target_id_key"
|
||||
ON "LiteLLM_MemoryPolicy"("target_type", "target_id");
|
||||
|
||||
CREATE TABLE "LiteLLM_MemoryPreference" (
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -522,6 +522,7 @@ class RequestRateLimiterStash:
|
|||
owner_litellm_call_id: str | None = None
|
||||
rate_limit_response: RateLimitResponse | None = None
|
||||
parallel_slot: ParallelSlotAcquisition | None = None
|
||||
parallel_release_complete: asyncio.Event = field(default_factory=asyncio.Event)
|
||||
reserved_tokens: int = 0
|
||||
reserved_model: str | None = None
|
||||
reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
|
||||
|
|
@ -548,6 +549,13 @@ def get_request_stash() -> RequestRateLimiterStash | None:
|
|||
return _request_stash.get()
|
||||
|
||||
|
||||
async def wait_for_request_parallel_release() -> None:
|
||||
"""Let sequential internal requests wait for their deferred slot release."""
|
||||
stash: Final = get_request_stash()
|
||||
if stash is not None and stash.parallel_slot is not None:
|
||||
await stash.parallel_release_complete.wait()
|
||||
|
||||
|
||||
def get_or_create_request_stash() -> RequestRateLimiterStash:
|
||||
stash = _request_stash.get()
|
||||
if stash is None:
|
||||
|
|
@ -3375,6 +3383,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
stash.parallel_release_complete.set()
|
||||
self._handle_rate_limit_error(
|
||||
response=io_response,
|
||||
descriptors=descriptors,
|
||||
|
|
@ -3547,6 +3556,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
slot_id=parallel_slot_id,
|
||||
counter_keys=parallel_counter_keys,
|
||||
)
|
||||
stash.parallel_release_complete.clear()
|
||||
|
||||
# ----------------------------------------------------------------
|
||||
# TPM token reservation
|
||||
|
|
@ -3638,6 +3648,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
stash.parallel_release_complete.set()
|
||||
self._handle_rate_limit_error(
|
||||
response=tpm_response,
|
||||
descriptors=descriptors,
|
||||
|
|
@ -4457,6 +4468,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span=litellm_parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
stash.parallel_release_complete.set()
|
||||
|
||||
pipeline_operations: Final = self._build_success_event_pipeline_operations(
|
||||
kwargs=kwargs,
|
||||
|
|
@ -4583,6 +4595,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span=litellm_parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
stash.parallel_release_complete.set()
|
||||
|
||||
# Skip the reservation refund if async_post_call_failure_hook
|
||||
# already released it (proxy-level rejection that also bubbles up
|
||||
|
|
@ -4699,6 +4712,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span=None,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
stash.parallel_release_complete.set()
|
||||
|
||||
async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response):
|
||||
"""
|
||||
|
|
@ -4780,6 +4794,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
stash.parallel_slot = None
|
||||
stash.parallel_release_complete.set()
|
||||
|
||||
if stash.batch_enqueued_reservation is not None:
|
||||
await self.batch_enqueued_token_store.refund(
|
||||
|
|
|
|||
|
|
@ -20,7 +20,8 @@ from litellm.litellm_core_utils.prompt_templates.server_tools import (
|
|||
prepare_server_tools,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.memory.policy import MemoryIdentity, resolve_memory_access
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import wait_for_request_parallel_release
|
||||
from litellm.proxy.memory.policy import MemoryIdentity, gateway_memory_is_configured, resolve_memory_access
|
||||
from litellm.proxy.memory.store import MemoryStore
|
||||
from litellm.types.memory_v2 import MemoryCapture, MemoryRead, MemorySearch
|
||||
|
||||
|
|
@ -155,11 +156,16 @@ async def prepare_gateway_memory(
|
|||
) -> dict[str, object]:
|
||||
if _memory_call.get() or route not in ("acompletion", "aresponses", "anthropic_messages"):
|
||||
return data
|
||||
from litellm.proxy.proxy_server import app, prisma_client
|
||||
from litellm.proxy.proxy_server import app, prisma_client, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
return data
|
||||
access: Final = await resolve_memory_access(prisma_client, MemoryIdentity.from_auth(auth))
|
||||
identity: Final = MemoryIdentity.from_auth(auth)
|
||||
if not identity.user_id and not identity.key_id:
|
||||
return data
|
||||
if not await gateway_memory_is_configured(prisma_client, user_api_key_cache):
|
||||
return data
|
||||
access: Final = await resolve_memory_access(prisma_client, identity)
|
||||
if not access.active:
|
||||
return data
|
||||
functions: Final = tuple(
|
||||
|
|
@ -184,18 +190,35 @@ async def prepare_gateway_memory(
|
|||
}
|
||||
token: Final = _memory_call.set(True)
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app), base_url="http://litellm-memory"
|
||||
) as client:
|
||||
# Dispatch in process through the existing authenticated endpoint. No
|
||||
# network client, TLS context, or connection pool is created here.
|
||||
async with httpx.ASGITransport(app=app) as transport:
|
||||
|
||||
async def dispatch_round(body: Mapping[str, object]) -> httpx.Response:
|
||||
result: Final = await transport.handle_async_request(
|
||||
httpx.Request(
|
||||
"POST",
|
||||
"http://litellm-memory" + request.url.path,
|
||||
json=body,
|
||||
headers=headers,
|
||||
params=request.query_params,
|
||||
)
|
||||
)
|
||||
await result.aread()
|
||||
# Success accounting runs asynchronously. The next model round
|
||||
# must not compete with this completed call for the same slot.
|
||||
await wait_for_request_parallel_release()
|
||||
return result
|
||||
|
||||
async def call_model(body: Mapping[str, object]) -> Mapping[str, object]:
|
||||
round_body: Final = { # mutable-ok: HTTP JSON serialization requires a native dictionary.
|
||||
**body,
|
||||
"litellm_call_id": str(uuid4()),
|
||||
}
|
||||
result: Final = await client.post(
|
||||
request.url.path, json=round_body, headers=headers, params=request.query_params
|
||||
)
|
||||
# Each endpoint owns its request context, including the rate
|
||||
# limiter's mutable stash. Reusing this task would let the next
|
||||
# round overwrite the owner seen by deferred logging callbacks.
|
||||
result: Final = await asyncio.create_task(dispatch_round(round_body))
|
||||
if result.is_error:
|
||||
raise HTTPException(
|
||||
status_code=result.status_code,
|
||||
|
|
|
|||
|
|
@ -5,7 +5,14 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Response
|
|||
from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth, user_api_key_has_admin_view
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.memory.memory_endpoints import is_memory_team_admin, require_memory_prisma
|
||||
from litellm.proxy.memory.policy import MemoryAccess, MemoryIdentity, memory_digest, resolve_memory_access
|
||||
from litellm.proxy.memory.policy import (
|
||||
MemoryAccess,
|
||||
MemoryIdentity,
|
||||
invalidate_memory_configuration,
|
||||
memory_digest,
|
||||
memory_primary_client,
|
||||
resolve_memory_access,
|
||||
)
|
||||
from litellm.proxy.memory.store import MemoryStore
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.project_repository import ProjectRepository
|
||||
|
|
@ -41,7 +48,7 @@ router: Final = APIRouter(
|
|||
async def require_policy_admin(
|
||||
auth: UserAPIKeyAuth, target_type: MemoryTarget, target_id: str, *, write: bool = True
|
||||
) -> None:
|
||||
prisma: Final = require_memory_prisma()
|
||||
prisma: Final = memory_primary_client(require_memory_prisma())
|
||||
if write and auth.user_role in (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY):
|
||||
raise HTTPException(status_code=403, detail="Memory policies require administrator write access")
|
||||
proxy_admin: Final = (
|
||||
|
|
@ -92,7 +99,7 @@ async def list_policies(
|
|||
raise HTTPException(status_code=403, detail="Select a target you administer")
|
||||
elif target_type is not None or target_id is not None:
|
||||
raise HTTPException(status_code=400, detail="Provide both target_type and target_id")
|
||||
rows: Final = await MemoryPolicyRepository(require_memory_prisma()).table.find_many(
|
||||
rows: Final = await MemoryPolicyRepository(memory_primary_client(require_memory_prisma())).table.find_many(
|
||||
where={ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
"target_type": target_type,
|
||||
"target_id": target_id,
|
||||
|
|
@ -121,7 +128,7 @@ async def set_policy(policy: MemoryPolicyInput, auth: UserAPIKeyAuth = _AUTH) ->
|
|||
**policy.model_dump(),
|
||||
"updated_by": auth.user_id or "proxy-admin",
|
||||
}
|
||||
row: Final = await MemoryPolicyRepository(require_memory_prisma()).table.upsert(
|
||||
row: Final = await MemoryPolicyRepository(memory_primary_client(require_memory_prisma())).table.upsert(
|
||||
where={ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
"policy_id": policy_id
|
||||
},
|
||||
|
|
@ -133,12 +140,13 @@ async def set_policy(policy: MemoryPolicyInput, auth: UserAPIKeyAuth = _AUTH) ->
|
|||
"update": fields,
|
||||
},
|
||||
)
|
||||
await invalidate_memory_configuration()
|
||||
return MemoryPolicy.model_validate(row, from_attributes=True)
|
||||
|
||||
|
||||
@router.delete("/policies/{policy_id}", status_code=204)
|
||||
async def delete_policy(policy_id: str, auth: UserAPIKeyAuth = _AUTH) -> Response:
|
||||
table: Final = MemoryPolicyRepository(require_memory_prisma()).table
|
||||
table: Final = MemoryPolicyRepository(memory_primary_client(require_memory_prisma())).table
|
||||
row: Final = await table.find_unique(
|
||||
where={ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
"policy_id": policy_id
|
||||
|
|
@ -153,13 +161,14 @@ async def delete_policy(policy_id: str, auth: UserAPIKeyAuth = _AUTH) -> Respons
|
|||
"policy_id": policy_id
|
||||
}
|
||||
)
|
||||
await invalidate_memory_configuration()
|
||||
return Response(status_code=204)
|
||||
|
||||
|
||||
@router.get("/preference", response_model=MemoryPreference)
|
||||
async def get_preference(auth: UserAPIKeyAuth = _AUTH) -> MemoryPreference:
|
||||
subject: Final = MemoryIdentity.from_auth(auth).preference_subject
|
||||
row: Final = await MemoryPreferenceRepository(require_memory_prisma()).table.find_unique(
|
||||
row: Final = await MemoryPreferenceRepository(memory_primary_client(require_memory_prisma())).table.find_unique(
|
||||
where={ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
"subject": subject
|
||||
}
|
||||
|
|
@ -174,13 +183,13 @@ async def set_preference(preference: MemoryPreference, auth: UserAPIKeyAuth = _A
|
|||
raise HTTPException(status_code=403, detail="Read-only users cannot change memory preferences")
|
||||
subject: Final = identity.preference_subject
|
||||
if not preference.enabled:
|
||||
await MemoryPreferenceRepository(require_memory_prisma()).table.delete_many(
|
||||
await MemoryPreferenceRepository(memory_primary_client(require_memory_prisma())).table.delete_many(
|
||||
where={ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
"subject": subject
|
||||
}
|
||||
)
|
||||
return preference
|
||||
await MemoryPreferenceRepository(require_memory_prisma()).table.upsert(
|
||||
await MemoryPreferenceRepository(memory_primary_client(require_memory_prisma())).table.upsert(
|
||||
where={ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
"subject": subject
|
||||
},
|
||||
|
|
@ -206,7 +215,7 @@ async def get_status(
|
|||
|
||||
|
||||
async def access_for_key(auth: UserAPIKeyAuth, key_id: str | None) -> MemoryAccess:
|
||||
prisma: Final = require_memory_prisma()
|
||||
prisma: Final = memory_primary_client(require_memory_prisma())
|
||||
if key_id is None:
|
||||
return await resolve_memory_access(prisma, MemoryIdentity.from_auth(auth))
|
||||
key: Final = await VerificationTokenRepository(prisma).find_by_id(key_id, id_field="token")
|
||||
|
|
@ -224,7 +233,7 @@ async def access_for_key(auth: UserAPIKeyAuth, key_id: str | None) -> MemoryAcce
|
|||
user_id=key.user_id,
|
||||
team_id=key.team_id,
|
||||
project_id=key.project_id,
|
||||
organization_id=team.organization_id if team else key.org_id,
|
||||
organization_id=key.org_id or (team.organization_id if team else None),
|
||||
read_only=MemoryIdentity.from_auth(auth).read_only,
|
||||
)
|
||||
return await resolve_memory_access(prisma, identity)
|
||||
|
|
@ -238,7 +247,7 @@ async def list_entries(
|
|||
key_id: str | None = Query(None, pattern=r"^[a-f0-9]{64}$"),
|
||||
auth: UserAPIKeyAuth = _AUTH,
|
||||
) -> list[MemoryEntry]:
|
||||
prisma: Final = require_memory_prisma()
|
||||
prisma: Final = memory_primary_client(require_memory_prisma())
|
||||
access: Final = await access_for_key(auth, key_id)
|
||||
return await MemoryStore(prisma, access).search(
|
||||
MemorySearch(query=query, limit=limit, offset=offset), require_active=False
|
||||
|
|
@ -251,7 +260,7 @@ async def capture_entry(
|
|||
key_id: str | None = Query(None, pattern=r"^[a-f0-9]{64}$"),
|
||||
auth: UserAPIKeyAuth = _AUTH,
|
||||
) -> MemoryEntry:
|
||||
prisma: Final = require_memory_prisma()
|
||||
prisma: Final = memory_primary_client(require_memory_prisma())
|
||||
access: Final = await access_for_key(auth, key_id)
|
||||
return await MemoryStore(prisma, access).capture(capture)
|
||||
|
||||
|
|
@ -262,7 +271,7 @@ async def delete_entry(
|
|||
key_id: str | None = Query(None, pattern=r"^[a-f0-9]{64}$"),
|
||||
auth: UserAPIKeyAuth = _AUTH,
|
||||
) -> Response:
|
||||
prisma: Final = require_memory_prisma()
|
||||
prisma: Final = memory_primary_client(require_memory_prisma())
|
||||
access: Final = await access_for_key(auth, key_id)
|
||||
if not await MemoryStore(prisma, access).delete(memory_id):
|
||||
raise HTTPException(status_code=404, detail="Memory not found")
|
||||
|
|
|
|||
|
|
@ -5,10 +5,46 @@ from typing import Final
|
|||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth
|
||||
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
|
||||
|
||||
_CONFIGURED_CACHE_KEY: Final = "litellm:memory_v2:configured"
|
||||
|
||||
|
||||
async def gateway_memory_is_configured(prisma_client: object, cache: DualCache) -> bool:
|
||||
"""Avoid database work on ordinary requests when memory is not configured.
|
||||
|
||||
This is only a presence hint. Authorization is always checked against the
|
||||
primary before memory is used. The short TTL also covers separate workers
|
||||
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
|
||||
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)
|
||||
return configured
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def memory_primary_client(prisma_client: object) -> WriterPinnedClient:
|
||||
"""Memory authorization and read-after-write must never use a lagging replica."""
|
||||
db: Final = getattr(prisma_client, "db", None)
|
||||
if db is None:
|
||||
raise RuntimeError("Memory requires a connected Prisma database")
|
||||
# Pin the actual writer even while unavailable: memory must fail closed
|
||||
# instead of authorizing storage or recall from stale policy rows.
|
||||
return WriterPinnedClient(db.writer if isinstance(db, RoutingPrismaWrapper) else db)
|
||||
|
||||
|
||||
def memory_digest(*parts: str | None) -> str:
|
||||
return hashlib.sha256(json.dumps(parts, separators=(",", ":")).encode()).hexdigest()
|
||||
|
|
@ -119,7 +155,8 @@ class MemoryAccess:
|
|||
|
||||
|
||||
async def resolve_memory_access(prisma_client: object, identity: MemoryIdentity) -> MemoryAccess:
|
||||
rows: Final = await MemoryPolicyRepository(prisma_client).table.find_many(
|
||||
primary: Final = memory_primary_client(prisma_client)
|
||||
rows: Final = await MemoryPolicyRepository(primary).table.find_many(
|
||||
where={ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
"OR": [ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
{ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
|
|
@ -137,7 +174,7 @@ async def resolve_memory_access(prisma_client: object, identity: MemoryIdentity)
|
|||
policy: Final = next((policies[target] for target in identity.policy_targets if target in policies), None)
|
||||
if not identity.user_id and not identity.key_id:
|
||||
return MemoryAccess(identity=identity, policy=None, opted_in=False)
|
||||
preference: Final = await MemoryPreferenceRepository(prisma_client).table.find_unique(
|
||||
preference: Final = await MemoryPreferenceRepository(primary).table.find_unique(
|
||||
where={ # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
"subject": identity.preference_subject
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,10 +2,9 @@ import json
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from fastapi import HTTPException
|
||||
from prisma.errors import UniqueViolationError
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.proxy.memory.policy import MemoryAccess, memory_digest, resolve_memory_access
|
||||
from litellm.proxy.memory.policy import MemoryAccess, memory_digest, memory_primary_client, resolve_memory_access
|
||||
from litellm.repositories.table_repositories import MemoryRepository
|
||||
from litellm.types.memory_v2 import MemoryCapture, MemoryEntry, MemorySearch
|
||||
|
||||
|
|
@ -36,9 +35,9 @@ def memory_entry(row: "LiteLLM_MemoryTable") -> MemoryEntry:
|
|||
|
||||
class MemoryStore:
|
||||
def __init__(self, prisma_client: object, access: MemoryAccess) -> None:
|
||||
self.prisma_client = prisma_client
|
||||
self.prisma_client = memory_primary_client(prisma_client)
|
||||
self.access = access
|
||||
self.table = MemoryRepository(prisma_client).table
|
||||
self.table = MemoryRepository(self.prisma_client).table
|
||||
|
||||
async def _namespace(self, *, write: bool = False, require_active: bool = True) -> str:
|
||||
current: Final = await resolve_memory_access(self.prisma_client, self.access.identity)
|
||||
|
|
@ -115,6 +114,8 @@ class MemoryStore:
|
|||
return memory_entry(row)
|
||||
|
||||
async def capture(self, capture: MemoryCapture) -> MemoryEntry:
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
namespace: Final = await self._namespace(write=True)
|
||||
key: Final = f"memory-v2:{namespace}:{capture.key}"
|
||||
metadata: Final = { # mutable-ok: Prisma serializes these as native JSON containers.
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
from datetime import datetime
|
||||
from typing import Literal, Self, TypeAlias
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
from typing_extensions import Self
|
||||
|
||||
MemoryTarget: TypeAlias = Literal["gateway", "organization", "team", "project", "user", "key"]
|
||||
MemoryScope: TypeAlias = Literal["key", "user", "team", "project", "organization"]
|
||||
|
|
|
|||
|
|
@ -306,6 +306,7 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = (
|
|||
"/global",
|
||||
"/config",
|
||||
"/guardrails",
|
||||
"/v2/memory/policies",
|
||||
"/openapi.json",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1182,7 +1182,11 @@ async def test_end_user_jwt_auth(monkeypatch):
|
|||
"llm_router",
|
||||
router,
|
||||
)
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", {})
|
||||
# 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)
|
||||
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,7 +503,11 @@ 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")
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", "FAKE-VAR")
|
||||
# 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)
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", memory_db)
|
||||
setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)
|
||||
|
||||
with patch.object(
|
||||
|
|
|
|||
|
|
@ -4232,6 +4232,12 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch):
|
|||
await local_cache.async_get_cache(key=counter_key)
|
||||
) == 1
|
||||
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import wait_for_request_parallel_release
|
||||
|
||||
waiter = asyncio.create_task(wait_for_request_parallel_release())
|
||||
await asyncio.sleep(0)
|
||||
assert not waiter.done(), "Internal rounds must wait while the completed call still owns a slot"
|
||||
|
||||
await handler.async_log_success_event(
|
||||
kwargs={
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
|
||||
|
|
@ -4242,6 +4248,7 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch):
|
|||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
await asyncio.wait_for(waiter, timeout=1)
|
||||
assert handler._gauge_in_flight_from_cache_value(
|
||||
await local_cache.async_get_cache(key=counter_key)
|
||||
) == 0
|
||||
|
|
|
|||
366
tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py
Normal file
366
tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py
Normal file
|
|
@ -0,0 +1,366 @@
|
|||
"""Failure and authorization boundaries, with only the database/model edges replaced."""
|
||||
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from prisma.models import LiteLLM_MemoryTable
|
||||
|
||||
from litellm.proxy.memory.gateway import execute_memory_tool, run_memory_tools
|
||||
from litellm.proxy.memory.policy import MemoryAccess, MemoryIdentity, resolve_memory_access
|
||||
from litellm.proxy.memory.store import MemoryStore
|
||||
from litellm.types.memory_v2 import MemoryCapture, MemoryPolicy, MemorySearch
|
||||
|
||||
_NOW: Final = datetime(2026, 9, 12, tzinfo=timezone.utc)
|
||||
_IDENTITY: Final = MemoryIdentity("a" * 64, "owner", "team", "project", "org", False)
|
||||
_POLICY: Final = MemoryPolicy(
|
||||
policy_id="policy",
|
||||
target_type="team",
|
||||
target_id="team",
|
||||
activation="automatic",
|
||||
scope="key",
|
||||
updated_at=_NOW,
|
||||
updated_by="admin",
|
||||
)
|
||||
_CAPTURE: Final = MemoryCapture(key="demo", title="Demo", content="Use port 8347", evidence="User selected this port")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def prisma_edge() -> MagicMock:
|
||||
client = MagicMock()
|
||||
client.db.litellm_memorypolicy.find_many = AsyncMock(return_value=[_POLICY])
|
||||
client.db.litellm_memorypreference.find_unique = AsyncMock(return_value=None)
|
||||
table = client.db.litellm_memorytable
|
||||
table.find_unique = AsyncMock(return_value=None)
|
||||
table.find_first = AsyncMock(return_value=None)
|
||||
table.find_many = AsyncMock(return_value=[])
|
||||
table.create = AsyncMock()
|
||||
table.update_many = AsyncMock(return_value=1)
|
||||
table.delete_many = AsyncMock(return_value=1)
|
||||
return client
|
||||
|
||||
|
||||
def store(client: MagicMock, identity: MemoryIdentity = _IDENTITY) -> MemoryStore:
|
||||
return MemoryStore(client, MemoryAccess(identity, _POLICY, False))
|
||||
|
||||
|
||||
def row(**changes: object) -> LiteLLM_MemoryTable:
|
||||
return LiteLLM_MemoryTable.model_validate(
|
||||
{
|
||||
"memory_id": "entry",
|
||||
"key": f"memory-v2:{_IDENTITY.namespace('key')}:demo",
|
||||
"namespace": _IDENTITY.namespace("key"),
|
||||
"value": _CAPTURE.content,
|
||||
"metadata": json.dumps({"title": _CAPTURE.title, "evidence": _CAPTURE.evidence}),
|
||||
"created_at": _NOW,
|
||||
"updated_at": _NOW,
|
||||
**changes,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_policy_precedence_and_opt_in_are_resolved_from_database(prisma_edge: MagicMock) -> None:
|
||||
team = _POLICY.model_copy(update={"activation": "opt_in"})
|
||||
key = _POLICY.model_copy(update={"target_type": "key", "target_id": "a" * 64, "activation": "disabled"})
|
||||
policies = prisma_edge.db.litellm_memorypolicy.find_many
|
||||
policies.return_value = [team, key]
|
||||
prisma_edge.db.litellm_memorypreference.find_unique.return_value = SimpleNamespace(enabled=True)
|
||||
access = await resolve_memory_access(prisma_edge, _IDENTITY)
|
||||
assert access.policy == key and not access.active and access.opted_in
|
||||
policies.return_value = [team]
|
||||
assert (await resolve_memory_access(prisma_edge, _IDENTITY)).active
|
||||
prisma_edge.db.litellm_memorypreference.find_unique.return_value = None
|
||||
assert not (await resolve_memory_access(prisma_edge, _IDENTITY)).active
|
||||
policies.return_value = []
|
||||
assert (await resolve_memory_access(prisma_edge, _IDENTITY)).namespace is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("change", ["disabled", "scope", "missing", "readonly"])
|
||||
async def test_store_rechecks_policy_before_writing(prisma_edge: MagicMock, change: str) -> None:
|
||||
identity = _IDENTITY
|
||||
if change == "disabled":
|
||||
prisma_edge.db.litellm_memorypolicy.find_many.return_value = [
|
||||
_POLICY.model_copy(update={"activation": "disabled"})
|
||||
]
|
||||
elif change == "scope":
|
||||
prisma_edge.db.litellm_memorypolicy.find_many.return_value = [_POLICY.model_copy(update={"scope": "team"})]
|
||||
elif change == "missing":
|
||||
prisma_edge.db.litellm_memorypolicy.find_many.return_value = []
|
||||
else:
|
||||
identity = MemoryIdentity("a" * 64, "owner", "team", "project", "org", True)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await store(prisma_edge, identity).capture(_CAPTURE)
|
||||
assert exc.value.status_code == 403
|
||||
prisma_edge.db.litellm_memorytable.create.assert_not_awaited()
|
||||
prisma_edge.db.litellm_memorytable.update_many.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_requires_namespace_and_all_words_with_bounded_paging(prisma_edge: MagicMock) -> None:
|
||||
prisma_edge.db.litellm_memorytable.find_many.return_value = [row()]
|
||||
entries = await store(prisma_edge).search(MemorySearch(query="port demo port", limit=2, offset=3))
|
||||
assert entries[0].content == "Use port 8347"
|
||||
query = prisma_edge.db.litellm_memorytable.find_many.call_args.kwargs
|
||||
assert query["where"]["namespace"] == _IDENTITY.namespace("key")
|
||||
assert query["take"] == 2 and query["skip"] == 3
|
||||
clauses = query["where"]["AND"]
|
||||
assert len(clauses) == 2
|
||||
assert [clause["OR"][0]["value"]["contains"] for clause in clauses] == ["port", "demo"]
|
||||
assert query["order"] == [{"updated_at": "desc"}, {"memory_id": "asc"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_and_delete_cannot_address_another_namespace(prisma_edge: MagicMock) -> None:
|
||||
memory = store(prisma_edge)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await memory.read("foreign-entry")
|
||||
assert exc.value.status_code == 404
|
||||
assert prisma_edge.db.litellm_memorytable.find_first.call_args.kwargs["where"] == {
|
||||
"namespace": _IDENTITY.namespace("key"),
|
||||
"memory_id": "foreign-entry",
|
||||
}
|
||||
prisma_edge.db.litellm_memorypolicy.find_many.return_value = [_POLICY.model_copy(update={"activation": "disabled"})]
|
||||
assert await memory.delete("entry")
|
||||
assert prisma_edge.db.litellm_memorytable.delete_many.call_args.kwargs["where"] == {
|
||||
"namespace": _IDENTITY.namespace("key"),
|
||||
"memory_id": "entry",
|
||||
}
|
||||
with pytest.raises(HTTPException) as inactive:
|
||||
await memory.read("entry")
|
||||
assert inactive.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_identical_capture_is_idempotent_and_new_capture_has_scoped_identity(prisma_edge: MagicMock) -> None:
|
||||
table = prisma_edge.db.litellm_memorytable
|
||||
table.create.return_value = row()
|
||||
saved = await store(prisma_edge).capture(_CAPTURE)
|
||||
assert saved.content == _CAPTURE.content
|
||||
data = table.create.call_args.kwargs["data"]
|
||||
assert data["namespace"] == _IDENTITY.namespace("key") and data["user_id"] == "owner" and data["team_id"] == "team"
|
||||
assert data["key"].startswith("memory-v2:" + _IDENTITY.namespace("key") + ":")
|
||||
table.find_unique.return_value = row()
|
||||
assert await store(prisma_edge).capture(_CAPTURE) == saved
|
||||
table.create.assert_awaited_once()
|
||||
table.update_many.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("race", ["missing", "foreign", "stale", "concurrent"])
|
||||
async def test_capture_rejects_stale_or_conflicting_replacements(prisma_edge: MagicMock, race: str) -> None:
|
||||
table = prisma_edge.db.litellm_memorytable
|
||||
table.find_unique.return_value = (
|
||||
None if race == "missing" else row(namespace="foreign" if race == "foreign" else _IDENTITY.namespace("key"))
|
||||
)
|
||||
table.update_many.return_value = 0 if race == "concurrent" else 1
|
||||
correction = _CAPTURE.model_copy(
|
||||
update={
|
||||
"content": "Use port 8348",
|
||||
"expected_revision": _NOW - timedelta(seconds=1) if race == "stale" else _NOW,
|
||||
}
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await store(prisma_edge).capture(correction)
|
||||
assert exc.value.status_code == 409
|
||||
table.create.assert_not_awaited()
|
||||
if race == "concurrent":
|
||||
where = table.update_many.call_args.kwargs["where"]
|
||||
assert (
|
||||
where["namespace"] == _IDENTITY.namespace("key")
|
||||
and where["updated_at"] == _NOW
|
||||
and where["value"] == _CAPTURE.content
|
||||
)
|
||||
assert "metadata" in where
|
||||
else:
|
||||
table.update_many.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_replacement_reads_confirmed_updated_row(prisma_edge: MagicMock) -> None:
|
||||
table = prisma_edge.db.litellm_memorytable
|
||||
table.find_unique.return_value = row()
|
||||
table.find_first.return_value = row(value="Use port 8348", updated_at=_NOW + timedelta(seconds=1))
|
||||
result = await store(prisma_edge).capture(
|
||||
_CAPTURE.model_copy(update={"content": "Use port 8348", "expected_revision": _NOW})
|
||||
)
|
||||
assert result.content == "Use port 8348" and result.updated_at > _NOW
|
||||
table.update_many.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_argument_errors_are_recoverable_but_revocation_aborts(prisma_edge: MagicMock) -> None:
|
||||
memory = store(prisma_edge)
|
||||
invalid = await execute_memory_tool(
|
||||
memory, {"id": "a", "name": "litellm_memory_capture", "arguments": {"key": "missing-fields"}}
|
||||
)
|
||||
assert invalid.context == "" and "error" in invalid.output
|
||||
missing = await execute_memory_tool(
|
||||
memory, {"id": "a", "name": "litellm_memory_read", "arguments": {"memory_id": "missing"}}
|
||||
)
|
||||
assert missing.output == {"error": "Memory not found", "status": 404}
|
||||
unknown = await execute_memory_tool(memory, {"id": "a", "name": "other_tool", "arguments": {}})
|
||||
assert unknown.output == {"error": "Unknown memory tool"}
|
||||
prisma_edge.db.litellm_memorypolicy.find_many.return_value = []
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await execute_memory_tool(memory, {"id": "a", "name": "litellm_memory_search", "arguments": {}})
|
||||
assert exc.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_loop_is_bounded_and_search_results_reach_followup(prisma_edge: MagicMock) -> None:
|
||||
prisma_edge.db.litellm_memorytable.find_many.return_value = [row()]
|
||||
model = AsyncMock(
|
||||
return_value={"content": [{"type": "tool_use", "id": "search", "name": "litellm_memory_search", "input": {}}]}
|
||||
)
|
||||
context = await run_memory_tools(
|
||||
{"messages": [{"role": "user", "content": "My demo port?"}]}, "anthropic_messages", store(prisma_edge), model
|
||||
)
|
||||
assert model.await_count == 3 and len(context) == 3
|
||||
assert all("8347" in reference for reference in context)
|
||||
continuation = model.call_args_list[1].args[0]["messages"]
|
||||
assert continuation[-1]["content"][0]["tool_use_id"] == "search"
|
||||
assert "8347" in continuation[-1]["content"][0]["content"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("bad_id,count", [(True, 1), (False, 9)])
|
||||
async def test_invalid_model_calls_are_rejected_before_storage(
|
||||
prisma_edge: MagicMock, bad_id: bool, count: int
|
||||
) -> None:
|
||||
model = AsyncMock(
|
||||
return_value={
|
||||
"content": [
|
||||
{"type": "tool_use", "id": "" if bad_id else str(i), "name": "litellm_memory_search", "input": {}}
|
||||
for i in range(count)
|
||||
]
|
||||
}
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await run_memory_tools({"messages": []}, "anthropic_messages", store(prisma_edge), model)
|
||||
assert exc.value.status_code == 502
|
||||
prisma_edge.db.litellm_memorytable.find_many.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("writer_unavailable", [False, True])
|
||||
async def test_replica_lag_cannot_authorize_memory_after_primary_revocation(
|
||||
prisma_edge: MagicMock, writer_unavailable: bool
|
||||
) -> None:
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
|
||||
|
||||
writer = MagicMock(spec=PrismaWrapper)
|
||||
reader = MagicMock(spec=PrismaWrapper)
|
||||
writer.litellm_memorypolicy = SimpleNamespace(
|
||||
find_many=AsyncMock(return_value=[_POLICY.model_copy(update={"activation": "disabled"})])
|
||||
)
|
||||
reader.litellm_memorypolicy = SimpleNamespace(find_many=AsyncMock(return_value=[_POLICY]))
|
||||
writer.litellm_memorypreference = SimpleNamespace(find_unique=AsyncMock(return_value=None))
|
||||
writer.litellm_memorytable = prisma_edge.db.litellm_memorytable
|
||||
writer.is_connected = MagicMock(return_value=False)
|
||||
reader.is_connected = MagicMock(return_value=False)
|
||||
routed = RoutingPrismaWrapper(writer, reader)
|
||||
if writer_unavailable:
|
||||
writer.connect = AsyncMock(side_effect=RuntimeError("primary unavailable"))
|
||||
reader.connect = AsyncMock()
|
||||
await routed.connect()
|
||||
client = SimpleNamespace(db=routed)
|
||||
access = await resolve_memory_access(client, _IDENTITY)
|
||||
assert not access.active
|
||||
reader.litellm_memorypolicy.find_many.assert_not_awaited()
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await MemoryStore(client, MemoryAccess(_IDENTITY, _POLICY, False)).capture(_CAPTURE)
|
||||
assert exc.value.status_code == 403
|
||||
writer.litellm_memorytable.create.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unconfigured_gate_caches_presence_without_caching_authorization(prisma_edge: MagicMock) -> None:
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy.memory.policy import gateway_memory_is_configured
|
||||
|
||||
cache = DualCache()
|
||||
policies = prisma_edge.db.litellm_memorypolicy.find_many
|
||||
policies.return_value = []
|
||||
assert not await gateway_memory_is_configured(prisma_edge, cache)
|
||||
assert not await gateway_memory_is_configured(prisma_edge, cache)
|
||||
policies.assert_awaited_once()
|
||||
assert policies.call_args.kwargs == {"take": 1}
|
||||
enabled_cache = DualCache()
|
||||
policies.return_value = [_POLICY]
|
||||
assert await gateway_memory_is_configured(prisma_edge, enabled_cache)
|
||||
policies.return_value = [_POLICY.model_copy(update={"activation": "disabled"})]
|
||||
assert await gateway_memory_is_configured(prisma_edge, enabled_cache)
|
||||
assert not (await resolve_memory_access(prisma_edge, _IDENTITY)).active
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gateway_rounds_keep_separate_limiter_contexts_and_original_client_tools(prisma_edge: MagicMock) -> None:
|
||||
import asyncio
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import claim_request_stash_for_data, get_request_stash
|
||||
from litellm.proxy.memory.gateway import prepare_gateway_memory
|
||||
|
||||
provider = FastAPI()
|
||||
observed = []
|
||||
deferred = []
|
||||
release = asyncio.Event()
|
||||
prisma_edge.db.litellm_memorytable.find_many.return_value = [row()]
|
||||
|
||||
@provider.post("/v1/messages")
|
||||
async def model(request: Request):
|
||||
body = await request.json()
|
||||
call_id = body["litellm_call_id"]
|
||||
stash = claim_request_stash_for_data(body)
|
||||
observed.append((call_id, stash, body))
|
||||
|
||||
async def logged_owner():
|
||||
await release.wait()
|
||||
return get_request_stash().owner_litellm_call_id
|
||||
|
||||
deferred.append(asyncio.create_task(logged_owner()))
|
||||
return {"content": [{"type": "tool_use", "id": "search", "name": "litellm_memory_search", "input": {}}]}
|
||||
|
||||
original = {
|
||||
"model": "demo",
|
||||
"stream": True,
|
||||
"max_tokens": 1,
|
||||
"messages": [{"role": "user", "content": "My port?"}],
|
||||
"tools": [{"name": "client_tool"}],
|
||||
"tool_choice": {"type": "tool", "name": "client_tool"},
|
||||
}
|
||||
request = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"path": "/v1/messages",
|
||||
"query_string": b"",
|
||||
"headers": [(b"authorization", b"Bearer test"), (b"idempotency-key", b"visible-only")],
|
||||
}
|
||||
)
|
||||
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()):
|
||||
prepared = await prepare_gateway_memory(original, request, auth, "anthropic_messages")
|
||||
release.set()
|
||||
owners = await asyncio.gather(*deferred)
|
||||
assert len(observed) == 3 and len({id(stash) for _, stash, _ in observed}) == 3
|
||||
assert owners == [call_id for call_id, _, _ in observed]
|
||||
assert len(set(owners)) == 3
|
||||
assert get_request_stash() is before
|
||||
assert all(body["stream"] is False and body["max_tokens"] == 2048 for _, _, body in observed)
|
||||
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)
|
||||
|
|
@ -235,3 +235,14 @@ def test_every_app_mount_is_assigned_to_a_component():
|
|||
f"Add them to GATEWAY_MOUNT_PATHS, BACKEND_MOUNT_PATHS, or serve them "
|
||||
f"from the UI container:\n " + "\n ".join(sorted(unassigned))
|
||||
)
|
||||
|
||||
|
||||
def test_memory_v2_policies_stay_on_backend_while_own_entries_are_available_on_gateway():
|
||||
gateway = _component_paths(app.router.routes, GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES)
|
||||
backend = _component_paths(app.router.routes, BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES)
|
||||
for path in ("/v2/memory/policies", "/v2/memory/policies/{policy_id}"):
|
||||
assert path in backend
|
||||
assert path not in gateway
|
||||
for path in ("/v2/memory/status", "/v2/memory/preference", "/v2/memory/entries", "/v2/memory/entries/{memory_id}"):
|
||||
assert path in gateway
|
||||
assert path in backend
|
||||
|
|
|
|||
|
|
@ -384,7 +384,7 @@ describe("Sidebar (leftnav)", () => {
|
|||
});
|
||||
});
|
||||
|
||||
// Workflow Runs, Memory and Guardrails Monitor render a shell and then 401
|
||||
// Workflow Runs and Guardrails Monitor render a shell and then 401
|
||||
// for every non-proxy-admin role, because their page-load routes sit outside
|
||||
// internal_user_routes / self_managed_routes. Cost Optimization does not:
|
||||
// its primary call is /user/daily/activity, which every role may make, so
|
||||
|
|
@ -406,7 +406,7 @@ describe("Sidebar (leftnav)", () => {
|
|||
mockUseAuthorized.mockReset();
|
||||
});
|
||||
|
||||
it("hides Workflow Runs and Memory from an internal user under Agentic", async () => {
|
||||
it("shows self-service Memory and hides Workflow Runs from internal users", async () => {
|
||||
mockUseAuthorized.mockReturnValue(authFor("internal"));
|
||||
renderWithProviders(<Sidebar {...defaultProps} />);
|
||||
|
||||
|
|
@ -419,7 +419,7 @@ describe("Sidebar (leftnav)", () => {
|
|||
expect(screen.getByText("Agents")).toBeInTheDocument();
|
||||
});
|
||||
expect(screen.queryByText("Workflow Runs")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("Memory")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Memory")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// An org admin's session role is "Org Admin", which no capability list
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue