From 8675c51ba6f74ba91b74e847323b55054b4812c2 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 00:40:48 -0400 Subject: [PATCH 1/9] refactor: move OpenWebUI drafting to DirectiveDrafter API --- .../open_webui_pipe_with_directive_drafter.py | 126 ++++++++---- ...t_openwebui_pipe_with_directive_drafter.py | 188 ++++++++++++++---- 2 files changed, 233 insertions(+), 81 deletions(-) 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 81df95c..1edca7a 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 @@ -22,6 +22,7 @@ import json import logging import re +import threading from collections.abc import AsyncIterator from typing import Any, Literal, TypedDict, cast @@ -56,11 +57,13 @@ def Field(*, default: Any, description: str = "") -> Any: # type: ignore[no-red PolicyValue, ) from context_compiler.engine import Engine +from context_compiler.grammar import CanonicalDirective from context_compiler_directive_drafter import ( - DRAFT_OUTCOME_DIRECTIVE, + DirectiveDrafter, + DraftResult, + NoDirective, + UnknownDirective, get_converter_prompt, - parse_preprocessor_output, - preprocess_heuristic, ) logger = logging.getLogger(__name__) @@ -428,6 +431,27 @@ def _is_truthy_bool(value: object) -> bool: return False +def _run_coroutine_blocking(awaitable: object) -> Any: + result: dict[str, Any] = {} + error: list[BaseException] = [] + + def _runner() -> None: + import asyncio + + try: + result["value"] = asyncio.run(cast(Any, awaitable)) + except BaseException as exc: # pragma: no cover - exercised via caller tests + error.append(exc) + + thread = threading.Thread(target=_runner, daemon=True) + thread.start() + thread.join() + + if error: + raise error[0] + return result.get("value") + + class Pipe: """Map Context Compiler decisions into Open WebUI pipe behavior. @@ -466,6 +490,7 @@ class Valves(BaseModel): def __init__(self) -> None: self.valves = self.Valves() + self._last_preprocessor_error: str | None = None def _allow_missing_base_model_for_debug(self) -> bool: return _is_truthy_bool( @@ -686,20 +711,18 @@ async def _validate_configured_model_ids( ) return None - async def _llm_fallback_preprocess( + async def _llm_fallback_candidate( self, message: str, - state: _EngineSnapshot, *, request: Request, user_payload: dict[str, Any], - prompt_profile: str, model_id: str | None, - ) -> tuple[str | None, str | None]: - del state, prompt_profile + ) -> str | None: + self._last_preprocessor_error = None model_id = _normalize_model_id(model_id) if model_id is None: - return None, None + return None payload: dict[str, Any] = { "model": model_id, @@ -717,58 +740,72 @@ async def _llm_fallback_preprocess( except Exception as exc: normalized_exception = self._normalize_preprocessor_exception(exc) if normalized_exception is not None: - return None, normalized_exception - return None, None + self._last_preprocessor_error = normalized_exception + logger.warning("preprocessor: %s", normalized_exception) + return None normalized_error = self._normalize_preprocessor_error(response) if normalized_error is not None: - return None, normalized_error + self._last_preprocessor_error = normalized_error + logger.warning("preprocessor: %s", normalized_error) + return None - raw_output = _extract_completion_content(response) - parsed = parse_preprocessor_output(raw_output) - if parsed is None: - return None, None - return parsed.text, None + return _extract_completion_content(response) - async def _preprocess_user_input( + async def _draft_user_input( self, message: str, - state: _EngineSnapshot, *, request: Request, user_payload: dict[str, Any], - prompt_profile: str, model_id: str | None, - ) -> tuple[str | None, str | None]: - # Heuristic first for precision, determinism, and low latency. - # If heuristic does not produce a directive, try Open WebUI-native fallback. - heuristic_result = preprocess_heuristic(message) - - if ( - heuristic_result["outcome"] == DRAFT_OUTCOME_DIRECTIVE - and heuristic_result["directive"] - ): - parsed = parse_preprocessor_output(heuristic_result["directive"]) - if parsed is not None: - return parsed.text, None + ) -> DraftResult: + def fallback(candidate_message: str) -> str | None: + return cast( + str | None, + _run_coroutine_blocking( + self._llm_fallback_candidate( + candidate_message, + request=request, + user_payload=user_payload, + model_id=model_id, + ) + ), + ) - if _is_directive_shaped_input(message): - return None, None + drafter = DirectiveDrafter( + fallback=fallback, + fallback_source="openwebui_fallback", + ) + return drafter.draft_directive(message) - # In debug mode with missing base/preprocessor model ids, skip fallback - # preprocess entirely so we never attempt an empty-model LLM call. - model_id = _normalize_model_id(model_id) - if model_id is None: - return None, None + def _extract_drafted_text(self, drafted_result: DraftResult) -> str | None: + if isinstance(drafted_result.result, CanonicalDirective): + return drafted_result.result.text + if isinstance(drafted_result.result, NoDirective): + return None + if isinstance(drafted_result.result, UnknownDirective): + return None + return None - return await self._llm_fallback_preprocess( + async def _preprocess_user_input( + self, + message: str, + *, + request: Request, + user_payload: dict[str, Any], + prompt_profile: str, + model_id: str | None, + ) -> tuple[DraftResult, str | None]: + del prompt_profile + self._last_preprocessor_error = None + drafted_result = await self._draft_user_input( message, - state, request=request, user_payload=user_payload, - prompt_profile=prompt_profile, model_id=model_id, ) + return drafted_result, self._last_preprocessor_error async def _forward_passthrough( self, @@ -885,9 +922,8 @@ async def pipe( preprocessd: str | None = None preprocess_error: str | None = None - preprocessd, preprocess_error = await self._preprocess_user_input( + drafted_result, preprocess_error = await self._preprocess_user_input( latest_user_text, - _snapshot_engine_state(engine), request=__request__, user_payload=__user__, prompt_profile=self.valves.PREPROCESSOR_PROMPT_PROFILE, @@ -896,6 +932,8 @@ async def pipe( if preprocess_error is not None: return preprocess_error + preprocessd = self._extract_drafted_text(drafted_result) + logger.debug("preprocessor: drafted_result=%r", drafted_result) logger.debug("preprocessor: preprocessd=%r", preprocessd) # Preserve core behavior: if preprocess yields no directive, use raw user # text so the compiler still decides rejection/passthrough/update. diff --git a/python/tests/test_openwebui_pipe_with_directive_drafter.py b/python/tests/test_openwebui_pipe_with_directive_drafter.py index 53db5cb..84d2f85 100644 --- a/python/tests/test_openwebui_pipe_with_directive_drafter.py +++ b/python/tests/test_openwebui_pipe_with_directive_drafter.py @@ -4,8 +4,11 @@ import sys import types from pathlib import Path +from types import MappingProxyType import pytest +from context_compiler.grammar import CanonicalDirective, DirectiveKind +from context_compiler_directive_drafter import DraftResult, NoDirective, UnknownDirective REPO_ROOT = Path(__file__).resolve().parents[2] MODULE_PATH = ( @@ -103,10 +106,17 @@ def tracked_step(user_input: str): monkeypatch.setattr(module, "create_engine", create_engine_with_tracking) - async def fake_preprocess(*args, **kwargs): - return "use docker", None + async def fake_draft(*args, **kwargs): + return DraftResult( + source="test", + result=CanonicalDirective( + text="use docker", + kind=DirectiveKind.USE_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", fake_preprocess) + monkeypatch.setattr(module.Pipe, "_draft_user_input", fake_draft) pipe = module.Pipe() pipe.valves.BASE_MODEL_ID = "base-model" @@ -146,9 +156,16 @@ async def forward( pipe.valves.PREPROCESSOR_MODEL_ID = "prep-model" async def update_draft(*args, **kwargs): - return "use docker", None + return DraftResult( + source="test", + result=CanonicalDirective( + text="use docker", + kind=DirectiveKind.USE_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", update_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", update_draft) seed = asyncio.run( pipe.pipe( { @@ -162,9 +179,12 @@ async def update_draft(*args, **kwargs): ) async def no_draft(*args, **kwargs): - return None, None + return DraftResult( + source="test", + result=NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) rejected = asyncio.run( pipe.pipe( { @@ -202,12 +222,22 @@ def test_failed_transition_does_not_change_existing_engine_state(monkeypatch) -> pipe.valves.PREPROCESSOR_MODEL_ID = "prep-model" async def update_draft(*args, **kwargs): - return "use docker", None + return DraftResult( + source="test", + result=CanonicalDirective( + text="use docker", + kind=DirectiveKind.USE_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ) async def no_draft(*args, **kwargs): - return None, None + return DraftResult( + source="test", + result=NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", update_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", update_draft) asyncio.run( pipe.pipe( { @@ -220,7 +250,7 @@ async def no_draft(*args, **kwargs): ) ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) asyncio.run( pipe.pipe( { @@ -261,9 +291,12 @@ async def forward( module.generate_chat_completion = forward async def no_draft(*args, **kwargs): - return None, None + return DraftResult( + source="test", + result=NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) pipe = module.Pipe() pipe.valves.BASE_MODEL_ID = "base-model" @@ -300,9 +333,16 @@ async def forward( pipe.valves.PREPROCESSOR_MODEL_ID = "prep-model" async def update_draft(*args, **kwargs): - return "use docker", None + return DraftResult( + source="test", + result=CanonicalDirective( + text="use docker", + kind=DirectiveKind.USE_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", update_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", update_draft) update = asyncio.run( pipe.pipe( { @@ -316,9 +356,12 @@ async def update_draft(*args, **kwargs): ) async def no_draft(*args, **kwargs): - return None, None + return DraftResult( + source="test", + result=NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) rejection = asyncio.run( pipe.pipe( { @@ -354,9 +397,12 @@ async def forward( pipe.valves.PREPROCESSOR_MODEL_ID = "prep-model" async def no_draft(*args, **kwargs): - return None, None + return DraftResult( + source="test", + result=NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) chat_id = "chat-near-miss-followup" rejected = asyncio.run( @@ -414,9 +460,12 @@ async def forward( pipe.valves.PREPROCESSOR_MODEL_ID = "prep-model" async def compound_draft(*args, **kwargs): - return "use docker and prohibit peanuts", None + return DraftResult( + source="test", + result=UnknownDirective(reason="reject.multi_candidate_directive"), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", compound_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", compound_draft) result = asyncio.run( pipe.pipe( { @@ -457,9 +506,16 @@ async def forward( chat_id = "chat-passthrough" async def update_draft(*args, **kwargs): - return "use docker", None + return DraftResult( + source="test", + result=CanonicalDirective( + text="use docker", + kind=DirectiveKind.USE_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", update_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", update_draft) asyncio.run( pipe.pipe( { @@ -473,9 +529,12 @@ async def update_draft(*args, **kwargs): ) async def no_draft(*args, **kwargs): - return None, None + return DraftResult( + source="test", + result=NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) result = asyncio.run( pipe.pipe( { @@ -575,9 +634,12 @@ def test_debug_mode_missing_base_model_returns_deterministic_message( pipe.valves.ALLOW_MISSING_BASE_MODEL_FOR_DEBUG = True async def no_draft(*args, **kwargs): - return None, None + return DraftResult( + source="test", + result=NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) result = asyncio.run( pipe.pipe( @@ -608,11 +670,6 @@ async def generate( return {"choices": [{"message": {"content": "downstream"}}]} module.generate_chat_completion = generate - module.preprocess_heuristic = lambda _text: { - "outcome": "no_directive", - "directive": None, - } - result = asyncio.run( pipe.pipe( {"model": "pipe-model", "messages": [{"role": "user", "content": "hello"}]}, @@ -647,11 +704,6 @@ async def generate( return {"choices": [{"message": {"content": "downstream"}}]} module.generate_chat_completion = generate - module.preprocess_heuristic = lambda _text: { - "outcome": "no_directive", - "directive": None, - } - result = asyncio.run( pipe.pipe( { @@ -668,6 +720,68 @@ async def generate( assert calls == ["prep-model", "base-model"] +def test_extract_drafted_text_only_applies_canonical_directive(monkeypatch) -> None: + module = _load_module("owui_with_drafter_extract_text", monkeypatch) + pipe = module.Pipe() + + canonical = DraftResult( + source="test", + result=CanonicalDirective( + text="use docker", + kind=DirectiveKind.USE_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ) + no_directive = DraftResult( + source="test", + result=NoDirective(reason="reject.confident_non_directive"), + ) + unknown = DraftResult( + source="test", + result=UnknownDirective(reason="reject.multi_candidate_directive"), + ) + + assert pipe._extract_drafted_text(canonical) == "use docker" + assert pipe._extract_drafted_text(no_directive) is None + assert pipe._extract_drafted_text(unknown) is None + + +def test_unknown_directive_falls_back_to_normal_user_input_flow(monkeypatch) -> None: + module = _load_module("owui_with_drafter_unknown_falls_back", 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 unknown_draft(*args, **kwargs): + return DraftResult( + source="test", + result=UnknownDirective(reason="reject.multi_candidate_directive"), + ) + + monkeypatch.setattr(module.Pipe, "_draft_user_input", unknown_draft) + + result = asyncio.run( + pipe.pipe( + {"model": "pipe-model", "messages": [{"role": "user", "content": "hello"}]}, + __user__={"id": "u1"}, + __request__=object(), + __chat_id__="chat-unknown-raw", + ) + ) + + assert result == {"choices": [{"message": {"content": "downstream"}}]} + assert forwarded[0]["messages"] == [{"role": "user", "content": "hello"}] + + def test_validate_configured_model_ids_supports_async_user_lookup(monkeypatch) -> None: module = _load_module("owui_with_drafter_async_user_lookup", monkeypatch) pipe = module.Pipe() From e57f44e1bf8d15106bc0be9d0da13b70007838ec Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 00:42:59 -0400 Subject: [PATCH 2/9] fix: migrate openwebui directive drafter boundary --- .../open_webui_pipe_with_directive_drafter.py | 49 ++++++++----------- 1 file changed, 20 insertions(+), 29 deletions(-) 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 1edca7a..dd08951 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 @@ -126,21 +126,20 @@ def _restore_engine_from_snapshot(snapshot_json: str) -> Engine: return engine -def _render_compiler_state_block(state: _EngineSnapshot) -> str: +def _render_compiler_state_block(engine: Engine) -> str: lines: list[str] = [_CC_MARKER] - premise = state["premise"] - if premise is not None: - lines.append(f"Premise: {premise}") + if engine.premise is not None: + lines.append(f"Premise: {engine.premise}") use_items = sorted( - key for key, value in state["policies"].items() if value == POLICY_USE + key for key, value in engine.policies.items() if value == POLICY_USE ) if use_items: lines.append("Use: " + ", ".join(use_items)) prohibit_items = sorted( - key for key, value in state["policies"].items() if value == POLICY_PROHIBIT + key for key, value in engine.policies.items() if value == POLICY_PROHIBIT ) if prohibit_items: lines.append("Prohibit: " + ", ".join(prohibit_items)) @@ -149,18 +148,16 @@ def _render_compiler_state_block(state: _EngineSnapshot) -> str: def _render_show_state_summary(engine: Engine) -> str: - snapshot = _snapshot_engine_state(engine) - premise = snapshot["premise"] use_items = sorted( - key for key, value in snapshot["policies"].items() if value == POLICY_USE + key for key, value in engine.policies.items() if value == POLICY_USE ) prohibit_items = sorted( - key for key, value in snapshot["policies"].items() if value == POLICY_PROHIBIT + key for key, value in engine.policies.items() if value == POLICY_PROHIBIT ) use_text = ", ".join(use_items) if use_items else "none" prohibit_text = ", ".join(prohibit_items) if prohibit_items else "none" - premise_text = premise if premise is not None else "none" + premise_text = engine.premise if engine.premise is not None else "none" return f"Premise: {premise_text}\nUse: {use_text}\nProhibit: {prohibit_text}" @@ -215,10 +212,10 @@ def _normalize_state(value: object) -> _EngineSnapshot: } -def _has_non_empty_authoritative_state(state: _EngineSnapshot) -> bool: - if state["premise"] is not None: +def _has_non_empty_authoritative_state(engine: Engine) -> bool: + if engine.premise is not None: return True - return bool(state["policies"]) + return bool(engine.policies) def _render_state_summary_line(state: object) -> str: @@ -290,7 +287,7 @@ def _strip_trace_blocks_from_messages( def _build_forward_messages( raw_messages: object, *, - state: _EngineSnapshot | None = None, + engine: Engine | None = None, ) -> list[dict[str, Any]]: """Build forwarded messages with trace stripping and optional state injection.""" messages = ( @@ -300,10 +297,10 @@ def _build_forward_messages( if isinstance(raw_messages, list) else [] ) - if state is not None and _has_non_empty_authoritative_state(state): + if engine is not None and _has_non_empty_authoritative_state(engine): return _replace_compiler_system_message( messages, - _render_compiler_state_block(state), + _render_compiler_state_block(engine), ) return messages @@ -814,7 +811,7 @@ async def _forward_passthrough( request: Request, *, base_model_id: str | None, - state: _EngineSnapshot | None = None, + engine: Engine | None = None, ) -> Any: if base_model_id is None: if self._allow_missing_base_model_for_debug(): @@ -828,7 +825,7 @@ async def _forward_passthrough( ) payload = {**body} payload["model"] = base_model_id - payload["messages"] = _build_forward_messages(body.get("messages"), state=state) + payload["messages"] = _build_forward_messages(body.get("messages"), engine=engine) user = Users.get_user_by_id(user_payload["id"]) if inspect.isawaitable(user): user = await user @@ -986,16 +983,13 @@ async def pipe( llm_called=False, ) if decision["kind"] == DecisionKind.NO_DIRECTIVE: - compiled_state = _normalize_state(state_after) - state_injected = ( - "yes" if _has_non_empty_authoritative_state(compiled_state) else "no" - ) + state_injected = "yes" if _has_non_empty_authoritative_state(engine) else "no" response = await self._forward_passthrough( body, __user__, __request__, base_model_id=base_model_id, - state=compiled_state, + engine=engine, ) return self._with_trace( response, @@ -1020,16 +1014,13 @@ async def pipe( llm_called=False, ) - compiled_state = _normalize_state(state_after) - state_injected = ( - "yes" if _has_non_empty_authoritative_state(compiled_state) else "no" - ) + state_injected = "yes" if _has_non_empty_authoritative_state(engine) else "no" response = await self._forward_passthrough( body, __user__, __request__, base_model_id=base_model_id, - state=compiled_state, + engine=engine, ) return self._with_trace( response, From 755ad5054d3106ea764ab52c6f55bef08b76001c Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 00:47:48 -0400 Subject: [PATCH 3/9] refactor: remove openwebui directive parsing leftovers --- .../open_webui_pipe_with_directive_drafter.py | 105 +----------------- ...t_openwebui_pipe_with_directive_drafter.py | 24 ++-- 2 files changed, 14 insertions(+), 115 deletions(-) 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 dd08951..a87ec22 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 @@ -77,19 +77,6 @@ class _EngineSnapshot(TypedDict): policies: dict[str, PolicyValue] -def _is_directive_shaped_input(message: str) -> bool: - normalized = re.sub(r"\s+", " ", message.strip()).lower() - return ( - normalized.startswith("use") - or normalized.startswith("prohibit") - or normalized.startswith("remove policy") - or normalized.startswith("set premise") - or normalized.startswith("change premise") - or normalized.startswith("clear") - or normalized.startswith("reset") - ) - - def _resolve_chat_key( user: dict[str, Any], chat_id: str | None, @@ -319,73 +306,6 @@ def _render_item_label(value: str) -> str: return re.sub(r"\s+", " ", value).strip().lower() -def _near_miss_directive_rejection(value: str) -> str | None: - normalized = re.sub(r"\s+", " ", value.strip()) - lower = normalized.lower() - - if lower in {"reset premise", "reset premises", "clear premises"}: - return "Unknown directive.\nUse 'clear premise' or 'reset policies'." - if lower.startswith("set premise to "): - return "Invalid premise syntax.\nUse 'set premise '." - if lower.startswith("change premise ") and not lower.startswith( - "change premise to " - ): - return "Invalid premise syntax.\nUse 'change premise to '." - return None - - -def _summarize_update_from_input(user_input: str) -> str: - normalized = re.sub(r"\s+", " ", user_input.strip()) - lower = normalized.lower() - - if lower == "clear state": - return "State cleared." - if lower == "clear premise": - return "Premise cleared." - if lower == "reset policies": - return "Policies reset." - - replacement_match = re.match( - r"^use\s+(.+?)\s+instead\s+of\s+(.+)$", normalized, flags=re.IGNORECASE - ) - if replacement_match is not None: - item = _render_item_label(replacement_match.group(1).rstrip(" .!?")) - if item: - return f"State updated: Use {item}." - - use_match = re.match(r"^use\s+(.+)$", normalized, flags=re.IGNORECASE) - if use_match is not None: - item = _render_item_label(use_match.group(1).rstrip(" .!?")) - if item: - return f"State updated: Use {item}." - - prohibit_match = re.match(r"^prohibit\s+(.+)$", normalized, flags=re.IGNORECASE) - if prohibit_match is not None: - item = _render_item_label(prohibit_match.group(1).rstrip(" .!?")) - if item: - return f"State updated: Prohibit {item}." - - remove_policy_match = re.match( - r"^remove\s+policy\s+(.+)$", normalized, flags=re.IGNORECASE - ) - if remove_policy_match is not None: - item = _render_item_label(remove_policy_match.group(1).rstrip(" .!?")) - if item: - return f"State updated: Removed policy {item}." - - return "State updated." - - -def _is_administrative_update_input(user_input: str) -> bool: - normalized = re.sub(r"\s+", " ", user_input.strip()).lower() - return ( - normalized == "clear state" - or normalized == "clear premise" - or normalized == "reset policies" - or normalized.startswith("remove policy ") - ) - - def _extract_completion_content(response: object) -> str | None: choices_attr = getattr(response, "choices", None) if isinstance(choices_attr, list) and choices_attr: @@ -945,7 +865,6 @@ async def pipe( else: kind = DecisionKind.NO_DIRECTIVE.value logger.debug("preprocessor: decision=%s", kind) - near_miss_prompt = _near_miss_directive_rejection(latest_user_text) state_after = _snapshot_engine_state(engine) if decision["kind"] == DecisionKind.ERROR: @@ -953,7 +872,7 @@ async def pipe( engine_snapshot_json ) return self._with_trace( - near_miss_prompt or decision["message"] or "", + decision["message"] or "", original_input=latest_user_text, compiler_input=compile_input, decision=decision, @@ -962,26 +881,6 @@ async def pipe( preprocessor_output=preprocessd, llm_called=False, ) - if ( - 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, - compiler_input=compile_input, - decision={ - "kind": DecisionKind.ERROR.value, - "message": near_miss_prompt, - }, - state_before=state_before, - state_after=state_after, - preprocessor_output=preprocessd, - llm_called=False, - ) if decision["kind"] == DecisionKind.NO_DIRECTIVE: state_injected = "yes" if _has_non_empty_authoritative_state(engine) else "no" response = await self._forward_passthrough( @@ -1004,7 +903,7 @@ async def pipe( ) if is_update(decision): return self._with_trace( - _summarize_update_from_input(compile_input), + "State updated.", original_input=latest_user_text, compiler_input=compile_input, decision=decision, diff --git a/python/tests/test_openwebui_pipe_with_directive_drafter.py b/python/tests/test_openwebui_pipe_with_directive_drafter.py index 84d2f85..e5ce903 100644 --- a/python/tests/test_openwebui_pipe_with_directive_drafter.py +++ b/python/tests/test_openwebui_pipe_with_directive_drafter.py @@ -134,7 +134,7 @@ async def fake_draft(*args, **kwargs): ) ) - assert result == "State updated: Use docker." + assert result == "State updated." assert compile_inputs == ["use docker"] @@ -205,7 +205,7 @@ async def no_draft(*args, **kwargs): ) ) - assert seed == "State updated: Use docker." + assert seed == "State updated." assert rejected == ( '"docker" is currently in use.\nRemove or replace it before prohibiting it.' ) @@ -315,7 +315,7 @@ async def no_draft(*args, **kwargs): assert forwarded[0]["messages"] == [{"role": "user", "content": "hello"}] -def test_local_update_and_rejection_responses_skip_downstream_model( +def test_local_update_and_no_directive_passthrough_preserve_host_behavior( monkeypatch, ) -> None: module = _load_module("owui_with_drafter_local", monkeypatch) @@ -362,7 +362,7 @@ async def no_draft(*args, **kwargs): ) monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) - rejection = asyncio.run( + passthrough = asyncio.run( pipe.pipe( { "model": "pipe-model", @@ -376,13 +376,13 @@ async def no_draft(*args, **kwargs): ) ) - assert update == "State updated: Use docker." - assert rejection == "Invalid premise syntax.\nUse 'set premise '." - assert forwarded == [] + assert update == "State updated." + assert passthrough == {"choices": [{"message": {"content": "downstream"}}]} + assert len(forwarded) == 1 -def test_near_miss_rejection_does_not_change_existing_engine_state(monkeypatch) -> None: - module = _load_module("owui_with_drafter_near_miss_state_preserved", monkeypatch) +def test_no_directive_passthrough_does_not_change_existing_engine_state(monkeypatch) -> None: + module = _load_module("owui_with_drafter_no_directive_state_preserved", monkeypatch) forwarded: list[dict[str, object]] = [] async def forward( @@ -405,7 +405,7 @@ async def no_draft(*args, **kwargs): monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) chat_id = "chat-near-miss-followup" - rejected = asyncio.run( + passthrough = asyncio.run( pipe.pipe( { "model": "pipe-model", @@ -438,10 +438,10 @@ async def no_draft(*args, **kwargs): ) ) - assert rejected == "Invalid premise syntax.\nUse 'set premise '." + assert passthrough == {"choices": [{"message": {"content": "downstream"}}]} assert follow_up == {"choices": [{"message": {"content": "downstream"}}]} assert show_state == "Premise: none\nUse: none\nProhibit: none" - assert len(forwarded) == 1 + assert len(forwarded) == 2 def test_compound_directives_fall_through_to_normal_forwarding(monkeypatch) -> None: From 822a15158fc407152ad4bc2f42522c9fc58c37dd Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 00:50:02 -0400 Subject: [PATCH 4/9] fix: use async drafter in openwebui pipe --- .../open_webui_pipe_with_directive_drafter.py | 45 ++++--------------- 1 file changed, 9 insertions(+), 36 deletions(-) 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 a87ec22..07a8d45 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 @@ -22,7 +22,6 @@ import json import logging import re -import threading from collections.abc import AsyncIterator from typing import Any, Literal, TypedDict, cast @@ -348,27 +347,6 @@ def _is_truthy_bool(value: object) -> bool: return False -def _run_coroutine_blocking(awaitable: object) -> Any: - result: dict[str, Any] = {} - error: list[BaseException] = [] - - def _runner() -> None: - import asyncio - - try: - result["value"] = asyncio.run(cast(Any, awaitable)) - except BaseException as exc: # pragma: no cover - exercised via caller tests - error.append(exc) - - thread = threading.Thread(target=_runner, daemon=True) - thread.start() - thread.join() - - if error: - raise error[0] - return result.get("value") - - class Pipe: """Map Context Compiler decisions into Open WebUI pipe behavior. @@ -677,24 +655,19 @@ async def _draft_user_input( user_payload: dict[str, Any], model_id: str | None, ) -> DraftResult: - def fallback(candidate_message: str) -> str | None: - return cast( - str | None, - _run_coroutine_blocking( - self._llm_fallback_candidate( - candidate_message, - request=request, - user_payload=user_payload, - model_id=model_id, - ) - ), + async def fallback(candidate_message: str) -> str | None: + return await self._llm_fallback_candidate( + candidate_message, + request=request, + user_payload=user_payload, + model_id=model_id, ) drafter = DirectiveDrafter( - fallback=fallback, - fallback_source="openwebui_fallback", + async_fallback=fallback, + async_fallback_source="openwebui_fallback", ) - return drafter.draft_directive(message) + return await drafter.async_draft_directive(message) def _extract_drafted_text(self, drafted_result: DraftResult) -> str | None: if isinstance(drafted_result.result, CanonicalDirective): From 3ffd7773f02f77d4aabf8f06f679d47b8ca93829 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 01:10:03 -0400 Subject: [PATCH 5/9] fix: restrict drafter compiler input to canonical directives --- .../litellm/with_directive_drafter.py | 22 ++++++++--- ...ler_precall_hook_with_directive_drafter.py | 9 ++--- .../open_webui_pipe_with_directive_drafter.py | 38 +++++++++++++------ ...st_litellm_proxy_with_directive_drafter.py | 16 ++++++-- .../test_litellm_with_directive_drafter.py | 2 +- ...t_openwebui_pipe_with_directive_drafter.py | 13 ++++++- 6 files changed, 71 insertions(+), 29 deletions(-) diff --git a/python/examples/prompt_construction/litellm/with_directive_drafter.py b/python/examples/prompt_construction/litellm/with_directive_drafter.py index 9ff34fd..8d547e1 100644 --- a/python/examples/prompt_construction/litellm/with_directive_drafter.py +++ b/python/examples/prompt_construction/litellm/with_directive_drafter.py @@ -295,13 +295,23 @@ def _append_trace( def handle_turn(user_input: str, engine: Engine) -> str: state_before = (engine.premise, dict(engine.policies)) - preprocessd: str | None = None preprocessd = _preprocess_user_input(user_input) - compile_input = preprocessd if preprocessd else user_input - logger.debug( - "preprocessor: engine_input=%s", - "directive" if preprocessd else f"user_input len={len(user_input)}", - ) + if preprocessd is None: + messages = _build_messages(user_input, engine) + response_text = _call_litellm(messages) + return _append_trace( + response_text, + original_input=user_input, + compiler_input=user_input, + preprocessor_output=None, + decision={"kind": DecisionKind.NO_DIRECTIVE.value, "message": None}, + state_before=state_before, + state_after=(engine.premise, dict(engine.policies)), + llm_called=True, + ) + + compile_input = preprocessd + logger.debug("preprocessor: engine_input=directive") decision = engine.step(compile_input) if decision["kind"] == DecisionKind.ERROR: 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 1c3736e..7243902 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 @@ -179,17 +179,16 @@ async def async_pre_call_hook( logger.debug( "litellm_proxy: latest_user_text_present=%s", latest_user_text is not None ) - engine_input = latest_user_text drafted_result: DraftResult | None = None + decision: dict[str, object] if latest_user_text is not None: drafted_result = _draft_last_user_message(latest_user_text) logger.debug("litellm_proxy: drafted_result=%r", drafted_result) if isinstance(drafted_result.result, CanonicalDirective): - engine_input = drafted_result.result.text - - if engine_input is not None: - decision = engine.step(engine_input) + decision = engine.step(drafted_result.result.text) + else: + decision = {"kind": DecisionKind.NO_DIRECTIVE, "message": None} else: decision = {"kind": DecisionKind.NO_DIRECTIVE, "message": None} 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 07a8d45..501692e 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 @@ -808,9 +808,6 @@ 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 drafted_result, preprocess_error = await self._preprocess_user_input( latest_user_text, @@ -822,13 +819,30 @@ async def pipe( if preprocess_error is not None: return preprocess_error - preprocessd = self._extract_drafted_text(drafted_result) logger.debug("preprocessor: drafted_result=%r", drafted_result) - logger.debug("preprocessor: preprocessd=%r", preprocessd) - # Preserve core behavior: if preprocess yields no directive, use raw user - # text so the compiler still decides rejection/passthrough/update. - compile_input = preprocessd if preprocessd is not None else latest_user_text + if not isinstance(drafted_result.result, CanonicalDirective): + state_injected = "yes" if _has_non_empty_authoritative_state(engine) else "no" + response = await self._forward_passthrough( + body, + __user__, + __request__, + base_model_id=base_model_id, + engine=engine, + ) + return self._with_trace( + response, + original_input=latest_user_text, + compiler_input=latest_user_text, + decision={"kind": DecisionKind.NO_DIRECTIVE.value, "message": None}, + state_before=state_before, + state_after=state_before, + preprocessor_output=None, + llm_called=base_model_id is not None, + state_injected=state_injected, + ) + engine_snapshot_json = engine.export_json() + compile_input = drafted_result.result.text logger.debug("preprocessor: engine_input=%r", compile_input) decision = engine.step(compile_input) if decision["kind"] == DecisionKind.ERROR: @@ -851,7 +865,7 @@ async def pipe( decision=decision, state_before=state_before, state_after=state_after, - preprocessor_output=preprocessd, + preprocessor_output=compile_input, llm_called=False, ) if decision["kind"] == DecisionKind.NO_DIRECTIVE: @@ -870,7 +884,7 @@ async def pipe( decision=decision, state_before=state_before, state_after=state_after, - preprocessor_output=preprocessd, + preprocessor_output=compile_input, llm_called=base_model_id is not None, state_injected=state_injected, ) @@ -882,7 +896,7 @@ async def pipe( decision=decision, state_before=state_before, state_after=state_after, - preprocessor_output=preprocessd, + preprocessor_output=compile_input, llm_called=False, ) @@ -901,7 +915,7 @@ async def pipe( decision=decision, state_before=state_before, state_after=state_after, - preprocessor_output=preprocessd, + preprocessor_output=compile_input, llm_called=base_model_id is not None, state_injected=state_injected, ) diff --git a/python/tests/test_litellm_proxy_with_directive_drafter.py b/python/tests/test_litellm_proxy_with_directive_drafter.py index c4f3803..d5909b2 100644 --- a/python/tests/test_litellm_proxy_with_directive_drafter.py +++ b/python/tests/test_litellm_proxy_with_directive_drafter.py @@ -112,7 +112,7 @@ def fake_draft(message: str) -> object: drafted_inputs.append(message) return module.DraftResult( source="test", - result=NoDirective(reason="reject.confident_non_directive"), + result=decompose_directive("change premise to formal tone"), ) monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) @@ -207,13 +207,13 @@ def test_persistent_mode_with_drafter_preserves_existing_checkpoint_on_failure( module.CHECKPOINT_STORE.clear() hook = module.ContextCompilerPreCallHookWithPreprocessor() - def fake_draft(message: str) -> object: + def seed_draft(message: str) -> object: return module.DraftResult( source="test", - result=NoDirective(reason="reject.confident_non_directive"), + result=decompose_directive("use docker"), ) - monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) + monkeypatch.setattr(module, "_draft_last_user_message", seed_draft) seed_data = { "model": "demo", "context_compiler_mode": "persistent", @@ -232,6 +232,14 @@ def fake_draft(message: str) -> object: ) assert seed_result is seed_data + def reject_draft(message: str) -> object: + return module.DraftResult( + source="test", + result=decompose_directive("prohibit docker"), + ) + + monkeypatch.setattr(module, "_draft_last_user_message", reject_draft) + result = asyncio.run( hook.async_pre_call_hook(None, None, rejected_data, "completion") ) diff --git a/python/tests/test_litellm_with_directive_drafter.py b/python/tests/test_litellm_with_directive_drafter.py index 287fc3a..0467fe1 100644 --- a/python/tests/test_litellm_with_directive_drafter.py +++ b/python/tests/test_litellm_with_directive_drafter.py @@ -70,7 +70,7 @@ def downstream(messages: list[dict[str, str]]) -> str: result = module.handle_turn("hello there", engine) - assert compile_inputs == ["hello there"] + assert compile_inputs == [] assert result == "stubbed reply" assert len(llm_calls) == 1 diff --git a/python/tests/test_openwebui_pipe_with_directive_drafter.py b/python/tests/test_openwebui_pipe_with_directive_drafter.py index e5ce903..93a6dd2 100644 --- a/python/tests/test_openwebui_pipe_with_directive_drafter.py +++ b/python/tests/test_openwebui_pipe_with_directive_drafter.py @@ -181,7 +181,11 @@ async def update_draft(*args, **kwargs): async def no_draft(*args, **kwargs): return DraftResult( source="test", - result=NoDirective(reason="reject.confident_non_directive"), + result=CanonicalDirective( + text="prohibit docker", + kind=DirectiveKind.PROHIBIT_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), ) monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) @@ -196,6 +200,13 @@ async def no_draft(*args, **kwargs): __chat_id__="chat-failed-transition", ) ) + async def follow_up_no_directive(*args, **kwargs): + return DraftResult( + source="test", + result=NoDirective(reason="reject.confident_non_directive"), + ) + + monkeypatch.setattr(module.Pipe, "_draft_user_input", follow_up_no_directive) follow_up = asyncio.run( pipe.pipe( {"model": "pipe-model", "messages": [{"role": "user", "content": "yes"}]}, From dbb67d0ac0bc4824b626b68c2027c9ae15eabf47 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 01:14:12 -0400 Subject: [PATCH 6/9] refactor: use engine state directly in integrations --- .../litellm_proxy/_litellm_support.py | 34 +++---------- .../context_compiler_precall_hook.py | 4 +- ...ler_precall_hook_with_directive_drafter.py | 4 +- .../openwebui_pipe/open_webui_pipe.py | 49 ++++++++----------- 4 files changed, 28 insertions(+), 63 deletions(-) diff --git a/python/reference_integrations/litellm_proxy/_litellm_support.py b/python/reference_integrations/litellm_proxy/_litellm_support.py index 013069c..ca80eea 100644 --- a/python/reference_integrations/litellm_proxy/_litellm_support.py +++ b/python/reference_integrations/litellm_proxy/_litellm_support.py @@ -1,49 +1,27 @@ -"""Shared LiteLLM hook plumbing for request parsing and state rendering.""" +"""Shared LiteLLM hook plumbing for request parsing and engine-state rendering.""" from __future__ import annotations -from typing import TypedDict +from context_compiler.engine import Engine 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: +def render_compiled_state_contract(engine: Engine) -> str: prohibited = sorted( key - for key, value in compiled_state["policies"].items() + for key, value in engine.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: + if engine.premise: lines.append( "When the answer depends on user preference/style, " - f"treat the current premise as: {premise}." + f"treat the current premise as: {engine.premise}." ) lines.append( "If the user message conflicts with these constraints, follow them exactly." 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 f3ed44b..aed4189 100644 --- a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py +++ b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook.py @@ -39,7 +39,6 @@ class CustomLogger: # type: ignore[no-redef] 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__) @@ -115,12 +114,11 @@ async def async_pre_call_hook( checkpoint_to_jsonable(engine.export_json()), ) - 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(engine), } # 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 7243902..a9b07da 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 @@ -49,7 +49,6 @@ class CustomLogger: # type: ignore[no-redef] 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__) @@ -204,11 +203,10 @@ async def async_pre_call_hook( checkpoint_to_jsonable(engine.export_json()), ) - 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(engine), } logger.debug("litellm_proxy: inject_system_message=true") # Preserve original request messages; drafting changes only compiler input. diff --git a/python/reference_integrations/openwebui_pipe/open_webui_pipe.py b/python/reference_integrations/openwebui_pipe/open_webui_pipe.py index 76ba632..f036f2c 100644 --- a/python/reference_integrations/openwebui_pipe/open_webui_pipe.py +++ b/python/reference_integrations/openwebui_pipe/open_webui_pipe.py @@ -119,7 +119,7 @@ def _restore_engine_from_snapshot(snapshot_json: str) -> Engine: return engine -def _render_compiler_state_block(state: _EngineSnapshot) -> str: +def _render_compiler_state_block(engine: Engine) -> str: """Render deterministic compiler-owned state block text. The first line is ``[[cc_state]]``. Optional lines follow for ``Premise``, @@ -128,18 +128,17 @@ def _render_compiler_state_block(state: _EngineSnapshot) -> str: """ lines: list[str] = [_CC_MARKER] - premise = state["premise"] - if premise is not None: - lines.append(f"Premise: {premise}") + if engine.premise is not None: + lines.append(f"Premise: {engine.premise}") use_items = sorted( - key for key, value in state["policies"].items() if value == POLICY_USE + key for key, value in engine.policies.items() if value == POLICY_USE ) if use_items: lines.append("Use: " + ", ".join(use_items)) prohibit_items = sorted( - key for key, value in state["policies"].items() if value == POLICY_PROHIBIT + key for key, value in engine.policies.items() if value == POLICY_PROHIBIT ) if prohibit_items: lines.append("Prohibit: " + ", ".join(prohibit_items)) @@ -148,18 +147,16 @@ def _render_compiler_state_block(state: _EngineSnapshot) -> str: def _render_show_state_summary(engine: Engine) -> str: - snapshot = _snapshot_engine_state(engine) - premise = snapshot["premise"] use_items = sorted( - key for key, value in snapshot["policies"].items() if value == POLICY_USE + key for key, value in engine.policies.items() if value == POLICY_USE ) prohibit_items = sorted( - key for key, value in snapshot["policies"].items() if value == POLICY_PROHIBIT + key for key, value in engine.policies.items() if value == POLICY_PROHIBIT ) use_text = ", ".join(use_items) if use_items else "none" prohibit_text = ", ".join(prohibit_items) if prohibit_items else "none" - premise_text = premise if premise is not None else "none" + premise_text = engine.premise if engine.premise is not None else "none" return f"Premise: {premise_text}\nUse: {use_text}\nProhibit: {prohibit_text}" @@ -223,10 +220,10 @@ def _normalize_state(value: object) -> _EngineSnapshot: } -def _has_non_empty_authoritative_state(state: _EngineSnapshot) -> bool: - if state["premise"] is not None: +def _has_non_empty_authoritative_state(engine: Engine) -> bool: + if engine.premise is not None: return True - return bool(state["policies"]) + return bool(engine.policies) def _render_state_summary_line(state: object) -> str: @@ -298,7 +295,7 @@ def _strip_trace_blocks_from_messages( def _build_forward_messages( raw_messages: object, *, - state: _EngineSnapshot | None = None, + engine: Engine | None = None, ) -> list[dict[str, Any]]: """Build forwarded messages with trace stripping and optional state injection.""" messages = ( @@ -308,10 +305,10 @@ def _build_forward_messages( if isinstance(raw_messages, list) else [] ) - if state is not None and _has_non_empty_authoritative_state(state): + if engine is not None and _has_non_empty_authoritative_state(engine): return _replace_compiler_system_message( messages, - _render_compiler_state_block(state), + _render_compiler_state_block(engine), ) return messages @@ -575,12 +572,12 @@ async def _forward_passthrough( user_payload: dict[str, Any], request: Request, *, - state: _EngineSnapshot | None = None, + engine: Engine | None = None, ) -> Any: """Forward with model override and optional compiler-owned state injection.""" payload = {**body} payload["model"] = self.valves.BASE_MODEL_ID - payload["messages"] = _build_forward_messages(body.get("messages"), state=state) + payload["messages"] = _build_forward_messages(body.get("messages"), engine=engine) user = Users.get_user_by_id(user_payload["id"]) if inspect.isawaitable(user): user = await user @@ -692,12 +689,9 @@ async def pipe( llm_called=False, ) if decision["kind"] == DecisionKind.NO_DIRECTIVE: - compiled_state = _normalize_state(state_after) - state_injected = ( - "yes" if _has_non_empty_authoritative_state(compiled_state) else "no" - ) + state_injected = "yes" if _has_non_empty_authoritative_state(engine) else "no" response = await self._forward_passthrough( - body, __user__, __request__, state=compiled_state + body, __user__, __request__, engine=engine ) return self._with_trace( response, @@ -720,12 +714,9 @@ async def pipe( llm_called=False, ) - compiled_state = _normalize_state(state_after) - state_injected = ( - "yes" if _has_non_empty_authoritative_state(compiled_state) else "no" - ) + state_injected = "yes" if _has_non_empty_authoritative_state(engine) else "no" response = await self._forward_passthrough( - body, __user__, __request__, state=compiled_state + body, __user__, __request__, engine=engine ) return self._with_trace( response, From 8aa32eb464cc7a16b3cdbb28ceef9296457cb09b Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 01:18:11 -0400 Subject: [PATCH 7/9] refactor: use engine state directly in litellm basic example --- .../prompt_construction/litellm/basic.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/python/examples/prompt_construction/litellm/basic.py b/python/examples/prompt_construction/litellm/basic.py index c16d6f1..f1b5cce 100644 --- a/python/examples/prompt_construction/litellm/basic.py +++ b/python/examples/prompt_construction/litellm/basic.py @@ -139,14 +139,14 @@ def _build_trace_text( return "\n".join(lines) -def _render_compiled_state_contract(compiled_state: _EngineSnapshot) -> str: - premise = compiled_state["premise"] +def _render_compiled_state_contract(engine: Engine) -> str: + premise = engine.premise use_items = sorted( - key for key, value in compiled_state["policies"].items() if value == POLICY_USE + key for key, value in engine.policies.items() if value == POLICY_USE ) prohibit_items = sorted( key - for key, value in compiled_state["policies"].items() + for key, value in engine.policies.items() if value == POLICY_PROHIBIT ) @@ -162,14 +162,12 @@ def _render_compiled_state_contract(compiled_state: _EngineSnapshot) -> str: return "Host policy contract:\n" + "\n".join(f"- {line}" for line in lines) -def _build_messages( - user_input: str, compiled_state: _EngineSnapshot -) -> list[dict[str, str]]: +def _build_messages(user_input: str, engine: Engine) -> list[dict[str, str]]: return [ { "role": "system", "content": "You are a helpful assistant.\n" - + _render_compiled_state_contract(compiled_state), + + _render_compiled_state_contract(engine), }, {"role": "user", "content": user_input}, ] @@ -339,7 +337,7 @@ def handle_turn( llm_called=False, ) - messages = _build_messages(user_input, _snapshot_engine_state(engine)) + messages = _build_messages(user_input, engine) response_text = _call_litellm(messages) return _append_trace( response_text, From e3047f558963d39fca4b77ebd112c8d004b93d05 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 01:22:08 -0400 Subject: [PATCH 8/9] chore: perform ruff fixes --- python/examples/prompt_construction/litellm/basic.py | 4 +--- .../litellm_proxy/_litellm_support.py | 6 ++---- .../openwebui_pipe/open_webui_pipe.py | 8 ++++++-- .../open_webui_pipe_with_directive_drafter.py | 12 +++++++++--- .../test_openwebui_pipe_with_directive_drafter.py | 11 +++++++++-- 5 files changed, 27 insertions(+), 14 deletions(-) diff --git a/python/examples/prompt_construction/litellm/basic.py b/python/examples/prompt_construction/litellm/basic.py index f1b5cce..40b5ff4 100644 --- a/python/examples/prompt_construction/litellm/basic.py +++ b/python/examples/prompt_construction/litellm/basic.py @@ -145,9 +145,7 @@ def _render_compiled_state_contract(engine: Engine) -> str: key for key, value in engine.policies.items() if value == POLICY_USE ) prohibit_items = sorted( - key - for key, value in engine.policies.items() - if value == POLICY_PROHIBIT + key for key, value in engine.policies.items() if value == POLICY_PROHIBIT ) lines: list[str] = ["The following constraints are authoritative."] diff --git a/python/reference_integrations/litellm_proxy/_litellm_support.py b/python/reference_integrations/litellm_proxy/_litellm_support.py index ca80eea..2deaeab 100644 --- a/python/reference_integrations/litellm_proxy/_litellm_support.py +++ b/python/reference_integrations/litellm_proxy/_litellm_support.py @@ -4,14 +4,12 @@ from context_compiler.engine import Engine -from context_compiler import POLICY_PROHIBIT, PolicyValue +from context_compiler import POLICY_PROHIBIT def render_compiled_state_contract(engine: Engine) -> str: prohibited = sorted( - key - for key, value in engine.policies.items() - if value == POLICY_PROHIBIT + key for key, value in engine.policies.items() if value == POLICY_PROHIBIT ) lines: list[str] = ["The following constraints are authoritative."] diff --git a/python/reference_integrations/openwebui_pipe/open_webui_pipe.py b/python/reference_integrations/openwebui_pipe/open_webui_pipe.py index f036f2c..2874719 100644 --- a/python/reference_integrations/openwebui_pipe/open_webui_pipe.py +++ b/python/reference_integrations/openwebui_pipe/open_webui_pipe.py @@ -577,7 +577,9 @@ async def _forward_passthrough( """Forward with model override and optional compiler-owned state injection.""" payload = {**body} payload["model"] = self.valves.BASE_MODEL_ID - payload["messages"] = _build_forward_messages(body.get("messages"), engine=engine) + payload["messages"] = _build_forward_messages( + body.get("messages"), engine=engine + ) user = Users.get_user_by_id(user_payload["id"]) if inspect.isawaitable(user): user = await user @@ -689,7 +691,9 @@ async def pipe( llm_called=False, ) if decision["kind"] == DecisionKind.NO_DIRECTIVE: - state_injected = "yes" if _has_non_empty_authoritative_state(engine) else "no" + state_injected = ( + "yes" if _has_non_empty_authoritative_state(engine) else "no" + ) response = await self._forward_passthrough( body, __user__, __request__, engine=engine ) 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 501692e..d9a13aa 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 @@ -718,7 +718,9 @@ async def _forward_passthrough( ) payload = {**body} payload["model"] = base_model_id - payload["messages"] = _build_forward_messages(body.get("messages"), engine=engine) + payload["messages"] = _build_forward_messages( + body.get("messages"), engine=engine + ) user = Users.get_user_by_id(user_payload["id"]) if inspect.isawaitable(user): user = await user @@ -821,7 +823,9 @@ async def pipe( logger.debug("preprocessor: drafted_result=%r", drafted_result) if not isinstance(drafted_result.result, CanonicalDirective): - state_injected = "yes" if _has_non_empty_authoritative_state(engine) else "no" + state_injected = ( + "yes" if _has_non_empty_authoritative_state(engine) else "no" + ) response = await self._forward_passthrough( body, __user__, @@ -869,7 +873,9 @@ async def pipe( llm_called=False, ) if decision["kind"] == DecisionKind.NO_DIRECTIVE: - state_injected = "yes" if _has_non_empty_authoritative_state(engine) else "no" + state_injected = ( + "yes" if _has_non_empty_authoritative_state(engine) else "no" + ) response = await self._forward_passthrough( body, __user__, diff --git a/python/tests/test_openwebui_pipe_with_directive_drafter.py b/python/tests/test_openwebui_pipe_with_directive_drafter.py index 93a6dd2..6b0b14c 100644 --- a/python/tests/test_openwebui_pipe_with_directive_drafter.py +++ b/python/tests/test_openwebui_pipe_with_directive_drafter.py @@ -8,7 +8,11 @@ import pytest from context_compiler.grammar import CanonicalDirective, DirectiveKind -from context_compiler_directive_drafter import DraftResult, NoDirective, UnknownDirective +from context_compiler_directive_drafter import ( + DraftResult, + NoDirective, + UnknownDirective, +) REPO_ROOT = Path(__file__).resolve().parents[2] MODULE_PATH = ( @@ -200,6 +204,7 @@ async def no_draft(*args, **kwargs): __chat_id__="chat-failed-transition", ) ) + async def follow_up_no_directive(*args, **kwargs): return DraftResult( source="test", @@ -392,7 +397,9 @@ async def no_draft(*args, **kwargs): assert len(forwarded) == 1 -def test_no_directive_passthrough_does_not_change_existing_engine_state(monkeypatch) -> None: +def test_no_directive_passthrough_does_not_change_existing_engine_state( + monkeypatch, +) -> None: module = _load_module("owui_with_drafter_no_directive_state_preserved", monkeypatch) forwarded: list[dict[str, object]] = [] From 6724caed93de48009d61bf9244d365a7557a2d6f Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 01:28:05 -0400 Subject: [PATCH 9/9] chore: fix mypy failures --- ...context_compiler_precall_hook_with_directive_drafter.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) 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 a9b07da..9741c06 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 @@ -179,7 +179,7 @@ async def async_pre_call_hook( "litellm_proxy: latest_user_text_present=%s", latest_user_text is not None ) drafted_result: DraftResult | None = None - decision: dict[str, object] + decision: Any if latest_user_text is not None: drafted_result = _draft_last_user_message(latest_user_text) @@ -195,7 +195,10 @@ async def async_pre_call_hook( if decision["kind"] == DecisionKind.ERROR: logger.debug("litellm_proxy: rejecting_failed_application=true") - return decision.get("message") or "Request rejected." + message = decision.get("message") + return ( + message if isinstance(message, str) and message else "Request rejected." + ) if session.mode == MODE_PERSISTENT and session.session_key is not None: CHECKPOINT_STORE.save(