diff --git a/deploy/memory-pilot/README.md b/deploy/memory-pilot/README.md deleted file mode 100644 index ae46c31c953..00000000000 --- a/deploy/memory-pilot/README.md +++ /dev/null @@ -1,141 +0,0 @@ -# Memory gateway pilot on Render - -Run this branch as an isolated forwarding gateway. Colleagues keep their existing -upstream LiteLLM key and model name and change only their gateway base URL. They -need no plugin or client-side memory tools. Memory is off by default. Each person enables it explicitly for their key; -returning to the original URL stops using and collecting pilot memory. - -Every model call uses that caller's upstream key. The upstream -gateway continues to enforce its model permissions, budgets, rate limits, and -guardrails. The pilot checks the key against the upstream model catalog, then -registers its hash as a local virtual key so LiteLLM's normal authentication and -memory authorization still apply. No upstream key or provider credential is -configured on Render. The administrator credential belongs only to this pilot. - -The forwarding pilot isolates memories by virtual key. Upstream management APIs -may deny ordinary keys access to user/team/org details, so the pilot does not -infer those identities from client metadata. Install the feature directly in an -organization's gateway to use its existing user/team/project/org policies. -A regular gateway deployment reuses its existing PostgreSQL database with normal -schema migrations. It does not need a separate memory database or vector service. -This forwarding pilot has a separate database for isolation. Its memories are not -automatically available on the original gateway. Sharing requires both deployments -to run this feature against the same database and authenticated namespace; the -pilot must not be connected to an older gateway's production database. - -## Create the service - -1. Create a separate Render Postgres 16 database in the same region as the web - service. Restrict public database access; use its internal connection URL. -2. Create a Python web service from this repository and the Memory V2 branch. - A Standard service and the smallest paid Postgres plan are sufficient starting - points for a small pilot. They incur Render hosting charges. -3. Set the build command to `bash deploy/memory-pilot/build.sh`, the start command - to `bash deploy/memory-pilot/start.sh`, and the health path to - `/health/readiness`. The build includes the dashboard from this branch. - Set the service's maximum shutdown delay to 300 seconds so active requests - can drain during a deployment. Uvicorn allows 290 seconds before cleanup -4. Set these environment variables in Render: - - | Variable | Value | - | --- | --- | - | `DATABASE_URL` | The new database's internal connection URL | - | `UPSTREAM_LITELLM_BASE_URL` | Your original gateway URL, without `/v1` | - | `LITELLM_MASTER_KEY` | A new random `sk-` administrator key | - | `LITELLM_SALT_KEY` | A separate random encryption secret; preserve it across deploys | - | `PYTHON_VERSION` | `3.12.14` | - | `NODE_VERSION` | `24.19.0` | - | `NEXT_TELEMETRY_DISABLED` | `1` | - | `PORT` | `4000` | - -5. After deployment, open `/ui/memory`, sign in as `admin` using the pilot's master - key, expand **Advanced settings**, and save a policy for **Whole gateway**, **Users choose whether to opt in**, - **Private to each virtual key**. This policy persists across restarts. Memory - stays disabled until an administrator enables it. - -Equivalent activation through the API, with secrets supplied in shell variables: - -```bash -curl --fail-with-body "$PILOT_URL/v2/memory/policies" \ - -H "Authorization: Bearer $PILOT_ADMIN_KEY" \ - -H 'Content-Type: application/json' \ - -X PUT \ - -d '{"target_type":"gateway","target_id":"*","activation":"opt_in","scope":"key"}' -``` - -Administrators can instead enable memory automatically, disable a particular registered key, -or disable the whole gateway. Under an opt-in policy, callers set their preference -with `PUT /v2/memory/preference` and `{"enabled":true}` using their own key. -In a normal gateway, signed-in users can also select their key in Memory and turn -on the switch. Pilot administrators can do this for a registered key. The switch -shows the actual state; turning it off stops saving and recall but keeps saved -memories visible. Memories appear newest first, with optional details. - -## Try it - -Set an OpenAI-compatible client's base URL to `https://YOUR-SERVICE.onrender.com/v1`. -For Claude Code, set `ANTHROPIC_BASE_URL` to `https://YOUR-SERVICE.onrender.com`. -Retain the same gateway key and model setting. First opt in using the preference -API above, or ask the pilot administrator to turn on memory for your key. - -In one conversation, say “Remember that my demo project is Cobalt Heron and its -staging port is 8347.” In a **new conversation**, ask “What is my demo project and -its staging port?” Check actual saved entries with `GET /v2/memory/entries` using -the same key. An unrelated key must not see them. Administrators can inspect, -correct, or delete entries in Memory; callers can use the self-service API. - -## Behavior and limits - -- Streaming keeps LiteLLM's configured SSE keepalives across silent memory - rounds. The pilot sends comments every 15 seconds of silence and disables - proxy buffering. A failure after streaming starts arrives as a native SSE - error; before streaming starts, HTTP errors retain their retry delay -- Model calls retain LiteLLM's normal timeout and retry settings. The separate - upstream credential check has a 20-second timeout. Deployments drain existing - requests for up to five minutes; requests still running after that can be - interrupted. Schedule pilot updates outside active office usage -- Supported surfaces: Chat Completions, Responses, and Anthropic Messages, - including their native streaming responses and client tool continuation. -- The selected model must support function calling. The actual answering model - receives catalog, fuzzy search, full-read, and observation-capture tools beside - its normal client tools. The gateway executes only its own memory tools. -- A request allows at most eight model rounds and sixteen memory calls per round. - One final reflection round can acknowledge an empty observation batch. Additional - rounds use the same model and caller budget, and add latency and token spend. -- Captures are immediately visible after a confirmed save. Each observation keeps - its title, relevance guidance, scope, kind, certainty, evidence, source, and actor. - Corrections append observations. Agents receive no memory deletion tool. -- Search uses weighted fuzzy matching over the authorized scope. There is no vector - database, extraction model, or nightly consolidation. -- Searches accept up to 16 distinct terms. Fuzzy matching checks up to 256 distinct - words per field; exact terms still match anywhere in the field. -- Fixed instructions and tool definitions preserve prompt-prefix caching after - warm-up. Dynamic catalogs and checkpoint IDs stay at the conversation tail. - Complete-response caching is bypassed for memory rounds on both gateways so - permission checks, retrieval, and capture execute against current state. -- Hidden tool continuations expire after 24 hours, hold at most one megabyte each, - and are limited to 1,000 per key and scope. They contain gateway-added fragments, - not another copy of the complete incoming transcript. Responses retrieval and - continuation use gateway-owned response IDs; deleting one removes its model - responses and temporary continuation records, not saved memories. -- `/input_items` returns 501 for gateway-owned response IDs. Retain the original - client input; the hidden provider transcript is not a public input history. -- Foreground requests with one completion are supported. Use modern tools instead - of legacy functions. The special Cursor conversion route, background responses, - multiple completions, and WebSocket inference are outside this implementation. -- 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. -- Each memory scope can hold up to 1,000 entries. Creation checks this limit - atomically; correction and deletion remain available when the scope is full. -- 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. -- Invalid tool arguments return errors to the model. Infrastructure and model - failures fail the request or stream instead of reporting a successful save. Administrators can disable memory to restore ordinary calls. -- Switching away or disabling memory stops automatic use; it does not delete - existing entries. Delete memories explicitly through Memory or the API. -- Shared upstream keys share a pilot namespace. Give each person a distinct key - when their memories must be private from each other. -- Other API surfaces are outside this forwarding pilot. Use the original gateway - for embeddings, images, realtime, batches, and administration. diff --git a/deploy/memory-pilot/build.sh b/deploy/memory-pilot/build.sh deleted file mode 100755 index 46daf947233..00000000000 --- a/deploy/memory-pilot/build.sh +++ /dev/null @@ -1,19 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail -# Render's preinstalled Rust toolchain directory is read-only. Maturin needs -# 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 -export PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH" -prisma generate --schema schema.prisma -( - cd ui/litellm-dashboard - npm ci - npm run build -) -rm -rf litellm/proxy/_experimental/out -cp -R ui/litellm-dashboard/out litellm/proxy/_experimental/out diff --git a/deploy/memory-pilot/hooks.py b/deploy/memory-pilot/hooks.py deleted file mode 100644 index eee41cca7d3..00000000000 --- a/deploy/memory-pilot/hooks.py +++ /dev/null @@ -1 +0,0 @@ -from pilot import forward_credential as forward_credential diff --git a/deploy/memory-pilot/pilot.py b/deploy/memory-pilot/pilot.py deleted file mode 100644 index c328115e3f7..00000000000 --- a/deploy/memory-pilot/pilot.py +++ /dev/null @@ -1,181 +0,0 @@ -"""An isolated office pilot that preserves upstream gateway credentials.""" - -import asyncio -import os -import secrets -from collections.abc import Callable -from contextvars import ContextVar -from typing import Final - -import httpx -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.caching.in_memory_cache import InMemoryCache -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.prompt_templates.server_tool_responses import object_value -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.proxy._types import UI_TEAM_ID, UserAPIKeyAuth, hash_token -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.proxy.memory.transport import in_gateway_round -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("/") -_CREDENTIAL: Final[ContextVar[str | None]] = ContextVar("memory_pilot_credential", default=None) -_INFERENCE: Final = frozenset( - ("/chat/completions", "/v1/chat/completions", "/responses", "/v1/responses", "/v1/messages") -) -_SELF_SERVICE: Final = frozenset(("/v2/memory/status", "/v2/memory/preference", "/v2/memory/entries")) - - -class ForwardCredential(CustomLogger): - async def async_pre_call_hook( - self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral - ) -> dict[str, object]: - credential: Final = _CREDENTIAL.get() - if credential is None: - raise HTTPException(status_code=403, detail="Use your upstream gateway key for model calls") - return { # mutable-ok: The gateway hook returns native provider request JSON. - **data, - "api_key": credential, - "api_base": _UPSTREAM, - **( - { # mutable-ok: The second gateway must receive its own cache controls in the provider body. - "extra_body": { # mutable-ok: The proxy provider forwards this JSON unchanged. - **object_value(data.get("extra_body")), - "cache": { - "no-cache": True, - "no-store": True, - }, # mutable-ok: Native upstream gateway cache controls. - }, - } - if in_gateway_round() - else {} - ), # mutable-ok: Hook payload is native JSON. - } - - -forward_credential: Final = ForwardCredential() - - -class PilotGateway: - def __init__(self, app: ASGIApp, clock: Callable[[], float] | None = None) -> None: - self.app = app - self.registered_validation_slots = asyncio.Semaphore(12) - self.enrollment_validation_slots = asyncio.Semaphore(4) - self.recently_validated = InMemoryCache(max_size_in_memory=1000, default_ttl=60, clock=clock) - 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.close() - return - if scope["type"] != "http": - await self.app(scope, receive, send) - return - request: Final = Request(scope, receive) - 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 "" - ) - from litellm.proxy.proxy_server import master_key, prisma_client - - if not credential or master_key and secrets.compare_digest(credential, master_key): - await self.app(scope, receive, send) - return - if prisma_client is None: - await JSONResponse({"error": "Pilot database unavailable"}, status_code=503)(scope, receive, send) - return - digest: Final = hash_token(credential) - tokens: Final = VerificationTokenRepository(prisma_client) - local_key: Final = await tokens.find_by_id(digest) - if local_key and local_key.team_id == UI_TEAM_ID: - await self.app(scope, receive, send) - return - if ( - not credential.startswith("sk-") - and ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(credential) is not None - ): - await self.app(scope, receive, send) - return - path: Final = request.url.path.rstrip("/") - memory_response: Final = request.method in ("GET", "DELETE") and path.startswith( - ("/v1/responses/resp_litellm_memory_", "/responses/resp_litellm_memory_") - ) - if ( - path not in _INFERENCE | _SELF_SERVICE | {"/models", "/v1/models"} - and not path.startswith("/v2/memory/entries/") - and not memory_response - ): - await JSONResponse( - {"error": "Upstream keys can only use inference and their own memories"}, status_code=403 - )(scope, receive, send) - return - validation_slots: Final = ( - self.registered_validation_slots - if self.recently_validated.get_cache(digest) - else self.enrollment_validation_slots - ) - try: - await asyncio.wait_for(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} - ) - if models.status_code in (401, 403): - self.recently_validated.delete_cache(digest) - elif models.is_success: - self.recently_validated.delete_cache(digest) - self.recently_validated.set_cache(digest, True) - except httpx.HTTPError: - await JSONResponse({"error": "Upstream gateway unavailable"}, status_code=503)(scope, receive, send) - return - finally: - validation_slots.release() - if models.is_error: - await JSONResponse({"error": "Upstream gateway rejected this key"}, status_code=models.status_code)( - scope, receive, send - ) - return - if path in ("/models", "/v1/models"): - await JSONResponse(models.json())(scope, receive, send) - return - await tokens.table.upsert( - where={"token": digest}, - data={ - "create": {"token": digest, "models": [], "key_alias": "Memory pilot " + digest[:8]}, - "update": {}, - }, - ) - token: Final = _CREDENTIAL.set(credential) - try: - await self.app(scope, receive, send) - finally: - _CREDENTIAL.reset(token) - - -def create_app() -> PilotGateway: - from litellm.proxy.proxy_server import app - - return PilotGateway(app) diff --git a/deploy/memory-pilot/proxy_config.yaml b/deploy/memory-pilot/proxy_config.yaml deleted file mode 100644 index 160f98e1d0b..00000000000 --- a/deploy/memory-pilot/proxy_config.yaml +++ /dev/null @@ -1,16 +0,0 @@ -model_list: - - model_name: "*" - litellm_params: - model: litellm_proxy/* - api_base: os.environ/UPSTREAM_LITELLM_BASE_URL - model_info: - supports_function_calling: true -litellm_settings: - callbacks: - - hooks.forward_credential - drop_params: true - turn_off_message_logging: true - sse_keepalive_ping_interval_seconds: 15 -general_settings: - master_key: os.environ/LITELLM_MASTER_KEY - store_model_in_db: true diff --git a/deploy/memory-pilot/start.sh b/deploy/memory-pilot/start.sh deleted file mode 100755 index 9eac610730b..00000000000 --- a/deploy/memory-pilot/start.sh +++ /dev/null @@ -1,9 +0,0 @@ -#!/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}" -exec uvicorn pilot:create_app --factory --host 0.0.0.0 --port "${PORT:-4000}" --timeout-graceful-shutdown 290 diff --git a/tests/test_litellm/proxy/memory/test_memory_pilot.py b/tests/test_litellm/proxy/memory/test_memory_pilot.py deleted file mode 100644 index 93b0b57889a..00000000000 --- a/tests/test_litellm/proxy/memory/test_memory_pilot.py +++ /dev/null @@ -1,139 +0,0 @@ -"""Bound unauthenticated upstream validation without replacing gateway authentication.""" - -import asyncio -import hashlib -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_unknown_keys_cannot_consume_registered_validation_capacity_and_slots_release( - 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) - clock = MagicMock(return_value=0.0) - gateway = module.PilotGateway(Starlette(), clock=clock) - entered = asyncio.Event() - release = asyncio.Event() - count = 0 - - async def upstream_get(*args: object, **kwargs: object) -> httpx.Response: - if kwargs.get("headers") == {"Authorization": "Bearer sk-established"}: - return httpx.Response(200, json={"data": [{"id": "model"}]}) - nonlocal count - count += 1 - if count == 4: - entered.set() - await release.wait() - raise httpx.ConnectError("unavailable") - - upstream = MagicMock(get=AsyncMock(side_effect=upstream_get)) - gateway.upstream = upstream - database = MagicMock() - - async def key_lookup(*, where: dict[str, str]) -> dict[str, str] | None: - digest = hashlib.sha256(b"sk-established").hexdigest() - return {"token": digest} if where == {"token": digest} else None - - database.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=key_lookup) - 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: - initial = await client.get("/v1/models", headers={"Authorization": "Bearer sk-established"}) - assert initial.status_code == 200 - pending = [ - asyncio.create_task(client.get("/v1/models", headers={"Authorization": f"Bearer sk-invalid-{i}"})) - for i in range(4) - ] - 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 == 5 - clock.return_value = 59.0 - established = await client.get("/v1/models", headers={"Authorization": "Bearer sk-established"}) - assert established.status_code == 200 and established.json() == {"data": [{"id": "model"}]} - clock.return_value = 61.0 - refreshed = await client.get("/v1/models", headers={"Authorization": "Bearer sk-established"}) - assert refreshed.status_code == 200 - clock.return_value = 122.0 - expired = await client.get("/v1/models", headers={"Authorization": "Bearer sk-established"}) - assert expired.status_code == 503 and expired.headers["retry-after"] == "1" - 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 == 8 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("rejection_status", [401, 403]) -async def test_rejected_enrolled_keys_lose_reserved_capacity_and_can_revalidate( - monkeypatch: pytest.MonkeyPatch, rejection_status: int -) -> 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_revoked_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()) - revoked = asyncio.Event() - block = asyncio.Event() - entered = asyncio.Event() - release = asyncio.Event() - count = 0 - - async def upstream_get(*args: object, **kwargs: object) -> httpx.Response: - if kwargs.get("headers") == {"Authorization": "Bearer sk-revoked"}: - if block.is_set(): - nonlocal count - count += 1 - if count == 4: - entered.set() - await release.wait() - if revoked.is_set(): - return httpx.Response(rejection_status, json={"error": "rejected"}) - return httpx.Response(200, json={"data": []}) - - gateway.upstream = MagicMock(get=AsyncMock(side_effect=upstream_get)) - database = MagicMock() - database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value={"token": "enrolled"}) - with patch.multiple( # test-quality-ok: Keep enrollment rows present while upstream revokes access. - "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: - for credential in ("sk-established", "sk-revoked"): - initial = await client.get("/v1/models", headers={"Authorization": f"Bearer {credential}"}) - assert initial.status_code == 200 - revoked.set() - rejected = await client.get("/v1/models", headers={"Authorization": "Bearer sk-revoked"}) - assert rejected.status_code == rejection_status - block.set() - pending = [ - asyncio.create_task(client.get("/v1/models", headers={"Authorization": "Bearer sk-revoked"})) - for _ in range(4) - ] - await asyncio.wait_for(entered.wait(), timeout=2) - refused = await client.get("/v1/models", headers={"Authorization": "Bearer sk-revoked"}) - assert refused.status_code == 503 and refused.headers["retry-after"] == "1" - established = await client.get("/v1/models", headers={"Authorization": "Bearer sk-established"}) - assert established.status_code == 200 - release.set() - assert all(response.status_code == rejection_status for response in await asyncio.gather(*pending)) - revoked.clear() - restored = await client.get("/v1/models", headers={"Authorization": "Bearer sk-revoked"}) - assert restored.status_code == 200 - revoked.set() - rechecked = await client.get("/v1/models", headers={"Authorization": "Bearer sk-revoked"}) - assert rechecked.status_code == rejection_status