diff --git a/python/examples/prompt_construction/litellm/basic.py b/python/examples/prompt_construction/litellm/basic.py index c16d6f1..40b5ff4 100644 --- a/python/examples/prompt_construction/litellm/basic.py +++ b/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, 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/_litellm_support.py b/python/reference_integrations/litellm_proxy/_litellm_support.py index 013069c..2deaeab 100644 --- a/python/reference_integrations/litellm_proxy/_litellm_support.py +++ b/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." 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 1c3736e..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 @@ -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,17 +178,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: 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} @@ -197,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( @@ -205,11 +206,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..2874719 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,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, 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..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 @@ -56,11 +56,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__) @@ -74,19 +76,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, @@ -123,21 +112,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)) @@ -146,18 +134,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}" @@ -212,10 +198,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: @@ -287,7 +273,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 = ( @@ -297,10 +283,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 @@ -319,73 +305,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: @@ -466,6 +385,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 +606,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 +635,67 @@ 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: + 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, + ) - if _is_directive_shaped_input(message): - return None, None + drafter = DirectiveDrafter( + async_fallback=fallback, + async_fallback_source="openwebui_fallback", + ) + return await drafter.async_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, @@ -777,7 +704,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(): @@ -791,7 +718,9 @@ 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 @@ -881,13 +810,9 @@ 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 - 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,11 +821,32 @@ async def pipe( if preprocess_error is not None: return preprocess_error - 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 + 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" + ) + 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: @@ -910,7 +856,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: @@ -918,46 +863,25 @@ 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, state_before=state_before, state_after=state_after, - 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, + preprocessor_output=compile_input, 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__, base_model_id=base_model_id, - state=compiled_state, + engine=engine, ) return self._with_trace( response, @@ -966,32 +890,29 @@ 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, ) 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, state_before=state_before, state_after=state_after, - preprocessor_output=preprocessd, + preprocessor_output=compile_input, 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, @@ -1000,7 +921,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 53db5cb..6b0b14c 100644 --- a/python/tests/test_openwebui_pipe_with_directive_drafter.py +++ b/python/tests/test_openwebui_pipe_with_directive_drafter.py @@ -4,8 +4,15 @@ 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 +110,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" @@ -124,7 +138,7 @@ async def fake_preprocess(*args, **kwargs): ) ) - assert result == "State updated: Use docker." + assert result == "State updated." assert compile_inputs == ["use docker"] @@ -146,9 +160,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 +183,16 @@ async def update_draft(*args, **kwargs): ) async def no_draft(*args, **kwargs): - return None, None + return DraftResult( + source="test", + result=CanonicalDirective( + text="prohibit docker", + kind=DirectiveKind.PROHIBIT_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ) - monkeypatch.setattr(module.Pipe, "_preprocess_user_input", no_draft) + monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) rejected = asyncio.run( pipe.pipe( { @@ -176,6 +204,14 @@ 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"}]}, @@ -185,7 +221,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.' ) @@ -202,12 +238,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 +266,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 +307,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" @@ -282,7 +331,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) @@ -300,9 +349,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,10 +372,13 @@ 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) - rejection = asyncio.run( + monkeypatch.setattr(module.Pipe, "_draft_user_input", no_draft) + passthrough = asyncio.run( pipe.pipe( { "model": "pipe-model", @@ -333,13 +392,15 @@ 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( @@ -354,12 +415,15 @@ 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( + passthrough = asyncio.run( pipe.pipe( { "model": "pipe-model", @@ -392,10 +456,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: @@ -414,9 +478,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 +524,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 +547,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 +652,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 +688,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 +722,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 +738,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()