From 35ecb06beb45105fe13bfe7a99d090fab229da37 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sat, 8 Aug 2026 23:37:32 -0400 Subject: [PATCH 1/4] fix: preserve LiteLLM checkpoints after failed transitions --- .../litellm_proxy/README.md | 12 +-- .../context_compiler_precall_hook.py | 17 ++-- ...ler_precall_hook_with_directive_drafter.py | 17 ++-- python/tests/test_litellm_proxy_hooks.py | 89 ++++++++--------- python/tests/test_litellm_proxy_runtime.py | 40 +++----- ...st_litellm_proxy_with_directive_drafter.py | 96 +++++++++---------- 6 files changed, 120 insertions(+), 151 deletions(-) diff --git a/python/reference_integrations/litellm_proxy/README.md b/python/reference_integrations/litellm_proxy/README.md index 54dddcc..0bf1f86 100644 --- a/python/reference_integrations/litellm_proxy/README.md +++ b/python/reference_integrations/litellm_proxy/README.md @@ -22,10 +22,10 @@ Available hook files: turn with a fresh engine. - In explicit persistent mode, the hook resolves a session key, loads a saved checkpoint, restores the engine, processes the latest user turn once, and - saves the resulting checkpoint after every decision, including `clarify`. + saves the resulting checkpoint only after non-error decisions. - In stateless mode, no continuity is preserved across requests. -- If result is `clarify`, the proxy does not call the downstream model and - LiteLLM surfaces the clarification as an HTTP 400 response. +- If directive application fails, the proxy does not call the downstream model + and LiteLLM surfaces the rejection as an HTTP 400 response. - If result is `passthrough`, the proxy forwards the request normally. - If result is `update`, the proxy injects compiler state as a system message and then calls the model. @@ -49,7 +49,7 @@ The reference hooks support two explicit modes: - `persistent` - explicit mode - requires a stable session key - - preserves saved state and pending clarification across requests + - preserves saved authoritative state across requests - `stateless` - default mode - processes only the latest user turn @@ -217,8 +217,8 @@ Use `llama` only for LLM-only fallback drafting with Llama-family models. - In the directive-drafter hook, drafter state context now comes from restored checkpoint state rather than transcript-prefix reconstruction. - Compound directive-shaped input such as `use docker and prohibit peanuts` - should produce a local clarify response telling the user to submit each - directive separately, without mutating saved state or forwarding upstream. + should produce a local rejection telling the user to submit each directive + separately, without mutating saved state or forwarding upstream. ## Troubleshooting diff --git a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py index 28e19ca..639012a 100644 --- a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py +++ b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py @@ -4,8 +4,9 @@ - Resolve explicit persistent or stateless mode for the current request. - In persistent mode, restore compiler checkpoint by session key. - Process only the latest user turn exactly once. -- Save checkpoint after each decision, including clarify. -- If clarification is required, block upstream model call. +- Save checkpoints only after successful authoritative state transitions. +- If directive application fails, reject the current request without persisting + failed-turn engine state. - Otherwise inject compiled state guidance into a system message. """ @@ -151,18 +152,18 @@ async def async_pre_call_hook( else: decision = {"kind": DecisionKind.NO_DIRECTIVE, "message": None} + logger.debug("litellm_proxy: decision_kind=%s", decision["kind"]) + + if decision["kind"] == DecisionKind.ERROR: + logger.debug("litellm_proxy: rejecting_failed_application=true") + return decision.get("message") or "Request rejected." + if session.mode == MODE_PERSISTENT and session.session_key is not None: CHECKPOINT_STORE.save( session.session_key, checkpoint_to_jsonable(engine.export_json()), ) - logger.debug("litellm_proxy: decision_kind=%s", decision["kind"]) - - if decision["kind"] == DecisionKind.ERROR: - logger.debug("litellm_proxy: blocking_on_clarify=true") - return decision.get("message") or "Request rejected." - compiled_state = _snapshot_engine_state(engine) # For long-running conversations, you can optionally compact transcripts by removing user inputs that were compiled into state. See Demo 6. # noqa: E501 system_message: dict[str, object] = { diff --git a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py index b155ae2..f521e03 100644 --- a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py +++ b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py @@ -5,8 +5,9 @@ - In persistent mode, restore compiler checkpoint by session key. - Draft only the latest user message after restore. - Call ``engine.step(...)`` exactly once for the current turn. -- Save checkpoint after each decision, including clarify. -- If clarification is required, block upstream model call. +- Save checkpoints only after successful authoritative state transitions. +- If directive application fails, reject the current request without persisting + failed-turn engine state. - Otherwise inject compiled state guidance into a system message. """ @@ -281,18 +282,18 @@ async def async_pre_call_hook( else: decision = {"kind": DecisionKind.NO_DIRECTIVE, "message": None} + logger.debug("litellm_proxy: decision_kind=%s", decision["kind"]) + + if decision["kind"] == DecisionKind.ERROR: + logger.debug("litellm_proxy: rejecting_failed_application=true") + return decision.get("message") or "Request rejected." + if session.mode == MODE_PERSISTENT and session.session_key is not None: CHECKPOINT_STORE.save( session.session_key, checkpoint_to_jsonable(engine.export_json()), ) - logger.debug("litellm_proxy: decision_kind=%s", decision["kind"]) - - if decision["kind"] == DecisionKind.ERROR: - logger.debug("litellm_proxy: blocking_on_clarify=true") - return decision.get("message") or "Request rejected." - compiled_state = _snapshot_engine_state(engine) system_message: dict[str, object] = { "role": "system", diff --git a/python/tests/test_litellm_proxy_hooks.py b/python/tests/test_litellm_proxy_hooks.py index cdee1ce..ee8e3b3 100644 --- a/python/tests/test_litellm_proxy_hooks.py +++ b/python/tests/test_litellm_proxy_hooks.py @@ -132,81 +132,74 @@ def test_persistent_mode_restores_checkpoint_and_isolates_sessions(monkeypatch) assert "peanuts" not in str(other["messages"][0]["content"]) -def test_persistent_mode_saves_updated_authoritative_state_for_follow_up_turns( +def test_failed_application_rejects_request_and_does_not_persist_state( monkeypatch, ) -> None: - module = _load_proxy_module(monkeypatch, "litellm_proxy_pending") + module = _load_proxy_module(monkeypatch, "litellm_proxy_failed_application") module.CHECKPOINT_STORE.clear() hook = module.ContextCompilerPreCallHook() - clarify_data = { + rejected_data = { "model": "demo", "context_compiler_mode": "persistent", - "context_compiler_session_key": "chat-clarify", - "messages": [{"role": "user", "content": "use kubectl instead of docker"}], - } - confirm_data = { - "model": "demo", - "context_compiler_mode": "persistent", - "context_compiler_session_key": "chat-clarify", - "messages": [ - {"role": "user", "content": "use kubectl instead of docker"}, - {"role": "assistant", "content": "question asked"}, - {"role": "user", "content": "yes"}, - ], + "context_compiler_session_key": "chat-failed-apply", + "messages": [{"role": "user", "content": "change premise to formal tone"}], } - first = asyncio.run( - hook.async_pre_call_hook(None, None, clarify_data, "completion") - ) - second = asyncio.run( - hook.async_pre_call_hook(None, None, confirm_data, "completion") + result = asyncio.run( + hook.async_pre_call_hook(None, None, rejected_data, "completion") ) - assert first is clarify_data - assert second is confirm_data - checkpoint = module.CHECKPOINT_STORE.load("chat-clarify") - assert checkpoint is not None - assert checkpoint["policies"] == {"kubectl": "use"} + assert isinstance(result, str) + assert "No premise is set." in result + assert module.CHECKPOINT_STORE.load("chat-failed-apply") is None -def test_persistent_mode_does_not_treat_confirmation_text_as_removed_resume_flow( - monkeypatch, -) -> None: - module = _load_proxy_module(monkeypatch, "litellm_proxy_pending_no") +def test_failed_application_preserves_existing_checkpoint_state(monkeypatch) -> None: + module = _load_proxy_module(monkeypatch, "litellm_proxy_preserve_checkpoint") module.CHECKPOINT_STORE.clear() hook = module.ContextCompilerPreCallHook() - clarify_data = { + seed_data = { "model": "demo", "context_compiler_mode": "persistent", - "context_compiler_session_key": "chat-clarify-no", - "messages": [{"role": "user", "content": "use kubectl instead of docker"}], + "context_compiler_session_key": "chat-preserve-checkpoint", + "messages": [{"role": "user", "content": "use docker"}], } - reject_data = { + rejected_data = { "model": "demo", "context_compiler_mode": "persistent", - "context_compiler_session_key": "chat-clarify-no", - "messages": [ - {"role": "user", "content": "use kubectl instead of docker"}, - {"role": "assistant", "content": "question asked"}, - {"role": "user", "content": "no"}, - ], + "context_compiler_session_key": "chat-preserve-checkpoint", + "messages": [{"role": "user", "content": "prohibit docker"}], } - first = asyncio.run( - hook.async_pre_call_hook(None, None, clarify_data, "completion") + seed_result = asyncio.run( + hook.async_pre_call_hook(None, None, seed_data, "completion") ) - second = asyncio.run( - hook.async_pre_call_hook(None, None, reject_data, "completion") + assert seed_result is seed_data + + result = asyncio.run( + hook.async_pre_call_hook(None, None, rejected_data, "completion") ) - assert first is clarify_data - assert second is reject_data - checkpoint = module.CHECKPOINT_STORE.load("chat-clarify-no") + assert isinstance(result, str) + checkpoint = module.CHECKPOINT_STORE.load("chat-preserve-checkpoint") assert checkpoint is not None - assert checkpoint["policies"] == {"kubectl": "use"} - assert "docker" not in str(reject_data["messages"][0]["content"]) + assert checkpoint["policies"] == {"docker": "use"} + + follow_up_data = { + "model": "demo", + "context_compiler_mode": "persistent", + "context_compiler_session_key": "chat-preserve-checkpoint", + "messages": [{"role": "user", "content": "prohibit docker"}], + } + + follow_up_result = asyncio.run( + hook.async_pre_call_hook(None, None, follow_up_data, "completion") + ) + + assert isinstance(follow_up_result, str) + assert '"docker" is currently in use.' in follow_up_result def test_state_is_saved_after_update(monkeypatch) -> None: diff --git a/python/tests/test_litellm_proxy_runtime.py b/python/tests/test_litellm_proxy_runtime.py index 138f8e9..8c689e4 100644 --- a/python/tests/test_litellm_proxy_runtime.py +++ b/python/tests/test_litellm_proxy_runtime.py @@ -206,27 +206,19 @@ def _start_proxy_runtime( config_path.unlink(missing_ok=True) -def test_litellm_proxy_runtime_persists_state_without_removed_confirmation_flow( +def test_litellm_proxy_runtime_rejects_failed_application_without_persisting_resume_state( litellm_proxy_runtime_basic: _ProxyRuntime, litellm_runtime_stub: _ThreadedStubServer, ) -> None: response = _post_chat_completion( port=litellm_proxy_runtime_basic.port, - messages=[{"role": "user", "content": "use kubectl instead of docker"}], - session_key="runtime-basic-confirm", + messages=[{"role": "user", "content": "change premise to formal tone"}], + session_key="runtime-basic-reject", ) - assert response.status_code == 200 - assert response.json()["choices"][0]["message"]["content"] == "stubbed reply" - assert len(litellm_runtime_stub.captured_requests) == 1 - - forwarded_payload = litellm_runtime_stub.captured_requests[0] - forwarded_messages = forwarded_payload["messages"] - assert isinstance(forwarded_messages, list) - assert len(forwarded_messages) == 2 - assert forwarded_messages[1:] == [ - {"role": "user", "content": "use kubectl instead of docker"} - ] + assert response.status_code == 400 + assert "No premise is set." in response.text + assert litellm_runtime_stub.captured_requests == [] def test_litellm_proxy_runtime_forwards_allowed_request_with_contract( @@ -275,27 +267,19 @@ def test_litellm_proxy_runtime_forwards_allowed_request_with_contract( assert forwarded_messages[1:] == original_messages -def test_litellm_proxy_runtime_with_directive_drafter_persists_state_without_removed_confirmation_flow( +def test_litellm_proxy_runtime_with_directive_drafter_rejects_failed_application_without_persisting_resume_state( litellm_proxy_runtime_with_directive_drafter: _ProxyRuntime, litellm_runtime_stub: _ThreadedStubServer, ) -> None: response = _post_chat_completion( port=litellm_proxy_runtime_with_directive_drafter.port, - messages=[{"role": "user", "content": "use kubectl instead of docker"}], - session_key="runtime-drafter-confirm", + messages=[{"role": "user", "content": "change premise to formal tone"}], + session_key="runtime-drafter-reject", ) - assert response.status_code == 200 - assert response.json()["choices"][0]["message"]["content"] == "stubbed reply" - assert len(litellm_runtime_stub.captured_requests) == 1 - - forwarded_payload = litellm_runtime_stub.captured_requests[0] - forwarded_messages = forwarded_payload["messages"] - assert isinstance(forwarded_messages, list) - assert len(forwarded_messages) == 2 - assert forwarded_messages[1:] == [ - {"role": "user", "content": "use kubectl instead of docker"} - ] + assert response.status_code == 400 + assert "No premise is set." in response.text + assert litellm_runtime_stub.captured_requests == [] def test_litellm_proxy_runtime_with_directive_drafter_forwards_allowed_request_with_contract( diff --git a/python/tests/test_litellm_proxy_with_directive_drafter.py b/python/tests/test_litellm_proxy_with_directive_drafter.py index ac06013..1f50696 100644 --- a/python/tests/test_litellm_proxy_with_directive_drafter.py +++ b/python/tests/test_litellm_proxy_with_directive_drafter.py @@ -89,51 +89,34 @@ def test_drafter_output_applies_to_current_turn_only(monkeypatch) -> None: assert "peanuts" not in str(data["messages"][0]["content"]) -def test_persistent_mode_with_drafter_saves_updated_state_for_follow_up_turns( +def test_persistent_mode_with_drafter_rejects_failed_application_without_persisting( monkeypatch, ) -> None: - module = _load_module(monkeypatch, "litellm_proxy_with_drafter_pending") + module = _load_module(monkeypatch, "litellm_proxy_with_drafter_failed_apply") module.CHECKPOINT_STORE.clear() hook = module.ContextCompilerPreCallHookWithPreprocessor() drafted_inputs: list[str] = [] def fake_preprocess(message: str, state: dict[str, object] | None) -> str | None: drafted_inputs.append(message) - if message == "use kubectl instead of docker": - return None - return "use docker" + return None monkeypatch.setattr(module, "_preprocess_last_user_message", fake_preprocess) - first = { - "model": "demo", - "context_compiler_mode": "persistent", - "context_compiler_session_key": "chat-drafter-pending", - "messages": [{"role": "user", "content": "use kubectl instead of docker"}], - } - second = { + rejected_data = { "model": "demo", "context_compiler_mode": "persistent", - "context_compiler_session_key": "chat-drafter-pending", - "messages": [ - {"role": "user", "content": "use kubectl instead of docker"}, - {"role": "assistant", "content": "question asked"}, - {"role": "user", "content": "yes"}, - ], + "context_compiler_session_key": "chat-drafter-failed-apply", + "messages": [{"role": "user", "content": "change premise to formal tone"}], } - first_result = asyncio.run( - hook.async_pre_call_hook(None, None, first, "completion") - ) - second_result = asyncio.run( - hook.async_pre_call_hook(None, None, second, "completion") + result = asyncio.run( + hook.async_pre_call_hook(None, None, rejected_data, "completion") ) - assert first_result is first - assert second_result is second - assert drafted_inputs == ["use kubectl instead of docker", "yes"] - checkpoint = module.CHECKPOINT_STORE.load("chat-drafter-pending") - assert checkpoint is not None - assert checkpoint["policies"] == {"kubectl": "use", "docker": "use"} + assert isinstance(result, str) + assert "No premise is set." in result + assert drafted_inputs == ["change premise to formal tone"] + assert module.CHECKPOINT_STORE.load("chat-drafter-failed-apply") is None def test_missing_session_key_fails_clearly_in_persistent_mode(monkeypatch) -> None: @@ -193,50 +176,57 @@ def test_stateless_mode_has_no_cross_call_continuity(monkeypatch) -> None: assert "peanuts" not in str(second["messages"][0]["content"]) -def test_persistent_mode_with_drafter_does_not_resume_removed_confirmation_flow( +def test_persistent_mode_with_drafter_preserves_existing_checkpoint_on_failure( monkeypatch, ) -> None: - module = _load_module(monkeypatch, "litellm_proxy_with_drafter_pending_no") + module = _load_module(monkeypatch, "litellm_proxy_with_drafter_preserve_checkpoint") module.CHECKPOINT_STORE.clear() hook = module.ContextCompilerPreCallHookWithPreprocessor() - drafted_inputs: list[str] = [] def fake_preprocess(message: str, state: dict[str, object] | None) -> str | None: - drafted_inputs.append(message) return None monkeypatch.setattr(module, "_preprocess_last_user_message", fake_preprocess) - first = { + seed_data = { "model": "demo", "context_compiler_mode": "persistent", - "context_compiler_session_key": "chat-drafter-pending-no", - "messages": [{"role": "user", "content": "use kubectl instead of docker"}], + "context_compiler_session_key": "chat-drafter-preserve-checkpoint", + "messages": [{"role": "user", "content": "use docker"}], } - second = { + rejected_data = { "model": "demo", "context_compiler_mode": "persistent", - "context_compiler_session_key": "chat-drafter-pending-no", - "messages": [ - {"role": "user", "content": "use kubectl instead of docker"}, - {"role": "assistant", "content": "question asked"}, - {"role": "user", "content": "no"}, - ], + "context_compiler_session_key": "chat-drafter-preserve-checkpoint", + "messages": [{"role": "user", "content": "prohibit docker"}], } - first_result = asyncio.run( - hook.async_pre_call_hook(None, None, first, "completion") + seed_result = asyncio.run( + hook.async_pre_call_hook(None, None, seed_data, "completion") ) - second_result = asyncio.run( - hook.async_pre_call_hook(None, None, second, "completion") + assert seed_result is seed_data + + result = asyncio.run( + hook.async_pre_call_hook(None, None, rejected_data, "completion") ) - assert first_result is first - assert second_result is second - assert drafted_inputs == ["use kubectl instead of docker", "no"] - checkpoint = module.CHECKPOINT_STORE.load("chat-drafter-pending-no") + assert isinstance(result, str) + checkpoint = module.CHECKPOINT_STORE.load("chat-drafter-preserve-checkpoint") assert checkpoint is not None - assert checkpoint["policies"] == {"kubectl": "use"} - assert "docker" not in str(second["messages"][0]["content"]) + assert checkpoint["policies"] == {"docker": "use"} + + follow_up_data = { + "model": "demo", + "context_compiler_mode": "persistent", + "context_compiler_session_key": "chat-drafter-preserve-checkpoint", + "messages": [{"role": "user", "content": "prohibit docker"}], + } + + follow_up_result = asyncio.run( + hook.async_pre_call_hook(None, None, follow_up_data, "completion") + ) + + assert isinstance(follow_up_result, str) + assert '"docker" is currently in use.' in follow_up_result def test_normal_update_explicitly_saves_checkpoint(monkeypatch) -> None: From 5065d337860984fe2c4c49c41a37bbaa30ed9830 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sat, 8 Aug 2026 23:54:25 -0400 Subject: [PATCH 2/4] fix: remove OpenWebUI continuation semantics --- .../openwebui_pipe/README.md | 17 +- .../openwebui_pipe/open_webui_pipe.py | 24 ++- .../open_webui_pipe_with_directive_drafter.py | 21 ++- python/tests/test_openwebui_pipe.py | 134 ++++++++++++++-- ...t_openwebui_pipe_with_directive_drafter.py | 147 ++++++++++++++++-- 5 files changed, 308 insertions(+), 35 deletions(-) diff --git a/python/reference_integrations/openwebui_pipe/README.md b/python/reference_integrations/openwebui_pipe/README.md index 06182ca..cb98f60 100644 --- a/python/reference_integrations/openwebui_pipe/README.md +++ b/python/reference_integrations/openwebui_pipe/README.md @@ -10,7 +10,8 @@ without Directive Drafter preprocessing. - The pipe forwards normal chat turns to the backend model. - When compiler state is non-empty, passthrough includes exactly one compiler-owned `[[cc_state]]` system message in the forwarded request. -- Conflicting or ambiguous updates ask for clarification before state changes. +- Conflicting or invalid updates are rejected for the current request and do not + create resumable continuation state. - The pipe handles exact `show state` locally. The pipe treats near matches such as `show state please` as normal chat input. @@ -117,7 +118,7 @@ Suggested verification: - Send `use docker` and confirm you get `State updated: Use docker.` with trace showing a local turn - Send a normal prompt such as `what should I run?` and confirm trace shows a forwarded turn with compiler state included -- Send `use kubectl instead of docker` and confirm Open WebUI asks for clarification instead of changing state +- Send `use docker`, then `prohibit docker`, and confirm Open WebUI rejects the second request instead of changing state - Optionally send `show state` and confirm the state summary is returned locally Advanced check: @@ -136,7 +137,7 @@ Suggested verification: - Send `please use docker` and confirm either: - the Directive Drafter converts it into a local state update, or - trace shows the turn followed the normal compiler path without a silent state change -- Send `use kubectl instead of docker`, then reply `yes`, and confirm the saved clarification flow resumes locally +- Send `please use docker`, then `prohibit docker`, and confirm the second request is rejected without creating resumable state - Send `use docker and prohibit peanuts` and confirm the pipe responds locally that multiple directives are not supported and must be submitted separately - Send a normal prompt such as `what should I run?` and confirm trace shows a forwarded turn with compiler state included @@ -182,15 +183,15 @@ If you want a slightly broader manual pass: - prompt(s): `clear state` → `use docker` → `prohibit docker` - base model: generic Docker/prohibition guidance text - basic pipe: `'docker' is already in use. Only one policy per item is allowed. Use 'reset policies' to change it.` -- directive-drafter pipe: same conflict clarify -- why this matters: the app asks before applying a conflicting change. +- directive-drafter pipe: same deterministic rejection +- why this matters: the app rejects conflicting changes instead of creating a resumable flow. ### Case 3 - prompt(s): `clear state` → `use podman instead of docker` - base model: generic “how to switch to Podman” tutorial - basic pipe: `No exact policy found for "docker". Replacement requires an exact policy match...` -- directive-drafter pipe: same replacement clarify +- directive-drafter pipe: same deterministic rejection - why this matters: the app only replaces a policy when the old item already exists. ### Case 4 @@ -198,7 +199,7 @@ If you want a slightly broader manual pass: - prompt(s): `clear state` → `set premise to concise replies` - base model: accepts conversational style phrasing - basic pipe: `Did you mean 'set premise concise replies'?` -- directive-drafter pipe: same clarify (near-miss is not rewritten) +- directive-drafter pipe: same deterministic rejection (near-miss is not rewritten) - why this matters: near-miss text is not silently rewritten. ### Case 5 @@ -206,7 +207,7 @@ If you want a slightly broader manual pass: - prompt(s): `clear state` → `change premise concise replies` - base model: generic “please clarify changes” response - basic pipe: `Did you mean 'change premise to concise replies'?` -- directive-drafter pipe: same clarify (near-miss is passed through unchanged) +- directive-drafter pipe: same deterministic rejection (near-miss is passed through unchanged) - why this matters: the app waits for explicit, valid directive text before changing state. ## Compatibility diff --git a/python/reference_integrations/openwebui_pipe/open_webui_pipe.py b/python/reference_integrations/openwebui_pipe/open_webui_pipe.py index 5f07dc4..0a7b741 100644 --- a/python/reference_integrations/openwebui_pipe/open_webui_pipe.py +++ b/python/reference_integrations/openwebui_pipe/open_webui_pipe.py @@ -14,6 +14,8 @@ - Single Pipe Function for Open WebUI 0.8.x and 0.9.x. - In-memory per-process engine map keyed by chat key. - No persistence, no multi-worker coordination, no external storage. +- Failed transitions are rejected for the current request and do not leave + resumable in-memory engine state behind. """ import inspect @@ -111,6 +113,12 @@ def _snapshot_engine_state(engine: Engine) -> _EngineSnapshot: return {"premise": engine.premise, "policies": dict(engine.policies)} +def _restore_engine_from_snapshot(snapshot_json: str) -> Engine: + engine = create_engine() + engine.import_json(snapshot_json) + return engine + + def _render_compiler_state_block(state: _EngineSnapshot) -> str: """Render deterministic compiler-owned state block text. @@ -322,7 +330,7 @@ def _render_item_label(value: str) -> str: return re.sub(r"\s+", " ", value).strip().lower() -def _near_miss_directive_clarify(value: str) -> str | None: +def _near_miss_directive_rejection(value: str) -> str | None: normalized = re.sub(r"\s+", " ", value.strip()) lower = normalized.lower() @@ -392,7 +400,7 @@ def _is_administrative_update_input(user_input: str) -> bool: class Pipe: """Map Context Compiler decisions into Open WebUI pipe behavior. - - ``clarify`` returns plain text and skips model forwarding. + - failed transitions return plain-text rejection and skip model forwarding. - ``passthrough`` forwards with minimal mutation. - ``update`` returns deterministic local acknowledgement (no model call). """ @@ -635,7 +643,8 @@ async def pipe( - Bypass compiler for non-text or missing-user turns. - Resolve chat key and get/create per-chat engine. - Call ``engine.step(...)``. - - Map ``clarify`` / ``passthrough`` / ``update`` outcomes. + - Reject failed transitions without preserving resumable in-memory state. + - Map rejection / ``passthrough`` / ``update`` outcomes. """ raw_messages = body.get("messages") messages = ( @@ -669,6 +678,7 @@ async def pipe( return _render_show_state_summary(engine) state_before = _snapshot_engine_state(engine) + engine_snapshot_json = engine.export_json() logger.debug("pipe: engine_input=%r", latest_user_text) decision = engine.step(latest_user_text) if decision["kind"] == DecisionKind.ERROR: @@ -678,10 +688,13 @@ async def pipe( else: kind = DecisionKind.NO_DIRECTIVE.value logger.debug("pipe: decision=%s", kind) - near_miss_prompt = _near_miss_directive_clarify(latest_user_text) + near_miss_prompt = _near_miss_directive_rejection(latest_user_text) state_after = _snapshot_engine_state(engine) if decision["kind"] == DecisionKind.ERROR: + _ENGINES_BY_CHAT_KEY[chat_key] = _restore_engine_from_snapshot( + engine_snapshot_json + ) return self._with_trace( near_miss_prompt or decision["message"] or "", original_input=latest_user_text, @@ -695,6 +708,9 @@ async def pipe( near_miss_prompt is not None and decision["kind"] == DecisionKind.NO_DIRECTIVE ): + _ENGINES_BY_CHAT_KEY[chat_key] = _restore_engine_from_snapshot( + engine_snapshot_json + ) return self._with_trace( near_miss_prompt, original_input=latest_user_text, diff --git a/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py b/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py index 4e2b3ce..9c728a9 100644 --- a/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py +++ b/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py @@ -14,6 +14,8 @@ 3. Pass resulting directive (or original input) to `engine.step(...)` Core decision handling remains the same as the base integration. +Failed transitions are rejected for the current request and do not leave +resumable in-memory engine state behind. """ import inspect @@ -127,6 +129,12 @@ def _snapshot_engine_state(engine: Engine) -> _EngineSnapshot: return {"premise": engine.premise, "policies": dict(engine.policies)} +def _restore_engine_from_snapshot(snapshot_json: str) -> Engine: + engine = create_engine() + engine.import_json(snapshot_json) + return engine + + def _render_compiler_state_block(state: _EngineSnapshot) -> str: lines: list[str] = [_CC_MARKER] @@ -323,7 +331,7 @@ def _render_item_label(value: str) -> str: return re.sub(r"\s+", " ", value).strip().lower() -def _near_miss_directive_clarify(value: str) -> str | None: +def _near_miss_directive_rejection(value: str) -> str | None: normalized = re.sub(r"\s+", " ", value.strip()) lower = normalized.lower() @@ -927,6 +935,7 @@ async def pipe( return _render_show_state_summary(engine) state_before = _snapshot_engine_state(engine) + engine_snapshot_json = engine.export_json() preprocessd: str | None = None preprocess_error: str | None = None @@ -943,7 +952,7 @@ async def pipe( logger.debug("preprocessor: preprocessd=%r", preprocessd) # Preserve core behavior: if preprocess yields no directive, use raw user - # text so the compiler still decides clarify/passthrough/update. + # text so the compiler still decides rejection/passthrough/update. compile_input = preprocessd if preprocessd is not None else latest_user_text logger.debug("preprocessor: engine_input=%r", compile_input) @@ -955,10 +964,13 @@ async def pipe( else: kind = DecisionKind.NO_DIRECTIVE.value logger.debug("preprocessor: decision=%s", kind) - near_miss_prompt = _near_miss_directive_clarify(latest_user_text) + near_miss_prompt = _near_miss_directive_rejection(latest_user_text) state_after = _snapshot_engine_state(engine) if decision["kind"] == DecisionKind.ERROR: + _ENGINES_BY_CHAT_KEY[chat_key] = _restore_engine_from_snapshot( + engine_snapshot_json + ) return self._with_trace( near_miss_prompt or decision["message"] or "", original_input=latest_user_text, @@ -973,6 +985,9 @@ async def pipe( near_miss_prompt is not None and decision["kind"] == DecisionKind.NO_DIRECTIVE ): + _ENGINES_BY_CHAT_KEY[chat_key] = _restore_engine_from_snapshot( + engine_snapshot_json + ) return self._with_trace( near_miss_prompt, original_input=latest_user_text, diff --git a/python/tests/test_openwebui_pipe.py b/python/tests/test_openwebui_pipe.py index 635cbb7..79b53d5 100644 --- a/python/tests/test_openwebui_pipe.py +++ b/python/tests/test_openwebui_pipe.py @@ -199,10 +199,10 @@ async def _forward( assert forwarded == [] -def test_confirmation_text_is_not_treated_as_removed_resume_flow( +def test_failed_transition_is_rejected_and_follow_up_is_a_new_request( monkeypatch, ) -> None: - module = _load_module_with_stubs("owui_confirmation_resume", monkeypatch) + module = _load_module_with_stubs("owui_failed_transition_followup", monkeypatch) forwarded: list[dict[str, object]] = [] async def _forward( @@ -214,22 +214,31 @@ async def _forward( module.generate_chat_completion = _forward pipe = module.Pipe() pipe.valves.BASE_MODEL_ID = "base-model" - chat_id = "chat-confirm" + chat_id = "chat-failed-transition" - update = asyncio.run( + seed = asyncio.run( pipe.pipe( { "model": "pipe-model", - "messages": [ - {"role": "user", "content": "use docker instead of kubectl"} - ], + "messages": [{"role": "user", "content": "use docker"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + rejected = asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "prohibit docker"}], }, __user__={"id": "u1"}, __request__=object(), __chat_id__=chat_id, ) ) - resumed = asyncio.run( + follow_up = asyncio.run( pipe.pipe( {"model": "pipe-model", "messages": [{"role": "user", "content": "yes"}]}, __user__={"id": "u1"}, @@ -238,11 +247,60 @@ async def _forward( ) ) - assert update == "State updated: Use docker." - assert resumed == {"ok": True} + assert seed == "State updated: Use docker." + assert rejected == ( + '"docker" is currently in use.\nRemove or replace it before prohibiting it.' + ) + assert follow_up == {"ok": True} assert len(forwarded) == 1 +def test_failed_transition_does_not_create_resumable_engine_state(monkeypatch) -> None: + module = _load_module_with_stubs( + "owui_failed_transition_no_resume_state", monkeypatch + ) + pipe = module.Pipe() + pipe.valves.BASE_MODEL_ID = "base-model" + chat_id = "chat-no-resume-state" + + asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "use docker"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "prohibit docker"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + + show_state = asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "show state"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + + assert show_state == "Premise: none\nUse: docker\nProhibit: none" + + def test_exact_show_state_is_local_and_non_exact_forwards_normally(monkeypatch) -> None: module = _load_module_with_stubs("owui_show_state", monkeypatch) forwarded: list[dict[str, object]] = [] @@ -285,7 +343,7 @@ async def _forward( assert len(forwarded) == 1 -def test_near_miss_directive_clarify_returns_deterministic_text_and_skips_downstream( +def test_near_miss_directive_is_rejected_and_skips_downstream( monkeypatch, ) -> None: module = _load_module_with_stubs("owui_near_miss", monkeypatch) @@ -319,6 +377,60 @@ async def _forward( assert forwarded == [] +def test_near_miss_rejection_does_not_create_pending_state(monkeypatch) -> None: + module = _load_module_with_stubs("owui_near_miss_no_pending_state", monkeypatch) + forwarded: list[dict[str, object]] = [] + + async def _forward( + _: object, payload: dict[str, object], __: object + ) -> dict[str, object]: + forwarded.append(payload) + return {"ok": True} + + module.generate_chat_completion = _forward + pipe = module.Pipe() + pipe.valves.BASE_MODEL_ID = "base-model" + chat_id = "chat-near-miss-followup" + + rejected = asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [ + {"role": "user", "content": "set premise to concise replies"} + ], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + follow_up = asyncio.run( + pipe.pipe( + {"model": "pipe-model", "messages": [{"role": "user", "content": "yes"}]}, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + show_state = asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "show state"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + + assert rejected == "Invalid premise syntax.\nUse 'set premise '." + assert follow_up == {"ok": True} + assert show_state == "Premise: none\nUse: none\nProhibit: none" + assert len(forwarded) == 1 + + def test_passthrough_with_non_empty_state_injects_exactly_one_cc_state_system_message( monkeypatch, ) -> None: diff --git a/python/tests/test_openwebui_pipe_with_directive_drafter.py b/python/tests/test_openwebui_pipe_with_directive_drafter.py index 48136eb..7464667 100644 --- a/python/tests/test_openwebui_pipe_with_directive_drafter.py +++ b/python/tests/test_openwebui_pipe_with_directive_drafter.py @@ -128,10 +128,10 @@ async def fake_preprocess(*args, **kwargs): assert compile_inputs == ["use docker"] -def test_confirmation_text_is_not_treated_as_removed_pending_resume( +def test_failed_transition_is_rejected_and_follow_up_is_a_new_request( monkeypatch, ) -> None: - module = _load_module("owui_with_drafter_confirmation_followup", monkeypatch) + module = _load_module("owui_with_drafter_failed_transition_followup", monkeypatch) forwarded: list[dict[str, object]] = [] async def forward( @@ -149,7 +149,7 @@ async def update_draft(*args, **kwargs): return "use docker", None monkeypatch.setattr(module.Pipe, "_preprocess_user_input", update_draft) - update = asyncio.run( + seed = asyncio.run( pipe.pipe( { "model": "pipe-model", @@ -157,7 +157,7 @@ async def update_draft(*args, **kwargs): }, __user__={"id": "u1"}, __request__=object(), - __chat_id__="chat-pending", + __chat_id__="chat-failed-transition", ) ) @@ -165,20 +165,87 @@ async def no_draft(*args, **kwargs): return None, None monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + rejected = asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "prohibit docker"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__="chat-failed-transition", + ) + ) follow_up = asyncio.run( pipe.pipe( {"model": "pipe-model", "messages": [{"role": "user", "content": "yes"}]}, __user__={"id": "u1"}, __request__=object(), - __chat_id__="chat-pending", + __chat_id__="chat-failed-transition", ) ) - assert update == "State updated: Use docker." + assert seed == "State updated: Use docker." + assert rejected == ( + '"docker" is currently in use.\nRemove or replace it before prohibiting it.' + ) assert follow_up == {"choices": [{"message": {"content": "downstream"}}]} assert len(forwarded) == 1 +def test_failed_transition_does_not_create_resumable_engine_state(monkeypatch) -> None: + module = _load_module("owui_with_drafter_failed_transition_no_resume", monkeypatch) + pipe = module.Pipe() + pipe.valves.BASE_MODEL_ID = "base-model" + pipe.valves.PREPROCESSOR_MODEL_ID = "prep-model" + + async def update_draft(*args, **kwargs): + return "use docker", None + + async def no_draft(*args, **kwargs): + return None, None + + monkeypatch.setattr(module.Pipe, "_preprocess_user_input", update_draft) + asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "please use docker"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__="chat-no-resume-state", + ) + ) + + monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "prohibit docker"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__="chat-no-resume-state", + ) + ) + + show_state = asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "show state"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__="chat-no-resume-state", + ) + ) + + assert show_state == "Premise: none\nUse: docker\nProhibit: none" + + def test_fallback_to_raw_input_path_preserves_host_behavior(monkeypatch) -> None: module = _load_module("owui_with_drafter_raw", monkeypatch) forwarded: list[dict[str, object]] = [] @@ -213,7 +280,9 @@ async def no_draft(*args, **kwargs): assert forwarded[0]["messages"] == [{"role": "user", "content": "hello"}] -def test_local_update_and_clarify_responses_skip_downstream_model(monkeypatch) -> None: +def test_local_update_and_rejection_responses_skip_downstream_model( + monkeypatch, +) -> None: module = _load_module("owui_with_drafter_local", monkeypatch) forwarded: list[dict[str, object]] = [] @@ -248,7 +317,7 @@ async def no_draft(*args, **kwargs): return None, None monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) - clarify = asyncio.run( + rejection = asyncio.run( pipe.pipe( { "model": "pipe-model", @@ -263,10 +332,70 @@ async def no_draft(*args, **kwargs): ) assert update == "State updated: Use docker." - assert clarify == "Invalid premise syntax.\nUse 'set premise '." + assert rejection == "Invalid premise syntax.\nUse 'set premise '." assert forwarded == [] +def test_near_miss_rejection_does_not_create_pending_state(monkeypatch) -> None: + module = _load_module("owui_with_drafter_near_miss_no_pending", monkeypatch) + forwarded: list[dict[str, object]] = [] + + async def forward( + _: object, payload: dict[str, object], __: object + ) -> dict[str, object]: + forwarded.append(payload) + return {"choices": [{"message": {"content": "downstream"}}]} + + module.generate_chat_completion = forward + pipe = module.Pipe() + pipe.valves.BASE_MODEL_ID = "base-model" + pipe.valves.PREPROCESSOR_MODEL_ID = "prep-model" + + async def no_draft(*args, **kwargs): + return None, None + + monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + chat_id = "chat-near-miss-followup" + + rejected = asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [ + {"role": "user", "content": "set premise to concise replies"} + ], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + follow_up = asyncio.run( + pipe.pipe( + {"model": "pipe-model", "messages": [{"role": "user", "content": "yes"}]}, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + show_state = asyncio.run( + pipe.pipe( + { + "model": "pipe-model", + "messages": [{"role": "user", "content": "show state"}], + }, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__=chat_id, + ) + ) + + assert rejected == "Invalid premise syntax.\nUse 'set premise '." + assert follow_up == {"choices": [{"message": {"content": "downstream"}}]} + assert show_state == "Premise: none\nUse: none\nProhibit: none" + assert len(forwarded) == 1 + + def test_compound_directives_fall_through_to_normal_forwarding(monkeypatch) -> None: module = _load_module("owui_with_drafter_compound", monkeypatch) forwarded: list[dict[str, object]] = [] From deb26a2b692ad35f0e24a19a383c986c59139fe8 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 00:08:32 -0400 Subject: [PATCH 3/4] fix: remove stale OpenWebUI continuation references --- .../openwebui_pipe/README.md | 4 +- .../openwebui_pipe/open_webui_pipe.py | 32 --------------- .../open_webui_pipe_with_directive_drafter.py | 39 ------------------- python/tests/test_openwebui_pipe.py | 10 ++--- ...t_openwebui_pipe_with_directive_drafter.py | 18 +++++---- 5 files changed, 17 insertions(+), 86 deletions(-) diff --git a/python/reference_integrations/openwebui_pipe/README.md b/python/reference_integrations/openwebui_pipe/README.md index cb98f60..d535df4 100644 --- a/python/reference_integrations/openwebui_pipe/README.md +++ b/python/reference_integrations/openwebui_pipe/README.md @@ -133,7 +133,7 @@ Use this pipe when you want the same runtime behavior plus Directive Drafter pre Suggested verification: - Send `use docker` and confirm you get `State updated: Use docker.` with trace showing a local turn -- Send `set premise to concise replies` and confirm Open WebUI clarifies locally with `Use 'set premise '.` +- Send `set premise to concise replies` and confirm Open WebUI rejects the request with `Use 'set premise '.` - Send `please use docker` and confirm either: - the Directive Drafter converts it into a local state update, or - trace shows the turn followed the normal compiler path without a silent state change @@ -205,7 +205,7 @@ If you want a slightly broader manual pass: ### Case 5 - prompt(s): `clear state` → `change premise concise replies` -- base model: generic “please clarify changes” response +- base model: generic “please revise the request” response - basic pipe: `Did you mean 'change premise to concise replies'?` - directive-drafter pipe: same deterministic rejection (near-miss is passed through unchanged) - why this matters: the app waits for explicit, valid directive text before changing state. diff --git a/python/reference_integrations/openwebui_pipe/open_webui_pipe.py b/python/reference_integrations/openwebui_pipe/open_webui_pipe.py index 0a7b741..76ba632 100644 --- a/python/reference_integrations/openwebui_pipe/open_webui_pipe.py +++ b/python/reference_integrations/openwebui_pipe/open_webui_pipe.py @@ -596,38 +596,6 @@ async def _forward_passthrough( return normalized_error return response - async def _forward_update( - self, - body: dict[str, Any], - user_payload: dict[str, Any], - request: Request, - state: _EngineSnapshot, - ) -> Any: - """Forward with one compiler-owned state message based on current state. - - The body is shallow-copied, ``model`` is overridden, and exactly one - compiler-owned message is inserted/replaced before forwarding. - """ - payload = {**body} - payload["model"] = self.valves.BASE_MODEL_ID - - payload["messages"] = _build_forward_messages(body.get("messages"), state=state) - - user = Users.get_user_by_id(user_payload["id"]) - if inspect.isawaitable(user): - user = await user - try: - response = await generate_chat_completion(request, payload, user) - except Exception as exc: - normalized_exception = self._normalize_forward_exception(exc) - if normalized_exception is not None: - return normalized_exception - raise - normalized_error = self._normalize_forward_error(response) - if normalized_error is not None: - return normalized_error - return response - async def pipe( self, body: dict[str, Any], diff --git a/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py b/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py index 9c728a9..2f73a9b 100644 --- a/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py +++ b/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py @@ -822,45 +822,6 @@ async def _forward_passthrough( return normalized_error return response - async def _forward_update( - self, - body: dict[str, Any], - user_payload: dict[str, Any], - request: Request, - state: _EngineSnapshot, - *, - base_model_id: str | None, - ) -> Any: - if base_model_id is None: - if self._allow_missing_base_model_for_debug(): - return ( - "Context Compiler debug mode: BASE_MODEL_ID is empty; " - "skipping model passthrough." - ) - return ( - "Context Compiler pipe misconfigured: BASE_MODEL_ID is required " - "(or set ALLOW_MISSING_BASE_MODEL_FOR_DEBUG=true for testing)." - ) - payload = {**body} - payload["model"] = base_model_id - - payload["messages"] = _build_forward_messages(body.get("messages"), state=state) - - user = Users.get_user_by_id(user_payload["id"]) - if inspect.isawaitable(user): - user = await user - try: - response = await generate_chat_completion(request, payload, user) - except Exception as exc: - normalized_exception = self._normalize_forward_exception(exc) - if normalized_exception is not None: - return normalized_exception - raise - normalized_error = self._normalize_forward_error(response) - if normalized_error is not None: - return normalized_error - return response - async def pipe( self, body: dict[str, Any], diff --git a/python/tests/test_openwebui_pipe.py b/python/tests/test_openwebui_pipe.py index 79b53d5..9c57e36 100644 --- a/python/tests/test_openwebui_pipe.py +++ b/python/tests/test_openwebui_pipe.py @@ -255,13 +255,13 @@ async def _forward( assert len(forwarded) == 1 -def test_failed_transition_does_not_create_resumable_engine_state(monkeypatch) -> None: +def test_failed_transition_does_not_change_existing_engine_state(monkeypatch) -> None: module = _load_module_with_stubs( - "owui_failed_transition_no_resume_state", monkeypatch + "owui_failed_transition_state_preserved", monkeypatch ) pipe = module.Pipe() pipe.valves.BASE_MODEL_ID = "base-model" - chat_id = "chat-no-resume-state" + chat_id = "chat-state-preserved" asyncio.run( pipe.pipe( @@ -377,8 +377,8 @@ async def _forward( assert forwarded == [] -def test_near_miss_rejection_does_not_create_pending_state(monkeypatch) -> None: - module = _load_module_with_stubs("owui_near_miss_no_pending_state", monkeypatch) +def test_near_miss_rejection_does_not_change_existing_engine_state(monkeypatch) -> None: + module = _load_module_with_stubs("owui_near_miss_state_preserved", monkeypatch) forwarded: list[dict[str, object]] = [] async def _forward( diff --git a/python/tests/test_openwebui_pipe_with_directive_drafter.py b/python/tests/test_openwebui_pipe_with_directive_drafter.py index 7464667..53db5cb 100644 --- a/python/tests/test_openwebui_pipe_with_directive_drafter.py +++ b/python/tests/test_openwebui_pipe_with_directive_drafter.py @@ -193,8 +193,10 @@ async def no_draft(*args, **kwargs): assert len(forwarded) == 1 -def test_failed_transition_does_not_create_resumable_engine_state(monkeypatch) -> None: - module = _load_module("owui_with_drafter_failed_transition_no_resume", monkeypatch) +def test_failed_transition_does_not_change_existing_engine_state(monkeypatch) -> None: + module = _load_module( + "owui_with_drafter_failed_transition_state_preserved", monkeypatch + ) pipe = module.Pipe() pipe.valves.BASE_MODEL_ID = "base-model" pipe.valves.PREPROCESSOR_MODEL_ID = "prep-model" @@ -214,7 +216,7 @@ async def no_draft(*args, **kwargs): }, __user__={"id": "u1"}, __request__=object(), - __chat_id__="chat-no-resume-state", + __chat_id__="chat-state-preserved", ) ) @@ -227,7 +229,7 @@ async def no_draft(*args, **kwargs): }, __user__={"id": "u1"}, __request__=object(), - __chat_id__="chat-no-resume-state", + __chat_id__="chat-state-preserved", ) ) @@ -239,7 +241,7 @@ async def no_draft(*args, **kwargs): }, __user__={"id": "u1"}, __request__=object(), - __chat_id__="chat-no-resume-state", + __chat_id__="chat-state-preserved", ) ) @@ -327,7 +329,7 @@ async def no_draft(*args, **kwargs): }, __user__={"id": "u1"}, __request__=object(), - __chat_id__="chat-clarify", + __chat_id__="chat-invalid-request", ) ) @@ -336,8 +338,8 @@ async def no_draft(*args, **kwargs): assert forwarded == [] -def test_near_miss_rejection_does_not_create_pending_state(monkeypatch) -> None: - module = _load_module("owui_with_drafter_near_miss_no_pending", monkeypatch) +def test_near_miss_rejection_does_not_change_existing_engine_state(monkeypatch) -> None: + module = _load_module("owui_with_drafter_near_miss_state_preserved", monkeypatch) forwarded: list[dict[str, object]] = [] async def forward( From d2dd885d8feaa8f08c023fe3a04ab0c9f9890e9a Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 00:13:40 -0400 Subject: [PATCH 4/4] refactor: share LiteLLM hook support helpers --- .../litellm_proxy/_litellm_support.py | 59 +++++++++++++++ .../context_compiler_precall_hook.py | 67 +++-------------- ...ler_precall_hook_with_directive_drafter.py | 74 ++++--------------- 3 files changed, 81 insertions(+), 119 deletions(-) create mode 100644 python/reference_integrations/litellm_proxy/_litellm_support.py diff --git a/python/reference_integrations/litellm_proxy/_litellm_support.py b/python/reference_integrations/litellm_proxy/_litellm_support.py new file mode 100644 index 0000000..013069c --- /dev/null +++ b/python/reference_integrations/litellm_proxy/_litellm_support.py @@ -0,0 +1,59 @@ +"""Shared LiteLLM hook plumbing for request parsing and state rendering.""" + +from __future__ import annotations + +from typing import TypedDict + +from context_compiler import POLICY_PROHIBIT, PolicyValue + + +class EngineSnapshot(TypedDict): + premise: str | None + policies: dict[str, PolicyValue] + + +def snapshot_engine_state(engine: object) -> EngineSnapshot: + premise = getattr(engine, "premise", None) + policies = getattr(engine, "policies", {}) + normalized_policies = ( + dict(policies) + if isinstance(policies, dict) + else dict(policies) + if hasattr(policies, "items") + else {} + ) + return { + "premise": premise if isinstance(premise, str) else None, + "policies": normalized_policies, + } + + +def render_compiled_state_contract(compiled_state: EngineSnapshot) -> str: + prohibited = sorted( + key + for key, value in compiled_state["policies"].items() + if value == POLICY_PROHIBIT + ) + premise = compiled_state["premise"] + + lines: list[str] = ["The following constraints are authoritative."] + if prohibited: + items = ", ".join(prohibited) + lines.append(f"Never recommend or use prohibited items: {items}.") + if premise: + lines.append( + "When the answer depends on user preference/style, " + f"treat the current premise as: {premise}." + ) + lines.append( + "If the user message conflicts with these constraints, follow them exactly." + ) + + return "Host policy contract:\n" + "\n".join(f"- {line}" for line in lines) + + +def extract_request_messages(data: dict[str, object]) -> list[dict[str, object]]: + raw_messages = data.get("messages") + if not isinstance(raw_messages, list): + return [] + return [msg for msg in raw_messages if isinstance(msg, dict)] diff --git a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py index 639012a..f3ed44b 100644 --- a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py +++ b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py @@ -11,7 +11,7 @@ """ import logging -from typing import Any, TypedDict +from typing import Any try: from litellm.integrations.custom_logger import CustomLogger @@ -25,8 +25,6 @@ class CustomLogger: # type: ignore[no-redef] from context_compiler import ( DecisionKind, - POLICY_PROHIBIT, - PolicyValue, create_engine, ) from context_compiler_example_integrations.reference_integrations.litellm_proxy._checkpoint_support import ( @@ -38,6 +36,11 @@ class CustomLogger: # type: ignore[no-redef] extract_latest_user_text, resolve_session_context, ) +from context_compiler_example_integrations.reference_integrations.litellm_proxy._litellm_support import ( + extract_request_messages, + render_compiled_state_contract, + snapshot_engine_state, +) logger = logging.getLogger(__name__) @@ -50,58 +53,6 @@ class CustomLogger: # type: ignore[no-redef] CHECKPOINT_STORE: CheckpointStore = InMemoryCheckpointStore() -class _EngineSnapshot(TypedDict): - premise: str | None - policies: dict[str, PolicyValue] - - -def _snapshot_engine_state(engine: object) -> _EngineSnapshot: - premise = getattr(engine, "premise", None) - policies = getattr(engine, "policies", {}) - normalized_policies = ( - dict(policies) - if isinstance(policies, dict) - else dict(policies) - if hasattr(policies, "items") - else {} - ) - return { - "premise": premise if isinstance(premise, str) else None, - "policies": normalized_policies, - } - - -def _render_compiled_state_contract(compiled_state: _EngineSnapshot) -> str: - prohibited = sorted( - key - for key, value in compiled_state["policies"].items() - if value == POLICY_PROHIBIT - ) - premise = compiled_state["premise"] - - lines: list[str] = ["The following constraints are authoritative."] - if prohibited: - items = ", ".join(prohibited) - lines.append(f"Never recommend or use prohibited items: {items}.") - if premise: - lines.append( - "When the answer depends on user preference/style, " - f"treat the current premise as: {premise}." - ) - lines.append( - "If the user message conflicts with these constraints, follow them exactly." - ) - - return "Host policy contract:\n" + "\n".join(f"- {line}" for line in lines) - - -def _extract_request_messages(data: dict[str, object]) -> list[dict[str, object]]: - raw_messages = data.get("messages") - if not isinstance(raw_messages, list): - return [] - return [msg for msg in raw_messages if isinstance(msg, dict)] - - class ContextCompilerPreCallHook(CustomLogger): async def async_pre_call_hook( self, @@ -115,7 +66,7 @@ async def async_pre_call_hook( if call_type not in _SUPPORTED_CALL_TYPES: return data - request_messages = _extract_request_messages(data) + request_messages = extract_request_messages(data) logger.debug("litellm_proxy: message_count=%d", len(request_messages)) session = resolve_session_context(data) logger.debug( @@ -164,12 +115,12 @@ async def async_pre_call_hook( checkpoint_to_jsonable(engine.export_json()), ) - compiled_state = _snapshot_engine_state(engine) + compiled_state = snapshot_engine_state(engine) # For long-running conversations, you can optionally compact transcripts by removing user inputs that were compiled into state. See Demo 6. # noqa: E501 system_message: dict[str, object] = { "role": "system", "content": "You are a helpful assistant.\n" - + _render_compiled_state_contract(compiled_state), + + render_compiled_state_contract(compiled_state), } # Prepend one compiler contract system message, then forward the original # request messages unchanged. Existing system messages are preserved. diff --git a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py index f521e03..9d03aa6 100644 --- a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py +++ b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py @@ -17,7 +17,7 @@ from importlib import import_module from importlib.resources import as_file, files from importlib.resources.abc import Traversable -from typing import Any, TypedDict, cast +from typing import Any, cast try: from litellm.integrations.custom_logger import CustomLogger @@ -31,8 +31,6 @@ class CustomLogger: # type: ignore[no-redef] from context_compiler import ( DecisionKind, - POLICY_PROHIBIT, - PolicyValue, create_engine, ) from context_compiler_directive_drafter import ( @@ -50,6 +48,12 @@ class CustomLogger: # type: ignore[no-redef] extract_latest_user_text, resolve_session_context, ) +from context_compiler_example_integrations.reference_integrations.litellm_proxy._litellm_support import ( + EngineSnapshot, + extract_request_messages, + render_compiled_state_contract, + snapshot_engine_state, +) logger = logging.getLogger(__name__) @@ -64,58 +68,6 @@ class CustomLogger: # type: ignore[no-redef] CHECKPOINT_STORE: CheckpointStore = InMemoryCheckpointStore() -class _EngineSnapshot(TypedDict): - premise: str | None - policies: dict[str, PolicyValue] - - -def _snapshot_engine_state(engine: object) -> _EngineSnapshot: - premise = getattr(engine, "premise", None) - policies = getattr(engine, "policies", {}) - normalized_policies = ( - dict(policies) - if isinstance(policies, dict) - else dict(policies) - if hasattr(policies, "items") - else {} - ) - return { - "premise": premise if isinstance(premise, str) else None, - "policies": normalized_policies, - } - - -def _render_compiled_state_contract(compiled_state: _EngineSnapshot) -> str: - prohibited = sorted( - key - for key, value in compiled_state["policies"].items() - if value == POLICY_PROHIBIT - ) - premise = compiled_state["premise"] - - lines: list[str] = ["The following constraints are authoritative."] - if prohibited: - items = ", ".join(prohibited) - lines.append(f"Never recommend or use prohibited items: {items}.") - if premise: - lines.append( - "When the answer depends on user preference/style, " - f"treat the current premise as: {premise}." - ) - lines.append( - "If the user message conflicts with these constraints, follow them exactly." - ) - - return "Host policy contract:\n" + "\n".join(f"- {line}" for line in lines) - - -def _extract_request_messages(data: dict[str, object]) -> list[dict[str, object]]: - raw_messages = data.get("messages") - if not isinstance(raw_messages, list): - return [] - return [msg for msg in raw_messages if isinstance(msg, dict)] - - def _extract_response_content(response: object) -> str | None: if isinstance(response, Mapping): choices = response.get("choices") @@ -151,7 +103,7 @@ def _get_litellm_completion() -> Callable[..., object]: return cast(Callable[..., object], litellm_module.completion) -def _llm_fallback_preprocess(message: str, state: _EngineSnapshot) -> str | None: +def _llm_fallback_preprocess(message: str, state: EngineSnapshot) -> str | None: with as_file(_prompt_file_path()) as prompt_path: prompt = render_prompt(prompt_path, state["premise"], state["policies"]) if prompt is None: @@ -198,7 +150,7 @@ def _llm_fallback_preprocess(message: str, state: _EngineSnapshot) -> str | None def _preprocess_last_user_message( - message: str, state: _EngineSnapshot | None + message: str, state: EngineSnapshot | None ) -> str | None: try: heuristic_result = preprocess_heuristic(message) @@ -235,7 +187,7 @@ async def async_pre_call_hook( if call_type not in _SUPPORTED_CALL_TYPES: return data - request_messages = _extract_request_messages(data) + request_messages = extract_request_messages(data) logger.debug("litellm_proxy: message_count=%d", len(request_messages)) session = resolve_session_context(data) logger.debug( @@ -271,7 +223,7 @@ async def async_pre_call_hook( if latest_user_text is not None: drafted_input = _preprocess_last_user_message( - latest_user_text, _snapshot_engine_state(engine) + latest_user_text, snapshot_engine_state(engine) ) logger.debug("litellm_proxy: drafted_input=%r", drafted_input) if drafted_input is not None: @@ -294,11 +246,11 @@ async def async_pre_call_hook( checkpoint_to_jsonable(engine.export_json()), ) - compiled_state = _snapshot_engine_state(engine) + compiled_state = snapshot_engine_state(engine) system_message: dict[str, object] = { "role": "system", "content": "You are a helpful assistant.\n" - + _render_compiled_state_contract(compiled_state), + + render_compiled_state_contract(compiled_state), } logger.debug("litellm_proxy: inject_system_message=true") # Preserve original request messages; drafting changes only compiler input.