fix: migrate OpenWebUI drafter boundary and clean runtime state usage

python/examples/prompt_construction/litellm/basic.py

@@ -139,15 +139,13 @@ 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()
-        if value == POLICY_PROHIBIT
+        key for key, value in engine.policies.items() if value == POLICY_PROHIBIT
     )
 
     lines: list[str] = ["The following constraints are authoritative."]
@@ -162,14 +160,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 +335,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,

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:

python/reference_integrations/litellm_proxy/_litellm_support.py

@@ -1,49 +1,25 @@
-"""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
+from context_compiler import POLICY_PROHIBIT
 
 
-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()
-        if value == POLICY_PROHIBIT
+        key 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."

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.

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__)
@@ -179,37 +178,38 @@ 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: Any
 
         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}
 
         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."
+            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(
                 session.session_key,
                 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.

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,14 @@ 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 +691,11 @@ 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"
+                "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 +718,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,