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

fix(memory): preserve Claude directives and bound continuation storage

Browse source
This commit is contained in:
moe-berri 2026-09-14 13:52:18 -07:00
parent 57b6496a1e
commit 4010f471c8
9 changed files with 1242 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,7 +1,6 @@
"""An isolated office pilot that preserves upstream gateway credentials."""
import asyncio
import hashlib
import os
import secrets
from collections.abc import Callable
@ -19,7 +18,7 @@ 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
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
@ -101,7 +100,7 @@ class PilotGateway:
if prisma_client is None:
await JSONResponse({"error": "Pilot database unavailable"}, status_code=503)(scope, receive, send)
return
digest: Final = hashlib.sha256(credential.encode()).hexdigest()
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:

19
litellm/litellm_core_utils/prompt_templates/server_tools.py
View file

@ -137,16 +137,33 @@ def _tool_name(tool: object) -> object:
return _OBJECT.validate_python(function).get("name") if isinstance(function, dict) else definition.get("name")
def trailing_system_messages(data: Mapping[str, object], route: ServerToolRoute) -> int:
if route != "anthropic_messages":
return 0
messages: Final = _items(data.get("messages"))
return next(
(
index
for index, message in enumerate(reversed(messages))
if not isinstance(message, dict) or _OBJECT.validate_python(message).get("role") != "system"
),
len(messages),
)
def append_server_reference(data: Mapping[str, object], route: ServerToolRoute, reference: str) -> Mapping[str, object]:
field: Final = "input" if route == "aresponses" else "messages"
messages: Final = _items(data.get(field))
insertion: Final = len(messages) - trailing_system_messages(data, route)
return { # mutable-ok: Provider wire format requires native JSON containers.
**data,
field: [ # mutable-ok: Provider wire format requires native JSON containers.
*_items(data.get(field)),
*messages[:insertion],
{ # mutable-ok: Provider wire format requires native JSON containers.
"role": "user",
"content": reference,
},
*messages[insertion:],
],
}

5
litellm/proxy/_lazy_features.py
View file

@ -56,6 +56,11 @@ class LazyFeature:
LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
LazyFeature(
name="memory_v2",
module_path="litellm.proxy.memory.management",
path_prefixes=("/v2/memory",),
),
LazyFeature(
name="guardrails",
module_path="litellm.proxy.guardrails.guardrail_endpoints",

1074
litellm/proxy/_lazy_openapi_snapshot.json
View file

File diff suppressed because it is too large Load diff

33
litellm/proxy/memory/continuation.py
View file

@ -19,7 +19,9 @@ from litellm.repositories.unit_of_work import prisma_transaction
_ITEMS: Final = TypeAdapter(tuple[object, ...])
_OBJECT: Final = TypeAdapter(dict[str, object])
_MAX_PATCH_BYTES: Final = 1024 * 1024
_MAX_PATCHES: Final = 1000
_MAX_PATCHES: Final = 256
_MAX_NAMESPACE_BYTES: Final = 32 * 1024 * 1024
_USAGE: Final = TypeAdapter(tuple[dict[str, int], ...])
async def cleanup_memory_continuations(prisma_client: object) -> None:
@ -204,31 +206,34 @@ class MemoryContinuations:
key_id: Final = self.store.access.identity.key_id or self.store.access.identity.user_id or ""
now: Final = datetime.now(timezone.utc)
async with prisma_transaction(self.store.prisma_client) as transaction:
lock_key: Final = int(memory_digest("memory-continuation-quota", namespace, key_id)[:16], 16) - (1 << 63)
lock_key: Final = int(memory_digest("memory-continuation-quota", namespace)[:16], 16) - (1 << 63)
await transaction.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key)
table: Final = MemoryContinuationRepository(SimpleNamespace(db=transaction)).table
await table.delete_many(
where={ # mutable-ok: Prisma query and write JSON.
"namespace": namespace,
"key_id": key_id,
"expires_at": { # mutable-ok: Prisma query and write JSON.
"lte": now
},
}
)
count: Final = await table.count(
where={ # mutable-ok: Prisma query and write JSON.
"namespace": namespace,
"key_id": key_id,
"id": { # mutable-ok: Prisma query and write JSON.
"not_in": [ # mutable-ok: Prisma query and write JSON.
identifier for identifier, _ in payloads
]
},
}
usage: Final = _USAGE.validate_python(
await transaction.query_raw(
"SELECT COUNT(*) FILTER (WHERE key_id = $2)::int AS key_count, "
"COALESCE(SUM(octet_length(payload::text)), 0) + "
"(SELECT COALESCE(SUM(octet_length(value::text)), 0) "
"FROM jsonb_array_elements($4::jsonb)) AS bytes "
'FROM "LiteLLM_MemoryContinuation" WHERE namespace = $1 AND NOT (id = ANY($3::text[]))',
namespace,
key_id,
[identifier for identifier, _ in payloads], # mutable-ok: Native Prisma array parameter.
"[" + ",".join(payload for _, payload in payloads) + "]",
)
)
if count + len(payloads) > _MAX_PATCHES:
if usage[0]["key_count"] + len(payloads) > _MAX_PATCHES:
raise HTTPException(status_code=429, detail="Too many active memory continuations for this key")
if usage[0]["bytes"] > _MAX_NAMESPACE_BYTES:
raise HTTPException(status_code=429, detail="Memory continuations exceed 32 megabytes for this scope")
for identifier, payload in payloads:
await table.upsert(
where={ # mutable-ok: Prisma query and write JSON.

7
litellm/proxy/memory/gateway.py
View file

@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.server_tools import (
append_server_reference,
continue_server_tools,
inject_server_tools,
trailing_system_messages,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.sse_keepalive import wrap_passthrough_sse_bytes_with_keepalive_pings
@ -68,6 +69,7 @@ class GatewayMemoryLoop:
self.checkpoint = memory_digest(store.access.namespace, *prefix_hashes(self.visible_input, route)[-1:])
self.data: Mapping[str, object] = data
self.baseline_length = 0
self.replaced_input = 0
self.reflected = store.access.identity.read_only or (
data.get("tool_choice") not in (None, "auto")
and object_value(data.get("tool_choice")).get("type") != "auto"
@ -115,7 +117,8 @@ class GatewayMemoryLoop:
functions,
MEMORY_READ_ONLY_WORKFLOW if self.store.access.identity.read_only else MEMORY_WORKFLOW,
)
self.baseline_length = len(transcript_items(injected, self.route))
self.replaced_input = trailing_system_messages(injected, self.route)
self.baseline_length = len(transcript_items(injected, self.route)) - self.replaced_input
catalog: Final = await memory_catalog(self.store, MemoryCatalogRequest(limit=12))
self.data = append_server_reference(
injected,
@ -203,7 +206,7 @@ class GatewayMemoryLoop:
visible: Final = response_messages(response, self.route)
anchors: Final = prefix_hashes((*self.visible_input, *visible), self.route)
patch: Final = MemoryContinuation(
replaces=len(visible),
replaces=len(visible) + self.replaced_input,
replacement=transcript_items(self.data, self.route)[self.baseline_length :],
upstream_ids=self.upstream_ids,
pending_results=self.pending_results,

2
litellm/proxy/proxy_server.py
View file

@ -582,7 +582,6 @@ from litellm.proxy.management_helpers.team_metadata_validation import (
TEAM_METADATA_VALIDATOR_REGISTRY,
parse_team_metadata_schema,
)
from litellm.proxy.memory.management import router as memory_v2_router
from litellm.proxy.memory.memory_endpoints import router as memory_router
from litellm.proxy.middleware.billable_request_metrics_middleware import (
BillableRequestMetricsMiddleware,
@ -18803,7 +18802,6 @@ app.include_router(auto_router_management_router)
app.include_router(tag_management_router)
app.include_router(workflow_management_router)
app.include_router(memory_router)
app.include_router(memory_v2_router)
app.include_router(plugin_router)
app.include_router(cost_tracking_settings_router)
app.include_router(router_settings_router)

102
tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py
View file

@ -53,6 +53,7 @@ def prisma_edge() -> MagicMock:
table.count = AsyncMock(return_value=0)
client.db.tx.return_value.__aenter__.return_value = client.db
client.db.execute_raw = AsyncMock()
client.db.query_raw = AsyncMock(return_value=[{"key_count": 0, "bytes": 0}])
continuations = client.db.litellm_memorycontinuation
continuations.find_many = AsyncMock(return_value=[])
continuations.find_first = AsyncMock(return_value=None)
@ -512,6 +513,72 @@ async def test_model_loop_is_bounded_and_search_results_reach_the_active_model(p
prisma_edge.db.litellm_memorycontinuation.upsert.assert_not_awaited()
@pytest.mark.asyncio
async def test_trailing_system_messages_survive_client_tool_continuation(prisma_edge: MagicMock) -> None:
provider = FastAPI()
observed = []
client_call = {"type": "tool_use", "id": "client_read", "name": "Read", "input": {"path": "README.md"}}
@provider.post("/v1/messages")
async def model(incoming: Request):
body = await incoming.json()
observed.append(body)
messages = body["messages"]
assert all(
message["role"] != "system" or messages[index + 1]["role"] == "assistant"
for index, message in enumerate(messages[:-1])
)
return {
"id": "msg_" + str(len(observed)),
"role": "assistant",
"type": "message",
"stop_reason": "tool_use" if len(observed) == 1 else "end_turn",
"content": [client_call] if len(observed) == 1 else [{"type": "text", "text": "Read complete"}],
}
prefix = {
"role": "user",
"content": [{"type": "text", "text": "Read README.md", "cache_control": {"type": "ephemeral"}}],
}
directive = {"role": "system", "content": "Use concise answers"}
original = {
"messages": [prefix, directive],
"tools": [{"name": "Read", "input_schema": {"type": "object"}}],
"tool_choice": {"type": "tool", "name": "Read"},
}
first = GatewayMemoryLoop(provider, request(), original, "anthropic_messages", store(prisma_edge))
async for _ in first.run():
pass
saved = prisma_edge.db.litellm_memorycontinuation.upsert.call_args.kwargs
prisma_edge.db.litellm_memorycontinuation.find_many.return_value = [
SimpleNamespace(id=saved["where"]["id"], payload=json.loads(saved["data"]["create"]["payload"]))
]
following = {
**original,
"tool_choice": {"type": "none"},
"messages": [
prefix,
directive,
{"role": "assistant", "content": [client_call]},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "client_read", "content": "File content"}],
},
directive,
],
}
second = GatewayMemoryLoop(provider, request(), following, "anthropic_messages", store(prisma_edge))
async for _ in second.run():
pass
assert len(observed) == 2
assert observed[0]["messages"][0] == observed[1]["messages"][0] == prefix
assert observed[1]["messages"].count(directive) == 2
assert observed[1]["messages"].count({"role": "assistant", "content": [client_call]}) == 1
assert observed[1]["messages"][-1] == directive
assert observed[1]["messages"][-3]["content"][0]["tool_use_id"] == "client_read"
assert original["messages"] == [prefix, directive]
@pytest.mark.asyncio
@pytest.mark.parametrize("bad_id,count", [(True, 1), (False, 17)])
async def test_invalid_model_calls_are_rejected_before_storage(
@ -894,3 +961,38 @@ async def test_full_scope_blocks_creation_but_permits_correction_and_reclaimed_c
table.create.return_value = row()
assert (await store(prisma_edge).capture(_CAPTURE)).memory_id == "entry"
table.create.assert_awaited_once()
@pytest.mark.asyncio
@pytest.mark.parametrize("key_count,used_bytes", [(256, 0), (0, 32 * 1024 * 1024 + 1)])
async def test_continuation_quota_rejects_excess_without_writing(
prisma_edge: MagicMock, key_count: int, used_bytes: int
) -> None:
prisma_edge.db.query_raw.return_value = [{"key_count": key_count, "bytes": used_bytes}]
with pytest.raises(HTTPException) as exc:
await MemoryContinuations(store(prisma_edge), "aresponses").save("response", MemoryContinuation(replaces=1))
assert exc.value.status_code == 429
prisma_edge.db.litellm_memorycontinuation.upsert.assert_not_awaited()
@pytest.mark.asyncio
async def test_continuation_quota_shares_namespace_lock_across_keys_and_allows_replacements(
prisma_edge: MagicMock,
) -> None:
user_policy = _POLICY.model_copy(update={"scope": "user"})
prisma_edge.db.litellm_memorypolicy.find_many.return_value = [user_policy]
prisma_edge.db.query_raw.return_value = [{"key_count": 255, "bytes": 32 * 1024 * 1024}]
other_key = MemoryIdentity("b" * 64, "owner", "team", "project", "org", False)
for identity in (_IDENTITY, other_key):
continuations = MemoryContinuations(
MemoryStore(prisma_edge, MemoryAccess(identity, user_policy, False)), "aresponses"
)
await continuations.save("response", MemoryContinuation(replaces=1, response={"text": "é漢字"}))
query = prisma_edge.db.query_raw.call_args.args
assert query[1:4] == (identity.namespace("user"), identity.key_id, [continuations.identifier("response")])
assert json.loads(query[4])[0]["response"]["text"] == "é漢字"
locks = prisma_edge.db.execute_raw.call_args_list
assert locks[0] == locks[1]
cleanup = prisma_edge.db.litellm_memorycontinuation.delete_many.call_args.kwargs["where"]
assert cleanup["namespace"] == _IDENTITY.namespace("user") and "key_id" not in cleanup
assert prisma_edge.db.litellm_memorycontinuation.upsert.await_count == 2

17
tests/test_litellm/proxy/memory/test_memory_v2_protocols.py
View file

@ -91,6 +91,23 @@ def test_anthropic_continuation_preserves_signed_thinking_and_matches_tool_resul
]
@pytest.mark.parametrize("directive_only", [False, True])
def test_memory_reference_keeps_anthropic_trailing_system_directives_valid(directive_only: bool) -> None:
prefix: Final = {
"role": "user",
"content": [{"type": "text", "text": "Read a file", "cache_control": {"type": "ephemeral"}}],
}
directive: Final = {
"role": "system",
"content": [] if directive_only else "Use concise answers",
"output_config": {"effort": "low"},
}
original: Final = {"messages": [prefix, directive]}
result: Final = append_server_reference(original, "anthropic_messages", "Untrusted stored context")
assert result["messages"] == [prefix, {"role": "user", "content": "Untrusted stored context"}, directive]
assert original["messages"] == [prefix, directive]
def test_responses_continuation_keeps_reasoning_and_function_call_output() -> None:
output: Final = [
{"type": "reasoning", "id": "reason-1", "encrypted_content": "opaque-provider-data"},