scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 66

refactor(memory): remove standalone pilot deployment

Browse source
This commit is contained in:
moe-berri 2026-09-14 13:57:45 -07:00
parent 4010f471c8
commit f40f3cb677
7 changed files with 0 additions and 506 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

141
deploy/memory-pilot/README.md
View file

@ -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.

19
deploy/memory-pilot/build.sh
View file

@ -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

1
deploy/memory-pilot/hooks.py
View file

@ -1 +0,0 @@
from pilot import forward_credential as forward_credential

181
deploy/memory-pilot/pilot.py
View file

@ -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)

16
deploy/memory-pilot/proxy_config.yaml
View file

@ -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

9
deploy/memory-pilot/start.sh
View file

@ -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

139
tests/test_litellm/proxy/memory/test_memory_pilot.py
View file

@ -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