diff --git a/deploy/memory-pilot/README.md b/deploy/memory-pilot/README.md index 85c86934486..3beda2eac3a 100644 --- a/deploy/memory-pilot/README.md +++ b/deploy/memory-pilot/README.md @@ -92,6 +92,8 @@ correct, or delete entries in Memory; callers can use the self-service API. 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 @@ -101,6 +103,8 @@ correct, or delete entries in Memory; callers can use the self-service API. 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. diff --git a/litellm/litellm_core_utils/prompt_templates/server_tool_responses.py b/litellm/litellm_core_utils/prompt_templates/server_tool_responses.py index d939972c3c9..31a750aa79a 100644 --- a/litellm/litellm_core_utils/prompt_templates/server_tool_responses.py +++ b/litellm/litellm_core_utils/prompt_templates/server_tool_responses.py @@ -76,13 +76,15 @@ def public_tool_response( } -def combined_usage(usages: Sequence[Mapping[str, object]]) -> Mapping[str, object]: +def combined_usage(usages: Sequence[Mapping[str, object]], depth: int = 0) -> Mapping[str, object]: + if depth > 16: + raise ValueError("Server tool usage nesting exceeds 16 levels") names: Final = frozenset(key for usage in usages for key in usage) def combined(name: str) -> object: values: Final = tuple(usage[name] for usage in usages if usage.get(name) is not None) if any(isinstance(value, dict) for value in values): - return combined_usage(tuple(object_value(value) for value in values)) + return combined_usage(tuple(object_value(value) for value in values), depth + 1) numbers: Final = tuple( value for value in values if isinstance(value, (int, float)) and not isinstance(value, bool) ) diff --git a/litellm/proxy/memory/content.py b/litellm/proxy/memory/content.py index 14492866894..846afce1edd 100644 --- a/litellm/proxy/memory/content.py +++ b/litellm/proxy/memory/content.py @@ -43,7 +43,7 @@ def fuzzy_memories( (entry.content.casefold(), 0.15), ) indexed: Final = tuple( - (text, weight, tuple(frozenset(match.group() for match in _TOKEN.finditer(text)))) + (text, weight, tuple(dict.fromkeys(match.group() for match in _TOKEN.finditer(text)))[:256]) for text, weight in fields ) scores: Final = tuple( diff --git a/litellm/proxy/memory/continuation.py b/litellm/proxy/memory/continuation.py index 48f8c017388..96a679cf708 100644 --- a/litellm/proxy/memory/continuation.py +++ b/litellm/proxy/memory/continuation.py @@ -47,15 +47,17 @@ def _empty_array(value: object) -> bool: return isinstance(value, list) and not value -def _canonical(value: object) -> object: +def _canonical(value: object, depth: int = 0) -> object: + if depth > 64: + raise HTTPException(status_code=400, detail="Memory conversation nesting exceeds 64 levels") if isinstance(value, dict): return { # mutable-ok: Native provider JSON containers. - key: _canonical(item) + key: _canonical(item, depth + 1) for key, item in _OBJECT.validate_python(value).items() if key not in ("cache_control",) and item is not None and not _empty_array(item) } if isinstance(value, (list, tuple)): - return tuple(_canonical(item) for item in _ITEMS.validate_python(value)) + return tuple(_canonical(item, depth + 1) for item in _ITEMS.validate_python(value)) return value diff --git a/litellm/proxy/memory/management.py b/litellm/proxy/memory/management.py index ca28535c2e7..31e4d52dea4 100644 --- a/litellm/proxy/memory/management.py +++ b/litellm/proxy/memory/management.py @@ -1,4 +1,4 @@ -from typing import Final +from typing import Annotated, Final from fastapi import APIRouter, Depends, HTTPException, Query, Response @@ -30,6 +30,7 @@ from litellm.types.memory_v2 import ( MemoryPolicy, MemoryPolicyInput, MemoryPreference, + MemoryQuery, MemorySearch, MemoryStatus, MemoryTarget, @@ -241,7 +242,7 @@ async def access_for_key(auth: UserAPIKeyAuth, key_id: str | None) -> MemoryAcce @router.get("/entries", response_model=list[MemoryEntry]) async def list_entries( - query: str = Query("", max_length=500), + query: Annotated[MemoryQuery, Query(max_length=500)] = "", limit: int = Query(20, ge=1, le=20), offset: int = Query(0, ge=0), key_id: str | None = Query(None, pattern=r"^[a-f0-9]{64}$"), diff --git a/litellm/proxy/memory/responses.py b/litellm/proxy/memory/responses.py index 59407ffe85e..5154aadd158 100644 --- a/litellm/proxy/memory/responses.py +++ b/litellm/proxy/memory/responses.py @@ -5,9 +5,11 @@ from urllib.parse import quote from fastapi import HTTPException, Request from pydantic import TypeAdapter from starlette.responses import JSONResponse, Response +from starlette.types import ASGIApp from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.memory.continuation import MemoryContinuations +from litellm.proxy.memory.store import MemoryStore from litellm.proxy.memory.transport import gateway_round _OBJECT: Final = TypeAdapter(dict[str, object]) @@ -25,20 +27,28 @@ async def memory_response_operation( store: Final = await gateway_memory_store(auth) if store is None: raise HTTPException(status_code=404, detail="Memory response not found or expired") + return await serve_memory_response(response_id, request, route, store, app) + + +async def serve_memory_response( + response_id: str, request: Request, route: str, store: MemoryStore, app: ASGIApp +) -> Response: continuations: Final = MemoryContinuations(store, "aresponses") patch: Final = await continuations.load_response(response_id) if patch is None or patch.response is None or not patch.upstream_ids: raise HTTPException(status_code=404, detail="Memory response not found or expired") if route == "aget_responses": return JSONResponse(patch.response) - if route == "adelete_responses": - await store.authorize_namespace(write=True) - ids: Final = patch.upstream_ids if route == "adelete_responses" else patch.upstream_ids[-1:] + if route == "alist_input_items": + raise HTTPException( + status_code=501, + detail="Input history is unavailable for gateway memory responses; retain the original client input", + ) + await store.authorize_namespace(write=True) async def dispatch(identifier: str) -> Mapping[str, object]: - suffix: Final = "/input_items" if route == "alist_input_items" else "" - path: Final = "/v1/responses/" + identifier + suffix - raw_path: Final = ("/v1/responses/" + quote(identifier, safe="") + suffix).encode() + path: Final = "/v1/responses/" + identifier + raw_path: Final = ("/v1/responses/" + quote(identifier, safe="")).encode() inner: Final = Request( { # mutable-ok: Native ASGI or JSON payload. **request.scope, @@ -54,16 +64,15 @@ async def memory_response_operation( ) as call: start: Final = await call.started if start.status >= 400: - if route == "adelete_responses" and start.status == 404: + if start.status == 404: return { # mutable-ok: Native ASGI or JSON payload. } raise HTTPException(status_code=start.status, detail="The upstream response operation failed") content: Final = await call.read() return _OBJECT.validate_json(content) - results: Final = tuple([await dispatch(identifier) for identifier in ids]) - if route == "alist_input_items": - return JSONResponse(results[-1]) + for identifier in patch.upstream_ids: + await dispatch(identifier) await continuations.delete_response(response_id, patch) return JSONResponse( { # mutable-ok: Native provider JSON containers. diff --git a/litellm/proxy/memory/store.py b/litellm/proxy/memory/store.py index 049c36dde24..86a45bbad96 100644 --- a/litellm/proxy/memory/store.py +++ b/litellm/proxy/memory/store.py @@ -114,7 +114,7 @@ class MemoryStore: await transaction.execute_raw("SELECT pg_advisory_xact_lock($1::bigint)", lock_key) table: Final = MemoryRepository(SimpleNamespace(db=transaction)).table saved: Final = tuple([await self._capture(capture, namespace, table) for capture in captures]) - await self.authorize_namespace(write=True) + await MemoryStore(SimpleNamespace(db=transaction), self.access).authorize_namespace(write=True) return saved async def _capture( diff --git a/litellm/types/memory_v2.py b/litellm/types/memory_v2.py index 8a2e02de3b4..c36242be2bc 100644 --- a/litellm/types/memory_v2.py +++ b/litellm/types/memory_v2.py @@ -1,7 +1,8 @@ +import re from datetime import datetime -from typing import Literal, TypeAlias +from typing import Annotated, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, model_validator from typing_extensions import Self MemoryTarget: TypeAlias = Literal["gateway", "organization", "team", "project", "user", "key"] @@ -11,6 +12,17 @@ MemoryKind: TypeAlias = Literal["workflow", "decision", "correction", "learning" MemoryCertainty: TypeAlias = Literal["user_stated", "observed", "inferred"] +def _validate_search_query(value: str) -> str: + if len(frozenset(re.findall(r"[\w-]{2,}", value.casefold()))) > 16: + raise ValueError("Memory search accepts at most 16 distinct search terms") + return value + + +MemoryQuery: TypeAlias = Annotated[ + str, AfterValidator(_validate_search_query), Field(description="Use at most 16 distinct search terms") +] + + class MemoryPolicyInput(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) @@ -88,7 +100,7 @@ class MemoryEntry(BaseModel): class MemorySearch(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) - query: str = Field(default="", max_length=500) + query: MemoryQuery = Field(default="", max_length=500) limit: int = Field(default=8, ge=1, le=20) offset: int = Field(default=0, ge=0) @@ -129,7 +141,7 @@ class MemoryCatalogRequest(BaseModel): class MemoryRecallRequest(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) - query: str = Field(default="", max_length=2000) + query: MemoryQuery = Field(default="", max_length=2000) scope: str | None = Field(default=None, max_length=200) limit: int = Field(default=8, ge=1, le=30) diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index e9f87ba6cae..6d0aaf127f4 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -2,6 +2,8 @@ import ast import os IGNORE_FUNCTIONS = [ + "_canonical", # Memory transcript JSON walk raises at depth 64; excessive nesting is tested. + "combined_usage", # Server tool usage JSON walk raises at depth 16; cycles and excessive nesting are tested. "_format_type", "_remove_additional_properties", "_remove_strict_from_schema", diff --git a/tests/test_litellm/proxy/memory/test_content.py b/tests/test_litellm/proxy/memory/test_content.py index d109267afcf..2d98deac2ca 100644 --- a/tests/test_litellm/proxy/memory/test_content.py +++ b/tests/test_litellm/proxy/memory/test_content.py @@ -2,9 +2,20 @@ from datetime import datetime, timezone from typing import Final import pytest +from pydantic import ValidationError from litellm.proxy.memory.content import fuzzy_memories, redact_memory -from litellm.types.memory_v2 import MemoryEntry +from litellm.types.memory_v2 import MemoryEntry, MemoryRecallRequest, MemorySearch + + +@pytest.mark.parametrize("request_type", (MemoryRecallRequest, MemorySearch)) +def test_all_search_requests_reject_excessive_distinct_terms( + request_type: type[MemoryRecallRequest] | type[MemorySearch], +) -> None: + query: Final = ",".join(f"query{index}" for index in range(17)) + with pytest.raises(ValidationError, match="at most 16 distinct search terms"): + request_type(query=query) + assert request_type(query=" ".join(["repeat"] * 20)).query @pytest.mark.parametrize("query", ("autorouter clasifier rationle", "rout clasif", "clasifier")) diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py index 6436a142db7..2ad2554a25b 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py @@ -2,6 +2,7 @@ import json from datetime import datetime, timedelta, timezone +from functools import reduce from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -9,12 +10,19 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import FastAPI, HTTPException, Request from prisma.models import LiteLLM_MemoryTable +from starlette.responses import JSONResponse +from starlette.types import Receive, Scope, Send +from litellm.litellm_core_utils.prompt_templates.server_tool_responses import ( + combined_usage, + executable_server_calls, + object_items, +) +from litellm.proxy.memory.continuation import MemoryContinuation, MemoryContinuations, prefix_hashes from litellm.proxy.memory.gateway import GatewayMemoryLoop from litellm.proxy.memory.knowledge import MEMORY_TOOL_NAMES, execute_memory_tool -from litellm.litellm_core_utils.prompt_templates.server_tool_responses import executable_server_calls, object_items -from litellm.proxy.memory.continuation import MemoryContinuation, MemoryContinuations, prefix_hashes from litellm.proxy.memory.policy import MemoryAccess, MemoryIdentity, resolve_memory_access +from litellm.proxy.memory.responses import serve_memory_response from litellm.proxy.memory.store import MemoryStore from litellm.types.memory_v2 import MemoryCapture, MemoryPolicy, MemorySearch @@ -75,6 +83,102 @@ def row(**changes: object) -> LiteLLM_MemoryTable: ) +def test_nested_conversation_and_cyclic_usage_fail_with_bounded_errors() -> None: + nested = reduce(lambda value, _: {"nested": value}, range(70), {}) + with pytest.raises(HTTPException, match="nesting exceeds") as exc: + prefix_hashes(({"role": "user", "content": nested},), "acompletion") + assert exc.value.status_code == 400 + usage: dict[str, object] = {} + usage["details"] = usage + with pytest.raises(ValueError, match="nesting exceeds"): + combined_usage((usage,)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["get", "input_items", "missing"]) +async def test_saved_response_reads_are_scoped_and_never_return_internal_input( + prisma_edge: MagicMock, operation: str +) -> None: + patch = MemoryContinuation( + replaces=1, + response={"id": "resp_litellm_memory_test", "output": [{"type": "message", "content": []}]}, + upstream_ids=("native=one",), + ) + prisma_edge.db.litellm_memorycontinuation.find_first.return_value = ( + None if operation == "missing" else SimpleNamespace(payload=patch.model_dump()) + ) + + async def app(scope: Scope, receive: Receive, send: Send) -> None: + pytest.fail("Public reads must not fetch the hidden upstream transcript") + + request = Request({"type": "http", "method": "GET", "path": "/v1/responses/resp_litellm_memory_test"}) + route = "alist_input_items" if operation == "input_items" else "aget_responses" + if operation == "get": + response = await serve_memory_response("resp_litellm_memory_test", request, route, store(prisma_edge), app) + assert isinstance(response, JSONResponse) and json.loads(response.body) == patch.response + else: + with pytest.raises(HTTPException) as exc: + await serve_memory_response("resp_litellm_memory_test", request, route, store(prisma_edge), app) + assert exc.value.status_code == (404 if operation == "missing" else 501) + where = prisma_edge.db.litellm_memorycontinuation.find_first.call_args.kwargs["where"] + assert where["namespace"] == _IDENTITY.namespace("key") and where["key_id"] == _IDENTITY.key_id + assert where["expires_at"]["gt"] <= datetime.now(timezone.utc) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["success", "already_missing", "upstream_error", "readonly"]) +async def test_response_deletion_preserves_auth_paths_and_retry_state(prisma_edge: MagicMock, outcome: str) -> None: + patch = MemoryContinuation( + replaces=1, + response={"id": "resp_litellm_memory_test"}, + upstream_ids=("native=one", "native=two"), + transcript_anchor="transcript", + ) + prisma_edge.db.litellm_memorycontinuation.find_first.return_value = SimpleNamespace(payload=patch.model_dump()) + paths: list[str] = [] + + async def app(scope: Scope, receive: Receive, send: Send) -> None: + paths.append(scope["path"]) + assert scope["raw_path"] == scope["path"].replace("=", "%3D").encode() + assert Request(scope).headers["authorization"] == "Bearer synthetic-test-credential" + assert scope["method"] == "DELETE" and scope["query_string"] == b"api-version=test" + status = ( + 502 if outcome == "upstream_error" and len(paths) == 2 else 404 if outcome == "already_missing" else 200 + ) + await JSONResponse({"deleted": True}, status_code=status)(scope, receive, send) + + request = Request( + { + "type": "http", + "method": "DELETE", + "path": "/v1/responses/resp_litellm_memory_test", + "headers": [(b"authorization", b"Bearer synthetic-test-credential")], + "query_string": b"api-version=test", + } + ) + identity = MemoryIdentity("a" * 64, "owner", "team", "project", "org", outcome == "readonly") + if outcome in ("upstream_error", "readonly"): + with pytest.raises(HTTPException) as exc: + await serve_memory_response( + "resp_litellm_memory_test", request, "adelete_responses", store(prisma_edge, identity), app + ) + assert exc.value.status_code == (403 if outcome == "readonly" else 502) + prisma_edge.db.litellm_memorycontinuation.delete_many.assert_not_awaited() + else: + response = await serve_memory_response( + "resp_litellm_memory_test", request, "adelete_responses", store(prisma_edge), app + ) + assert isinstance(response, JSONResponse) + assert json.loads(response.body) == { + "id": "resp_litellm_memory_test", + "object": "response.deleted", + "deleted": True, + } + prisma_edge.db.litellm_memorycontinuation.delete_many.assert_awaited_once() + assert paths == ([] if outcome == "readonly" else ["/v1/responses/native=one", "/v1/responses/native=two"]) + prisma_edge.db.litellm_memorytable.delete_many.assert_not_awaited() + + @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"}) @@ -167,6 +271,33 @@ async def test_identical_capture_is_idempotent_and_new_capture_has_scoped_identi table.update_many.assert_not_awaited() +@pytest.mark.asyncio +@pytest.mark.parametrize("revoked", [False, True]) +async def test_capture_rechecks_policy_on_its_transaction_connection(prisma_edge: MagicMock, revoked: bool) -> None: + prisma_edge.db.litellm_memorytable.create.return_value = row() + prisma_edge.db.litellm_memorypolicy.find_many.side_effect = [ + [_POLICY], + RuntimeError("The only pooled connection belongs to the active transaction"), + ] + transaction = SimpleNamespace( + litellm_memorytable=prisma_edge.db.litellm_memorytable, + litellm_memorypolicy=SimpleNamespace( + find_many=AsyncMock( + return_value=[_POLICY.model_copy(update={"activation": "disabled"})] if revoked else [_POLICY] + ) + ), + litellm_memorypreference=prisma_edge.db.litellm_memorypreference, + execute_raw=AsyncMock(), + ) + prisma_edge.db.tx.return_value.__aenter__.return_value = transaction + if revoked: + with pytest.raises(HTTPException) as exc: + await store(prisma_edge).capture(_CAPTURE) + assert exc.value.status_code == 403 + else: + assert (await store(prisma_edge).capture(_CAPTURE)).content == _CAPTURE.content + + @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: diff --git a/tests/test_litellm/proxy/memory/test_memory_v2_management.py b/tests/test_litellm/proxy/memory/test_memory_v2_management.py index f503e6f751c..f4b09d23b0f 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_management.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_management.py @@ -6,9 +6,11 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import HTTPException +from fastapi import FastAPI, HTTPException +from fastapi.testclient import TestClient from litellm.proxy._types import UI_TEAM_ID, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.memory import management from litellm.proxy.memory.policy import MemoryIdentity, memory_digest from litellm.types.memory_v2 import MemoryCapture, MemoryPolicy, MemoryPolicyInput, MemoryPreference @@ -68,6 +70,17 @@ def auth(user: str = "owner", role: LitellmUserRoles = LitellmUserRoles.INTERNAL return UserAPIKeyAuth(token="a" * 64, user_id=user, user_role=role, team_id="team", org_id="explicit-org") +def test_management_search_rejects_excessive_terms_before_database_work(database: MagicMock) -> None: + app = FastAPI() + app.include_router(management.router) + app.dependency_overrides[user_api_key_auth] = auth + with TestClient(app) as client: + response = client.get("/v2/memory/entries", params={"query": ",".join(f"term{i}" for i in range(17))}) + assert response.status_code == 422 + assert "at most 16 distinct search terms" in response.text + database.db.litellm_memorytable.find_many.assert_not_awaited() + + def policy(**changes: object) -> MemoryPolicy: return MemoryPolicy.model_validate( {