From 39b98fdb895f58e8e7845585a07c6ecc8204fbde Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 14 Sep 2026 14:05:26 -0700 Subject: [PATCH] fix(memory): preserve round directives and component route discovery --- litellm/proxy/_lazy_features.py | 5 - litellm/proxy/_lazy_openapi_snapshot.json | 1074 ----------------- litellm/proxy/memory/gateway.py | 13 + litellm/proxy/proxy_server.py | 3 + .../proxy/memory/test_memory_v2_boundaries.py | 38 + 5 files changed, 54 insertions(+), 1079 deletions(-) diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index c7a8b4e0967..dd1180b30ad 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -56,11 +56,6 @@ 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", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7c98f4f0ea5..7a110eff080 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -31265,1080 +31265,6 @@ } } }, - "memory_v2": { - "components": { - "schemas": { - "HTTPValidationError": { - "properties": { - "detail": { - "items": { - "$ref": "#/components/schemas/ValidationError" - }, - "title": "Detail", - "type": "array" - } - }, - "title": "HTTPValidationError", - "type": "object" - }, - "MemoryCapture": { - "additionalProperties": false, - "properties": { - "certainty": { - "default": "observed", - "enum": [ - "user_stated", - "observed", - "inferred" - ], - "title": "Certainty", - "type": "string" - }, - "content": { - "maxLength": 8000, - "minLength": 1, - "title": "Content", - "type": "string" - }, - "evidence": { - "maxLength": 2000, - "minLength": 1, - "title": "Evidence", - "type": "string" - }, - "expected_revision": { - "anyOf": [ - { - "format": "date-time", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Expected Revision" - }, - "key": { - "maxLength": 160, - "minLength": 1, - "pattern": "^[a-zA-Z0-9_.-]+$", - "title": "Key", - "type": "string" - }, - "kind": { - "default": "context", - "enum": [ - "workflow", - "decision", - "correction", - "learning", - "context", - "disagreement" - ], - "title": "Kind", - "type": "string" - }, - "scope": { - "default": "", - "maxLength": 200, - "title": "Scope", - "type": "string" - }, - "source": { - "default": "", - "maxLength": 1000, - "title": "Source", - "type": "string" - }, - "title": { - "maxLength": 200, - "minLength": 1, - "title": "Title", - "type": "string" - }, - "when_to_use": { - "default": "", - "maxLength": 700, - "title": "When To Use", - "type": "string" - } - }, - "required": [ - "key", - "title", - "content", - "evidence" - ], - "title": "MemoryCapture", - "type": "object" - }, - "MemoryEntry": { - "properties": { - "actor": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Actor" - }, - "actor_name": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Actor Name" - }, - "certainty": { - "default": "observed", - "enum": [ - "user_stated", - "observed", - "inferred" - ], - "title": "Certainty", - "type": "string" - }, - "content": { - "title": "Content", - "type": "string" - }, - "created_at": { - "anyOf": [ - { - "format": "date-time", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Created At" - }, - "evidence": { - "title": "Evidence", - "type": "string" - }, - "key": { - "title": "Key", - "type": "string" - }, - "kind": { - "default": "context", - "enum": [ - "workflow", - "decision", - "correction", - "learning", - "context", - "disagreement" - ], - "title": "Kind", - "type": "string" - }, - "memory_id": { - "title": "Memory Id", - "type": "string" - }, - "scope": { - "default": "", - "title": "Scope", - "type": "string" - }, - "source": { - "default": "", - "title": "Source", - "type": "string" - }, - "title": { - "title": "Title", - "type": "string" - }, - "updated_at": { - "format": "date-time", - "title": "Updated At", - "type": "string" - }, - "when_to_use": { - "default": "", - "title": "When To Use", - "type": "string" - } - }, - "required": [ - "memory_id", - "key", - "title", - "content", - "evidence", - "updated_at" - ], - "title": "MemoryEntry", - "type": "object" - }, - "MemoryPolicy": { - "additionalProperties": false, - "properties": { - "activation": { - "enum": [ - "disabled", - "opt_in", - "automatic" - ], - "title": "Activation", - "type": "string" - }, - "policy_id": { - "title": "Policy Id", - "type": "string" - }, - "scope": { - "default": "user", - "enum": [ - "key", - "user", - "team", - "project", - "organization" - ], - "title": "Scope", - "type": "string" - }, - "target_id": { - "maxLength": 256, - "minLength": 1, - "title": "Target Id", - "type": "string" - }, - "target_type": { - "enum": [ - "gateway", - "organization", - "team", - "project", - "user", - "key" - ], - "title": "Target Type", - "type": "string" - }, - "updated_at": { - "format": "date-time", - "title": "Updated At", - "type": "string" - }, - "updated_by": { - "title": "Updated By", - "type": "string" - } - }, - "required": [ - "target_type", - "target_id", - "activation", - "policy_id", - "updated_at", - "updated_by" - ], - "title": "MemoryPolicy", - "type": "object" - }, - "MemoryPolicyInput": { - "additionalProperties": false, - "properties": { - "activation": { - "enum": [ - "disabled", - "opt_in", - "automatic" - ], - "title": "Activation", - "type": "string" - }, - "scope": { - "default": "user", - "enum": [ - "key", - "user", - "team", - "project", - "organization" - ], - "title": "Scope", - "type": "string" - }, - "target_id": { - "maxLength": 256, - "minLength": 1, - "title": "Target Id", - "type": "string" - }, - "target_type": { - "enum": [ - "gateway", - "organization", - "team", - "project", - "user", - "key" - ], - "title": "Target Type", - "type": "string" - } - }, - "required": [ - "target_type", - "target_id", - "activation" - ], - "title": "MemoryPolicyInput", - "type": "object" - }, - "MemoryPreference": { - "additionalProperties": false, - "properties": { - "enabled": { - "title": "Enabled", - "type": "boolean" - } - }, - "required": [ - "enabled" - ], - "title": "MemoryPreference", - "type": "object" - }, - "MemoryStatus": { - "properties": { - "activation": { - "enum": [ - "disabled", - "opt_in", - "automatic" - ], - "title": "Activation", - "type": "string" - }, - "active": { - "title": "Active", - "type": "boolean" - }, - "opted_in": { - "title": "Opted In", - "type": "boolean" - }, - "policy_id": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Policy Id" - }, - "scope": { - "anyOf": [ - { - "enum": [ - "key", - "user", - "team", - "project", - "organization" - ], - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Scope" - }, - "user_id": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "User Id" - }, - "user_name": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "User Name" - } - }, - "required": [ - "active", - "activation", - "scope", - "opted_in", - "policy_id" - ], - "title": "MemoryStatus", - "type": "object" - }, - "ValidationError": { - "properties": { - "ctx": { - "title": "Context", - "type": "object" - }, - "input": { - "title": "Input" - }, - "loc": { - "items": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "integer" - } - ] - }, - "title": "Location", - "type": "array" - }, - "msg": { - "title": "Message", - "type": "string" - }, - "type": { - "title": "Error Type", - "type": "string" - } - }, - "required": [ - "loc", - "msg", - "type" - ], - "title": "ValidationError", - "type": "object" - } - } - }, - "paths": { - "/v2/memory/entries": { - "get": { - "operationId": "list_entries_v2_memory_entries_get", - "parameters": [ - { - "in": "query", - "name": "query", - "required": false, - "schema": { - "default": "", - "maxLength": 500, - "title": "Query", - "type": "string" - } - }, - { - "in": "query", - "name": "limit", - "required": false, - "schema": { - "default": 20, - "maximum": 20, - "minimum": 1, - "title": "Limit", - "type": "integer" - } - }, - { - "in": "query", - "name": "offset", - "required": false, - "schema": { - "default": 0, - "minimum": 0, - "title": "Offset", - "type": "integer" - } - }, - { - "in": "query", - "name": "key_id", - "required": false, - "schema": { - "anyOf": [ - { - "pattern": "^[a-f0-9]{64}$", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Key Id" - } - }, - { - "in": "query", - "name": "before_updated_at", - "required": false, - "schema": { - "anyOf": [ - { - "format": "date-time", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Before Updated At" - } - }, - { - "in": "query", - "name": "before_memory_id", - "required": false, - "schema": { - "anyOf": [ - { - "maxLength": 128, - "minLength": 1, - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Before Memory Id" - } - } - ], - "responses": { - "200": { - "content": { - "application/json": { - "schema": { - "items": { - "$ref": "#/components/schemas/MemoryEntry" - }, - "title": "Response List Entries V2 Memory Entries Get", - "type": "array" - } - } - }, - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "List Entries", - "tags": [ - "memory_v2" - ] - }, - "post": { - "operationId": "capture_entry_v2_memory_entries_post", - "parameters": [ - { - "in": "query", - "name": "key_id", - "required": false, - "schema": { - "anyOf": [ - { - "pattern": "^[a-f0-9]{64}$", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Key Id" - } - } - ], - "requestBody": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/MemoryCapture" - } - } - }, - "required": true - }, - "responses": { - "200": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/MemoryEntry" - } - } - }, - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "Capture Entry", - "tags": [ - "memory_v2" - ] - } - }, - "/v2/memory/entries/{memory_id}": { - "delete": { - "operationId": "delete_entry_v2_memory_entries__memory_id__delete", - "parameters": [ - { - "in": "path", - "name": "memory_id", - "required": true, - "schema": { - "title": "Memory Id", - "type": "string" - } - }, - { - "in": "query", - "name": "key_id", - "required": false, - "schema": { - "anyOf": [ - { - "pattern": "^[a-f0-9]{64}$", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Key Id" - } - } - ], - "responses": { - "204": { - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "Delete Entry", - "tags": [ - "memory_v2" - ] - } - }, - "/v2/memory/policies": { - "get": { - "operationId": "list_policies_v2_memory_policies_get", - "parameters": [ - { - "in": "query", - "name": "target_type", - "required": false, - "schema": { - "anyOf": [ - { - "enum": [ - "gateway", - "organization", - "team", - "project", - "user", - "key" - ], - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Target Type" - } - }, - { - "in": "query", - "name": "target_id", - "required": false, - "schema": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Target Id" - } - }, - { - "in": "query", - "name": "offset", - "required": false, - "schema": { - "default": 0, - "minimum": 0, - "title": "Offset", - "type": "integer" - } - } - ], - "responses": { - "200": { - "content": { - "application/json": { - "schema": { - "items": { - "$ref": "#/components/schemas/MemoryPolicy" - }, - "title": "Response List Policies V2 Memory Policies Get", - "type": "array" - } - } - }, - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "List Policies", - "tags": [ - "memory_v2" - ] - }, - "put": { - "operationId": "set_policy_v2_memory_policies_put", - "requestBody": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/MemoryPolicyInput" - } - } - }, - "required": true - }, - "responses": { - "200": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/MemoryPolicy" - } - } - }, - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "Set Policy", - "tags": [ - "memory_v2" - ] - } - }, - "/v2/memory/policies/{policy_id}": { - "delete": { - "operationId": "delete_policy_v2_memory_policies__policy_id__delete", - "parameters": [ - { - "in": "path", - "name": "policy_id", - "required": true, - "schema": { - "title": "Policy Id", - "type": "string" - } - } - ], - "responses": { - "204": { - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "Delete Policy", - "tags": [ - "memory_v2" - ] - } - }, - "/v2/memory/preference": { - "get": { - "operationId": "get_preference_v2_memory_preference_get", - "parameters": [ - { - "in": "query", - "name": "key_id", - "required": false, - "schema": { - "anyOf": [ - { - "pattern": "^[a-f0-9]{64}$", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Key Id" - } - } - ], - "responses": { - "200": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/MemoryPreference" - } - } - }, - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "Get Preference", - "tags": [ - "memory_v2" - ] - }, - "put": { - "operationId": "set_preference_v2_memory_preference_put", - "parameters": [ - { - "in": "query", - "name": "key_id", - "required": false, - "schema": { - "anyOf": [ - { - "pattern": "^[a-f0-9]{64}$", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Key Id" - } - } - ], - "requestBody": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/MemoryPreference" - } - } - }, - "required": true - }, - "responses": { - "200": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/MemoryPreference" - } - } - }, - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "Set Preference", - "tags": [ - "memory_v2" - ] - } - }, - "/v2/memory/status": { - "get": { - "operationId": "get_status_v2_memory_status_get", - "parameters": [ - { - "in": "query", - "name": "key_id", - "required": false, - "schema": { - "anyOf": [ - { - "pattern": "^[a-f0-9]{64}$", - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Key Id" - } - } - ], - "responses": { - "200": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/MemoryStatus" - } - } - }, - "description": "Successful Response" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - } - }, - "security": [ - { - "APIKeyHeader": [] - } - ], - "summary": "Get Status", - "tags": [ - "memory_v2" - ] - } - } - } - }, "policies": { "components": { "schemas": { diff --git a/litellm/proxy/memory/gateway.py b/litellm/proxy/memory/gateway.py index 09e85b16db8..2fc7433e296 100644 --- a/litellm/proxy/memory/gateway.py +++ b/litellm/proxy/memory/gateway.py @@ -137,8 +137,21 @@ class GatewayMemoryLoop: async def _call(self) -> AsyncGenerator[bytes, None]: self.stream.begin_round() + # Claude output directives control the next generated turn. Repeat them + # on outgoing rounds without adding pending directives to saved history. + directives: Final = ( + transcript_items(self.original, self.route)[-self.replaced_input :] if self.replaced_input else () + ) + messages: Final = transcript_items(self.data, self.route) body: Final = { # mutable-ok: Native provider JSON containers. **self.data, + **( + { # mutable-ok: Native provider JSON containers. + "messages": [*messages, *directives], # mutable-ok: Provider request JSON. + } + if directives and messages[-len(directives) :] != directives + else {} # mutable-ok: Native provider JSON containers. + ), "cache": { # mutable-ok: Native provider JSON containers. **object_value(self.data.get("cache")), "no-cache": True, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 23de604a2e9..63a47eeeb7b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -18802,6 +18802,9 @@ app.include_router(auto_router_management_router) app.include_router(tag_management_router) app.include_router(workflow_management_router) app.include_router(memory_router) +from litellm.proxy.memory.management import router as memory_v2_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) 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 f886d44cee7..54d4dfbffcf 100644 --- a/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py +++ b/tests/test_litellm/proxy/memory/test_memory_v2_boundaries.py @@ -579,6 +579,44 @@ async def test_trailing_system_messages_survive_client_tool_continuation(prisma_ assert original["messages"] == [prefix, directive] +@pytest.mark.asyncio +async def test_claude_output_directives_reach_search_answer_and_reflection_rounds(prisma_edge: MagicMock) -> None: + provider = FastAPI() + observed = [] + directive = {"role": "system", "content": [], "output_config": {"effort": "low"}} + + @provider.post("/v1/messages") + async def model(incoming: Request): + body = await incoming.json() + observed.append(body) + assert body["messages"][-1] == directive + if len(observed) == 1: + content = [ + {"type": "tool_use", "id": "search", "name": "litellm_memory_search", "input": {"query": "demo"}} + ] + elif len(observed) == 2: + content = [{"type": "text", "text": "The port is 8347"}] + else: + content = [ + {"type": "tool_use", "id": "reflect", "name": "litellm_memory_capture", "input": {"observations": []}} + ] + return { + "id": "msg_" + str(len(observed)), + "stop_reason": "end_turn" if len(observed) == 2 else "tool_use", + "content": content, + } + + original = {"messages": [{"role": "user", "content": "My demo port?"}, directive]} + loop = GatewayMemoryLoop(provider, request(), original, "anthropic_messages", store(prisma_edge)) + async for _ in loop.run(): + pass + assert len(observed) == 3 + assert observed[1]["messages"][-2]["content"][0]["tool_use_id"] == "search" + assert "reflect once" in observed[2]["messages"][-2]["content"] + assert observed[0]["messages"][0] == original["messages"][0] + assert sum(message["role"] == "system" for message in object_items(loop.data.get("messages"))) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("bad_id,count", [(True, 1), (False, 17)]) async def test_invalid_model_calls_are_rejected_before_storage(