From 12f7434cb412a7905bc2c6e0a260898f4607aa4e Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 23:38:15 -0400 Subject: [PATCH 1/7] refactor: migrate LiteLLM example to DirectiveDrafter API --- .../litellm/with_directive_drafter.py | 67 +++------- .../test_litellm_with_directive_drafter.py | 119 +++++++----------- 2 files changed, 63 insertions(+), 123 deletions(-) diff --git a/python/examples/prompt_construction/litellm/with_directive_drafter.py b/python/examples/prompt_construction/litellm/with_directive_drafter.py index 1c3393c..0f86e58 100644 --- a/python/examples/prompt_construction/litellm/with_directive_drafter.py +++ b/python/examples/prompt_construction/litellm/with_directive_drafter.py @@ -31,11 +31,10 @@ is_update, ) from context_compiler.engine import Engine +from context_compiler.grammar import CanonicalDirective from context_compiler_directive_drafter import ( - DRAFT_OUTCOME_DIRECTIVE, + DirectiveDrafter, get_converter_prompt, - parse_preprocessor_output, - preprocess_heuristic, ) try: @@ -61,19 +60,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") - ) - - class _LiteLLMCallKwargs(TypedDict, total=False): model: str messages: list[dict[str, str]] @@ -109,6 +95,12 @@ def _snapshot_engine_state(engine: Engine) -> _EngineSnapshot: return {"premise": engine.premise, "policies": dict(engine.policies)} +_DIRECTIVE_DRAFTER = DirectiveDrafter( + fallback=lambda message: _llm_fallback_candidate(message), + fallback_source="litellm_fallback", +) + + def _render_state_lines(state: object) -> list[str]: if not isinstance(state, dict): return ["- unavailable"] @@ -232,9 +224,7 @@ def _call_litellm(messages: list[dict[str, str]]) -> str: return content -def _llm_fallback_preprocess(message: str, state: _EngineSnapshot) -> str | None: - del state - +def _llm_fallback_candidate(message: str) -> str | None: try: completion = _get_litellm_completion() except ModuleNotFoundError: @@ -264,44 +254,23 @@ def _llm_fallback_preprocess(message: str, state: _EngineSnapshot) -> str | None try: response = completion(**kwargs) - raw_output = _extract_response_content(response) + return _extract_response_content(response) except Exception: return None - parsed = parse_preprocessor_output(raw_output) - if parsed is None: - return None - return parsed.text - def _preprocess_user_input(message: str, state: _EngineSnapshot) -> str | None: - # Heuristic first (fast + high precision), then optional LLM fallback. - try: - heuristic_result = preprocess_heuristic(message) - logger.debug("preprocessor: heuristic_outcome=%s", heuristic_result["outcome"]) - if ( - heuristic_result["outcome"] == DRAFT_OUTCOME_DIRECTIVE - and heuristic_result["directive"] - ): - parsed = parse_preprocessor_output(heuristic_result["directive"]) - logger.debug( - "preprocessor: heuristic_directive=%r", heuristic_result["directive"] - ) - if parsed is not None: - return parsed.text - except Exception: - logger.debug("preprocessor: heuristic_exception", exc_info=True) - - if _is_directive_shaped_input(message): - return None - + del state try: - fallback_directive = _llm_fallback_preprocess(message, state) - logger.debug("preprocessor: fallback_directive=%r", fallback_directive) - return fallback_directive + drafted_result = _DIRECTIVE_DRAFTER.draft_directive(message) + logger.debug("preprocessor: drafted_result=%r", drafted_result) + if isinstance(drafted_result.result, CanonicalDirective): + return drafted_result.result.text except Exception: - # Safe no-op fallback: if preprocessor path fails, preserve basic behavior. + # Safe no-op fallback: if drafter path fails, preserve basic behavior. + logger.debug("preprocessor: drafter_exception", exc_info=True) return None + return None def _render_item_label(value: str) -> str: diff --git a/python/tests/test_litellm_with_directive_drafter.py b/python/tests/test_litellm_with_directive_drafter.py index 8b0ecb3..b345569 100644 --- a/python/tests/test_litellm_with_directive_drafter.py +++ b/python/tests/test_litellm_with_directive_drafter.py @@ -1,8 +1,10 @@ -from types import SimpleNamespace +from types import MappingProxyType from typing import Any import pytest from context_compiler import create_engine +from context_compiler.grammar import CanonicalDirective, DirectiveKind +from context_compiler_directive_drafter import DraftResult, NoDirective from context_compiler_example_integrations.examples.prompt_construction.litellm import ( with_directive_drafter as module, @@ -27,17 +29,16 @@ def step_with_capture(user_input: str): monkeypatch.setattr(engine, "step", step_with_capture) monkeypatch.setattr( - module, - "preprocess_heuristic", - lambda message: { - "outcome": module.DRAFT_OUTCOME_DIRECTIVE, - "directive": "use docker", - }, - ) - monkeypatch.setattr( - module, - "parse_preprocessor_output", - lambda value, **kwargs: SimpleNamespace(text=value), + module._DIRECTIVE_DRAFTER, + "draft_directive", + lambda message: DraftResult( + source="test", + result=CanonicalDirective( + text="use docker", + kind=DirectiveKind.USE_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ), ) result = module.handle_turn("please use docker", engine) @@ -77,11 +78,10 @@ def step_with_capture(user_input: str): monkeypatch.setattr(engine, "step", step_with_capture) monkeypatch.setattr( - module, - "preprocess_heuristic", - lambda message: {"outcome": "no_directive", "directive": None}, + module._DIRECTIVE_DRAFTER, + "draft_directive", + lambda message: DraftResult(source="test", result=NoDirective("not a directive")), ) - monkeypatch.setattr(module, "_llm_fallback_preprocess", lambda message, state: None) def downstream(messages: list[dict[str, str]]) -> str: llm_calls.append(messages) @@ -107,28 +107,26 @@ def should_not_call(messages: list[dict[str, str]]) -> str: monkeypatch.setattr(module, "_call_litellm", should_not_call) monkeypatch.setattr( - module, - "preprocess_heuristic", - lambda message: { - "outcome": module.DRAFT_OUTCOME_DIRECTIVE, - "directive": "use docker", - }, - ) - monkeypatch.setattr( - module, - "parse_preprocessor_output", - lambda value, **kwargs: SimpleNamespace(text=value), + module._DIRECTIVE_DRAFTER, + "draft_directive", + lambda message: DraftResult( + source="test", + result=CanonicalDirective( + text="use docker", + kind=DirectiveKind.USE_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ), ) update_engine = create_engine() update = module.handle_turn("please use docker", update_engine) monkeypatch.setattr( - module, - "preprocess_heuristic", - lambda message: {"outcome": "no_directive", "directive": None}, + module._DIRECTIVE_DRAFTER, + "draft_directive", + lambda message: DraftResult(source="test", result=NoDirective("not a directive")), ) - monkeypatch.setattr(module, "_llm_fallback_preprocess", lambda message, state: None) clarify_engine = create_engine() clarify = module.handle_turn("set premise to concise replies", clarify_engine) @@ -235,18 +233,8 @@ def completion(**kwargs: Any) -> dict[str, object]: monkeypatch.delenv("PREPROCESSOR_MODEL", raising=False) monkeypatch.setattr(module, "_get_litellm_completion", lambda: completion) monkeypatch.setattr(module, "get_converter_prompt", lambda: "prompt") - monkeypatch.setattr( - module, - "parse_preprocessor_output", - lambda value, **_kwargs: SimpleNamespace(text=value), - ) - assert ( - module._llm_fallback_preprocess( - "please use docker", {"premise": None, "policies": {}} - ) - == "use docker" - ) + assert module._llm_fallback_candidate("please use docker") == "use docker" assert seen["model"] == "openai/main-model" @@ -262,18 +250,8 @@ def completion(**kwargs: Any) -> dict[str, object]: monkeypatch.setenv("PREPROCESSOR_MODEL", "openai/preprocessor-model") monkeypatch.setattr(module, "_get_litellm_completion", lambda: completion) monkeypatch.setattr(module, "get_converter_prompt", lambda: "prompt") - monkeypatch.setattr( - module, - "parse_preprocessor_output", - lambda value, **_kwargs: SimpleNamespace(text=value), - ) - assert ( - module._llm_fallback_preprocess( - "please use docker", {"premise": None, "policies": {}} - ) - == "use docker" - ) + assert module._llm_fallback_candidate("please use docker") == "use docker" assert seen["model"] == "openai/preprocessor-model" @@ -295,44 +273,37 @@ def test_fallback_accepts_structurally_valid_output_without_source_awareness( monkeypatch.setattr(module, "get_converter_prompt", lambda: "prompt") assert ( - module._llm_fallback_preprocess( - "set premise to concise replies", {"premise": None, "policies": {}} - ) + module._llm_fallback_candidate("set premise to concise replies") == "set premise concise replies" ) -def test_directive_shaped_malformed_inputs_skip_fallback_and_use_normal_turn_flow( +def test_directive_shaped_malformed_inputs_can_fall_through_to_normal_turn_flow( monkeypatch, ) -> None: fallback_calls = 0 - downstream_calls = 0 - monkeypatch.setattr( - module, - "preprocess_heuristic", - lambda _text: {"outcome": "no_directive", "directive": None}, - ) - - def fallback(_message: str, _state: dict[str, object]) -> None: + def fallback(_message: str) -> str | None: nonlocal fallback_calls fallback_calls += 1 - raise AssertionError("fallback should not run") - - def downstream(_messages: list[dict[str, str]]) -> str: - nonlocal downstream_calls - downstream_calls += 1 - raise AssertionError("downstream should not run") + return None - monkeypatch.setattr(module, "_llm_fallback_preprocess", fallback) + monkeypatch.setattr(module, "_llm_fallback_candidate", fallback) + monkeypatch.setattr( + module, + "_DIRECTIVE_DRAFTER", + module.DirectiveDrafter( + fallback=module._llm_fallback_candidate, + fallback_source="litellm_fallback", + ), + ) monkeypatch.setattr(module, "_call_litellm", lambda _messages: "downstream reply") assert ( module.handle_turn("use docker instead of", create_engine()) == "downstream reply" ) - assert fallback_calls == 0 - assert downstream_calls == 0 + assert fallback_calls == 1 def test_compound_directives_fall_through_when_not_applied(monkeypatch) -> None: From 0e7fcaf21ffdc9d0de692cb7328a3611ff14f502 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 23:43:47 -0400 Subject: [PATCH 2/7] refactor: remove duplicated drafter parsing from LiteLLM example --- .../litellm/with_directive_drafter.py | 22 ++++---- .../test_litellm_with_directive_drafter.py | 50 +++++++------------ 2 files changed, 27 insertions(+), 45 deletions(-) diff --git a/python/examples/prompt_construction/litellm/with_directive_drafter.py b/python/examples/prompt_construction/litellm/with_directive_drafter.py index 0f86e58..073fab5 100644 --- a/python/examples/prompt_construction/litellm/with_directive_drafter.py +++ b/python/examples/prompt_construction/litellm/with_directive_drafter.py @@ -2,12 +2,11 @@ Flow: 1. Extract user input -2. Run heuristic directive drafter -3. If no directive, run LLM fallback directive drafter using prompt files -4. Pass directive (or original input) to engine.step(...) -5. clarify -> return prompt_to_user (no model call) -6. update -> return deterministic acknowledgment text (no model call) -7. passthrough -> call LiteLLM with compiled state + user input +2. Ask DirectiveDrafter to draft one directive +3. Pass drafted directive text (or original input) to engine.step(...) +4. clarify -> return prompt_to_user (no model call) +5. update -> return deterministic acknowledgment text (no model call) +6. passthrough -> call LiteLLM with compiled state + user input Intended host usage: - collect user input @@ -31,7 +30,6 @@ is_update, ) from context_compiler.engine import Engine -from context_compiler.grammar import CanonicalDirective from context_compiler_directive_drafter import ( DirectiveDrafter, get_converter_prompt, @@ -259,13 +257,13 @@ def _llm_fallback_candidate(message: str) -> str | None: return None -def _preprocess_user_input(message: str, state: _EngineSnapshot) -> str | None: - del state +def _preprocess_user_input(message: str) -> str | None: try: drafted_result = _DIRECTIVE_DRAFTER.draft_directive(message) logger.debug("preprocessor: drafted_result=%r", drafted_result) - if isinstance(drafted_result.result, CanonicalDirective): - return drafted_result.result.text + draft_text = getattr(drafted_result.result, "text", None) + if isinstance(draft_text, str): + return draft_text except Exception: # Safe no-op fallback: if drafter path fails, preserve basic behavior. logger.debug("preprocessor: drafter_exception", exc_info=True) @@ -365,7 +363,7 @@ def handle_turn( state_before = _snapshot_engine_state(engine) del session_key preprocessd: str | None = None - preprocessd = _preprocess_user_input(user_input, _snapshot_engine_state(engine)) + preprocessd = _preprocess_user_input(user_input) compile_input = preprocessd if preprocessd else user_input logger.debug( "preprocessor: engine_input=%s", diff --git a/python/tests/test_litellm_with_directive_drafter.py b/python/tests/test_litellm_with_directive_drafter.py index b345569..ab1732a 100644 --- a/python/tests/test_litellm_with_directive_drafter.py +++ b/python/tests/test_litellm_with_directive_drafter.py @@ -1,10 +1,8 @@ -from types import MappingProxyType from typing import Any import pytest from context_compiler import create_engine -from context_compiler.grammar import CanonicalDirective, DirectiveKind -from context_compiler_directive_drafter import DraftResult, NoDirective +from context_compiler_directive_drafter import DirectiveDrafter from context_compiler_example_integrations.examples.prompt_construction.litellm import ( with_directive_drafter as module, @@ -29,16 +27,9 @@ def step_with_capture(user_input: str): monkeypatch.setattr(engine, "step", step_with_capture) monkeypatch.setattr( - module._DIRECTIVE_DRAFTER, - "draft_directive", - lambda message: DraftResult( - source="test", - result=CanonicalDirective( - text="use docker", - kind=DirectiveKind.USE_ITEM, - operands=MappingProxyType({"item": "docker"}), - ), - ), + module, + "_DIRECTIVE_DRAFTER", + DirectiveDrafter(fallback=lambda _message: "use docker"), ) result = module.handle_turn("please use docker", engine) @@ -78,9 +69,9 @@ def step_with_capture(user_input: str): monkeypatch.setattr(engine, "step", step_with_capture) monkeypatch.setattr( - module._DIRECTIVE_DRAFTER, - "draft_directive", - lambda message: DraftResult(source="test", result=NoDirective("not a directive")), + module, + "_DIRECTIVE_DRAFTER", + DirectiveDrafter(fallback=lambda _message: None), ) def downstream(messages: list[dict[str, str]]) -> str: @@ -107,25 +98,18 @@ def should_not_call(messages: list[dict[str, str]]) -> str: monkeypatch.setattr(module, "_call_litellm", should_not_call) monkeypatch.setattr( - module._DIRECTIVE_DRAFTER, - "draft_directive", - lambda message: DraftResult( - source="test", - result=CanonicalDirective( - text="use docker", - kind=DirectiveKind.USE_ITEM, - operands=MappingProxyType({"item": "docker"}), - ), - ), + module, + "_DIRECTIVE_DRAFTER", + DirectiveDrafter(fallback=lambda _message: "use docker"), ) update_engine = create_engine() update = module.handle_turn("please use docker", update_engine) monkeypatch.setattr( - module._DIRECTIVE_DRAFTER, - "draft_directive", - lambda message: DraftResult(source="test", result=NoDirective("not a directive")), + module, + "_DIRECTIVE_DRAFTER", + DirectiveDrafter(fallback=lambda _message: None), ) clarify_engine = create_engine() clarify = module.handle_turn("set premise to concise replies", clarify_engine) @@ -292,7 +276,7 @@ def fallback(_message: str) -> str | None: monkeypatch.setattr( module, "_DIRECTIVE_DRAFTER", - module.DirectiveDrafter( + DirectiveDrafter( fallback=module._llm_fallback_candidate, fallback_source="litellm_fallback", ), @@ -317,8 +301,8 @@ def downstream(_messages: list[dict[str, str]]) -> str: monkeypatch.setattr(module, "_call_litellm", downstream) monkeypatch.setattr( module, - "_preprocess_user_input", - lambda _message, _state: "use docker and prohibit peanuts", + "_DIRECTIVE_DRAFTER", + DirectiveDrafter(fallback=lambda _message: "use docker and prohibit peanuts"), ) result = module.handle_turn( @@ -363,7 +347,7 @@ def test_session_key_no_longer_restores_or_persists_checkpoint_state( monkeypatch, ) -> None: monkeypatch.setattr(module, "_call_litellm", lambda _messages: "ok") - monkeypatch.setattr(module, "_preprocess_user_input", lambda _text, _state: None) + monkeypatch.setattr(module, "_preprocess_user_input", lambda _text: None) first_engine = create_engine() assert module.handle_turn("hello", first_engine, session_key="s1") == "ok" From b69bcadc1d9ff5d0f1d0120b709039258ec23510 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 23:48:40 -0400 Subject: [PATCH 3/7] refactor: simplify litellm directive drafter example --- .../litellm/with_directive_drafter.py | 42 ++++------ .../test_litellm_with_directive_drafter.py | 82 ++++--------------- 2 files changed, 33 insertions(+), 91 deletions(-) diff --git a/python/examples/prompt_construction/litellm/with_directive_drafter.py b/python/examples/prompt_construction/litellm/with_directive_drafter.py index 073fab5..fb0bea6 100644 --- a/python/examples/prompt_construction/litellm/with_directive_drafter.py +++ b/python/examples/prompt_construction/litellm/with_directive_drafter.py @@ -2,11 +2,12 @@ Flow: 1. Extract user input -2. Ask DirectiveDrafter to draft one directive -3. Pass drafted directive text (or original input) to engine.step(...) -4. clarify -> return prompt_to_user (no model call) -5. update -> return deterministic acknowledgment text (no model call) -6. passthrough -> call LiteLLM with compiled state + user input +2. Ask DirectiveDrafter to draft one directive, using LiteLLM only as fallback +3. Observe the returned DraftResult and extract drafted directive text when present +4. Pass drafted directive text, or the original input, to engine.step(...) +5. If the compiler returns an error or near-miss clarify, return that text locally +6. If the compiler applies an update, return a deterministic acknowledgment locally +7. Otherwise call LiteLLM with compiled state + user input Intended host usage: - collect user input @@ -31,19 +32,11 @@ ) from context_compiler.engine import Engine from context_compiler_directive_drafter import ( + DraftResult, DirectiveDrafter, get_converter_prompt, ) -try: - from .confirmation_helper import ( - is_confirmation_text, - ) -except ImportError: - from confirmation_helper import ( - is_confirmation_text, - ) - from context_compiler_example_integrations.examples._shared.provider_mode import ( print_startup_config, resolve_provider_config, @@ -261,9 +254,7 @@ def _preprocess_user_input(message: str) -> str | None: try: drafted_result = _DIRECTIVE_DRAFTER.draft_directive(message) logger.debug("preprocessor: drafted_result=%r", drafted_result) - draft_text = getattr(drafted_result.result, "text", None) - if isinstance(draft_text, str): - return draft_text + return _extract_drafted_text(drafted_result) except Exception: # Safe no-op fallback: if drafter path fails, preserve basic behavior. logger.debug("preprocessor: drafter_exception", exc_info=True) @@ -271,6 +262,13 @@ def _preprocess_user_input(message: str) -> str | None: return None +def _extract_drafted_text(drafted_result: DraftResult) -> str | None: + draft_text = getattr(drafted_result.result, "text", None) + if isinstance(draft_text, str): + return draft_text + return None + + def _render_item_label(value: str) -> str: return re.sub(r"\s+", " ", value).strip().lower() @@ -357,11 +355,8 @@ def _append_trace( return f"{response_text}\n\n{trace_text}" -def handle_turn( - user_input: str, engine: Engine, *, session_key: str | None = None -) -> str: +def handle_turn(user_input: str, engine: Engine) -> str: state_before = _snapshot_engine_state(engine) - del session_key preprocessd: str | None = None preprocessd = _preprocess_user_input(user_input) compile_input = preprocessd if preprocessd else user_input @@ -404,10 +399,7 @@ def handle_turn( llm_called=False, ) if is_update(decision): - if is_confirmation_text(user_input): - response_text = "State updated." - else: - response_text = _summarize_update_from_input(compile_input) + response_text = _summarize_update_from_input(compile_input) return _append_trace( response_text, original_input=user_input, diff --git a/python/tests/test_litellm_with_directive_drafter.py b/python/tests/test_litellm_with_directive_drafter.py index ab1732a..5d35b7f 100644 --- a/python/tests/test_litellm_with_directive_drafter.py +++ b/python/tests/test_litellm_with_directive_drafter.py @@ -38,25 +38,6 @@ def step_with_capture(user_input: str): assert compile_inputs == ["use docker"] -def test_follow_up_confirmation_is_not_treated_as_pending_resume(monkeypatch) -> None: - engine = create_engine() - first = module.handle_turn("use docker instead of kubectl", engine) - assert first == "State updated: Use docker." - - llm_calls: list[list[dict[str, str]]] = [] - - def downstream(messages: list[dict[str, str]]) -> str: - llm_calls.append(messages) - return "stubbed reply" - - monkeypatch.setattr(module, "_call_litellm", downstream) - - second = module.handle_turn("yes", engine) - - assert second == "stubbed reply" - assert llm_calls - - def test_unknown_or_unsafe_drafting_falls_back_to_raw_input(monkeypatch) -> None: engine = create_engine() compile_inputs: list[str] = [] @@ -87,6 +68,19 @@ def downstream(messages: list[dict[str, str]]) -> str: assert len(llm_calls) == 1 +def test_extract_drafted_text_observes_draft_result_behavior() -> None: + drafter = DirectiveDrafter(fallback=lambda _message: "use docker") + drafted_result = drafter.draft_directive("please use docker") + + assert module._extract_drafted_text(drafted_result) == "use docker" + + no_directive_result = DirectiveDrafter( + fallback=lambda _message: None + ).draft_directive("hello there") + + assert module._extract_drafted_text(no_directive_result) is None + + def test_local_update_and_clarify_responses_skip_downstream_litellm_call( monkeypatch, ) -> None: @@ -314,54 +308,10 @@ def downstream(_messages: list[dict[str, str]]) -> str: assert downstream_calls == 1 -def test_confirmation_follow_up_does_not_resume_removed_checkpoint_flow( - monkeypatch, -) -> None: - first_engine = create_engine() - clarify = module.handle_turn( - "use kubectl instead of docker", - first_engine, - session_key="resume-with-drafter", - ) - llm_calls = 0 - - def downstream(_messages: list[dict[str, str]]) -> str: - nonlocal llm_calls - llm_calls += 1 - return "downstream reply" - - monkeypatch.setattr(module, "_call_litellm", downstream) - resumed_engine = create_engine() - resumed = module.handle_turn( - "yes", resumed_engine, session_key="resume-with-drafter" - ) - - assert clarify == "State updated: Use kubectl." - assert resumed == "downstream reply" - assert llm_calls == 1 - assert dict(first_engine.policies) == {"kubectl": "use"} - assert dict(resumed_engine.policies) == {} - - -def test_session_key_no_longer_restores_or_persists_checkpoint_state( - monkeypatch, -) -> None: +def test_handle_turn_has_no_session_or_resume_behavior(monkeypatch) -> None: monkeypatch.setattr(module, "_call_litellm", lambda _messages: "ok") monkeypatch.setattr(module, "_preprocess_user_input", lambda _text: None) - first_engine = create_engine() - assert module.handle_turn("hello", first_engine, session_key="s1") == "ok" - - update_engine = create_engine() - assert ( - module.handle_turn("use docker", update_engine, session_key="s1") - == "State updated: Use docker." - ) + engine = create_engine() - clarify_engine = create_engine() - assert ( - module.handle_turn( - "use kubectl instead of docker", clarify_engine, session_key="s1" - ) - == "State updated: Use kubectl." - ) + assert module.handle_turn("hello", engine) == "ok" From d8aaf9543992bfc2fcf88a73b7948f5b103acc51 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 23:51:22 -0400 Subject: [PATCH 4/7] refactor: use engine state in litellm example --- .../litellm/with_directive_drafter.py | 71 +++++++++---------- 1 file changed, 33 insertions(+), 38 deletions(-) diff --git a/python/examples/prompt_construction/litellm/with_directive_drafter.py b/python/examples/prompt_construction/litellm/with_directive_drafter.py index fb0bea6..3083c37 100644 --- a/python/examples/prompt_construction/litellm/with_directive_drafter.py +++ b/python/examples/prompt_construction/litellm/with_directive_drafter.py @@ -46,11 +46,6 @@ SHOW_CONTEXT_COMPILER_TRACE = False -class _EngineSnapshot(TypedDict): - premise: str | None - policies: dict[str, PolicyValue] - - class _LiteLLMCallKwargs(TypedDict, total=False): model: str messages: list[dict[str, str]] @@ -82,22 +77,13 @@ def _extract_response_content(response: object) -> str | None: return None -def _snapshot_engine_state(engine: Engine) -> _EngineSnapshot: - return {"premise": engine.premise, "policies": dict(engine.policies)} - - _DIRECTIVE_DRAFTER = DirectiveDrafter( fallback=lambda message: _llm_fallback_candidate(message), fallback_source="litellm_fallback", ) -def _render_state_lines(state: object) -> list[str]: - if not isinstance(state, dict): - return ["- unavailable"] - raw_policies = state.get("policies") - policies = raw_policies if isinstance(raw_policies, dict) else {} - premise = state.get("premise") +def _render_state_lines(premise: str | None, policies: Mapping[str, PolicyValue]) -> list[str]: use_items = sorted( key for key, value in policies.items() @@ -123,8 +109,10 @@ def _build_trace_text( compiler_input: str, preprocessor_output: str | None, decision: object, - state_before: object, - state_after: object, + premise_before: str | None, + policies_before: Mapping[str, PolicyValue], + premise_after: str | None, + policies_after: Mapping[str, PolicyValue], llm_called: bool, ) -> str: kind = decision.get("kind", "unknown") if isinstance(decision, dict) else "unknown" @@ -136,14 +124,19 @@ def _build_trace_text( f"- decision: {kind}", f"- llm_called: {'yes' if llm_called else 'no'}", ] - if isinstance(state_before, dict) and isinstance(state_after, dict): - lines.append( - f"- state_changed: {'yes' if state_before != state_after else 'no'}" + lines.append( + "- state_changed: " + + ( + "yes" + if premise_before != premise_after + or dict(policies_before) != dict(policies_after) + else "no" ) + ) lines.append("state_before:") - lines.extend(_render_state_lines(state_before)) + lines.extend(_render_state_lines(premise_before, policies_before)) lines.append("state_after:") - lines.extend(_render_state_lines(state_after)) + lines.extend(_render_state_lines(premise_after, policies_after)) return "\n".join(lines) @@ -152,14 +145,14 @@ def _get_litellm_completion() -> Callable[..., object]: return cast(Callable[..., object], litellm_module.completion) -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 ) @@ -176,13 +169,13 @@ def _render_compiled_state_contract(compiled_state: _EngineSnapshot) -> str: def _build_messages( - user_input: str, compiled_state: _EngineSnapshot + 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}, ] @@ -337,8 +330,8 @@ def _append_trace( compiler_input: str, preprocessor_output: str | None, decision: object, - state_before: object, - state_after: object, + state_before: tuple[str | None, dict[str, PolicyValue]], + state_after: tuple[str | None, dict[str, PolicyValue]], llm_called: bool, ) -> str: if not SHOW_CONTEXT_COMPILER_TRACE: @@ -348,15 +341,17 @@ def _append_trace( compiler_input=compiler_input, preprocessor_output=preprocessor_output, decision=decision, - state_before=state_before, - state_after=state_after, + premise_before=state_before[0], + policies_before=state_before[1], + premise_after=state_after[0], + policies_after=state_after[1], llm_called=llm_called, ) return f"{response_text}\n\n{trace_text}" def handle_turn(user_input: str, engine: Engine) -> str: - state_before = _snapshot_engine_state(engine) + 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 @@ -384,7 +379,7 @@ def handle_turn(user_input: str, engine: Engine) -> str: preprocessor_output=preprocessd, decision=decision, state_before=state_before, - state_after=_snapshot_engine_state(engine), + state_after=(engine.premise, dict(engine.policies)), llm_called=False, ) if near_miss_prompt is not None and decision["kind"] == DecisionKind.NO_DIRECTIVE: @@ -395,7 +390,7 @@ def handle_turn(user_input: str, engine: Engine) -> str: preprocessor_output=preprocessd, decision={"kind": DecisionKind.ERROR, "message": near_miss_prompt}, state_before=state_before, - state_after=_snapshot_engine_state(engine), + state_after=(engine.premise, dict(engine.policies)), llm_called=False, ) if is_update(decision): @@ -407,10 +402,10 @@ def handle_turn(user_input: str, engine: Engine) -> str: preprocessor_output=preprocessd, decision=decision, state_before=state_before, - state_after=_snapshot_engine_state(engine), + state_after=(engine.premise, dict(engine.policies)), 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, @@ -419,6 +414,6 @@ def handle_turn(user_input: str, engine: Engine) -> str: preprocessor_output=preprocessd, decision=decision, state_before=state_before, - state_after=_snapshot_engine_state(engine), + state_after=(engine.premise, dict(engine.policies)), llm_called=True, ) From 94da375f3b49c031ce0317ce05b9ef0d31a9263f Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 23:54:09 -0400 Subject: [PATCH 5/7] refactor: remove litellm example grammar matching --- .../litellm/with_directive_drafter.py | 78 +------------------ .../test_litellm_with_directive_drafter.py | 30 +++++-- 2 files changed, 27 insertions(+), 81 deletions(-) diff --git a/python/examples/prompt_construction/litellm/with_directive_drafter.py b/python/examples/prompt_construction/litellm/with_directive_drafter.py index 3083c37..b95bc5d 100644 --- a/python/examples/prompt_construction/litellm/with_directive_drafter.py +++ b/python/examples/prompt_construction/litellm/with_directive_drafter.py @@ -17,7 +17,6 @@ import logging import os -import re from collections.abc import Callable, Mapping, Sequence from importlib import import_module from typing import TypedDict, cast @@ -262,67 +261,6 @@ def _extract_drafted_text(drafted_result: DraftResult) -> str | None: return None -def _render_item_label(value: str) -> str: - return re.sub(r"\s+", " ", value).strip().lower() - - -def _near_miss_directive_clarify(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 _append_trace( response_text: str, *, @@ -368,10 +306,9 @@ def handle_turn(user_input: str, engine: Engine) -> str: else: kind = DecisionKind.NO_DIRECTIVE.value logger.debug("preprocessor: decision=%s", kind) - near_miss_prompt = _near_miss_directive_clarify(user_input) if decision["kind"] == DecisionKind.ERROR: - response_text = near_miss_prompt or decision["message"] or "" + response_text = decision["message"] or "" return _append_trace( response_text, original_input=user_input, @@ -382,19 +319,8 @@ def handle_turn(user_input: str, engine: Engine) -> str: state_after=(engine.premise, dict(engine.policies)), llm_called=False, ) - if near_miss_prompt is not None and decision["kind"] == DecisionKind.NO_DIRECTIVE: - return _append_trace( - near_miss_prompt, - original_input=user_input, - compiler_input=compile_input, - preprocessor_output=preprocessd, - decision={"kind": DecisionKind.ERROR, "message": near_miss_prompt}, - state_before=state_before, - state_after=(engine.premise, dict(engine.policies)), - llm_called=False, - ) if is_update(decision): - response_text = _summarize_update_from_input(compile_input) + response_text = "State updated." return _append_trace( response_text, original_input=user_input, diff --git a/python/tests/test_litellm_with_directive_drafter.py b/python/tests/test_litellm_with_directive_drafter.py index 5d35b7f..6dea0a2 100644 --- a/python/tests/test_litellm_with_directive_drafter.py +++ b/python/tests/test_litellm_with_directive_drafter.py @@ -34,7 +34,7 @@ def step_with_capture(user_input: str): result = module.handle_turn("please use docker", engine) - assert result == "State updated: Use docker." + assert result == "State updated." assert compile_inputs == ["use docker"] @@ -81,7 +81,7 @@ def test_extract_drafted_text_observes_draft_result_behavior() -> None: assert module._extract_drafted_text(no_directive_result) is None -def test_local_update_and_clarify_responses_skip_downstream_litellm_call( +def test_local_update_responses_skip_downstream_litellm_call( monkeypatch, ) -> None: llm_calls: list[object] = [] @@ -105,12 +105,32 @@ def should_not_call(messages: list[dict[str, str]]) -> str: "_DIRECTIVE_DRAFTER", DirectiveDrafter(fallback=lambda _message: None), ) + + assert update == "State updated." + assert llm_calls == [] + + +def test_malformed_directive_like_input_falls_through_to_downstream_litellm( + monkeypatch, +) -> None: + llm_calls: list[object] = [] + + def downstream(messages: list[dict[str, str]]) -> str: + llm_calls.append(messages) + return "downstream reply" + + monkeypatch.setattr(module, "_call_litellm", downstream) + monkeypatch.setattr( + module, + "_DIRECTIVE_DRAFTER", + DirectiveDrafter(fallback=lambda _message: None), + ) + clarify_engine = create_engine() clarify = module.handle_turn("set premise to concise replies", clarify_engine) - assert update == "State updated: Use docker." - assert clarify == "Invalid premise syntax.\nUse 'set premise '." - assert llm_calls == [] + assert clarify == "downstream reply" + assert llm_calls def test_call_litellm_requires_api_key_in_openai_mode(monkeypatch) -> None: From 12bd80a048d1c9f4cb2e4783687c2d95b8b0fb36 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 00:22:14 -0400 Subject: [PATCH 6/7] refactor: handle DirectiveDrafter result variants explicitly --- .../litellm/with_directive_drafter.py | 13 ++++++-- .../test_litellm_with_directive_drafter.py | 32 ++++++++++++++++--- 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/python/examples/prompt_construction/litellm/with_directive_drafter.py b/python/examples/prompt_construction/litellm/with_directive_drafter.py index b95bc5d..f8ac13b 100644 --- a/python/examples/prompt_construction/litellm/with_directive_drafter.py +++ b/python/examples/prompt_construction/litellm/with_directive_drafter.py @@ -30,9 +30,12 @@ is_update, ) from context_compiler.engine import Engine +from context_compiler.grammar import CanonicalDirective from context_compiler_directive_drafter import ( DraftResult, DirectiveDrafter, + NoDirective, + UnknownDirective, get_converter_prompt, ) @@ -255,9 +258,13 @@ def _preprocess_user_input(message: str) -> str | None: def _extract_drafted_text(drafted_result: DraftResult) -> str | None: - draft_text = getattr(drafted_result.result, "text", None) - if isinstance(draft_text, str): - return draft_text + result = drafted_result.result + if isinstance(result, CanonicalDirective): + return result.text + if isinstance(result, NoDirective): + return None + if isinstance(result, UnknownDirective): + return None return None diff --git a/python/tests/test_litellm_with_directive_drafter.py b/python/tests/test_litellm_with_directive_drafter.py index 6dea0a2..d76a98b 100644 --- a/python/tests/test_litellm_with_directive_drafter.py +++ b/python/tests/test_litellm_with_directive_drafter.py @@ -1,8 +1,15 @@ +from types import MappingProxyType from typing import Any import pytest from context_compiler import create_engine -from context_compiler_directive_drafter import DirectiveDrafter +from context_compiler.grammar import CanonicalDirective, DirectiveKind +from context_compiler_directive_drafter import ( + DirectiveDrafter, + DraftResult, + NoDirective, + UnknownDirective, +) from context_compiler_example_integrations.examples.prompt_construction.litellm import ( with_directive_drafter as module, @@ -74,12 +81,29 @@ def test_extract_drafted_text_observes_draft_result_behavior() -> None: assert module._extract_drafted_text(drafted_result) == "use docker" - no_directive_result = DirectiveDrafter( - fallback=lambda _message: None - ).draft_directive("hello there") + no_directive_result = DraftResult(source="test", result=NoDirective("not a directive")) assert module._extract_drafted_text(no_directive_result) is None + unknown_directive_result = DraftResult( + source="test", result=UnknownDirective("unresolved") + ) + + assert module._extract_drafted_text(unknown_directive_result) is None + + +def test_extract_drafted_text_only_applies_canonical_directive() -> None: + drafted_result = DraftResult( + source="test", + result=CanonicalDirective( + text="use docker", + kind=DirectiveKind.USE_ITEM, + operands=MappingProxyType({"item": "docker"}), + ), + ) + + assert module._extract_drafted_text(drafted_result) == "use docker" + def test_local_update_responses_skip_downstream_litellm_call( monkeypatch, From 50a234c228f4bfe2f1d69dc99776bb360319a43b Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Mon, 10 Aug 2026 00:27:10 -0400 Subject: [PATCH 7/7] style: format litellm directive drafter files --- .../litellm/with_directive_drafter.py | 12 +++++------- python/tests/test_litellm_with_directive_drafter.py | 4 +++- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/python/examples/prompt_construction/litellm/with_directive_drafter.py b/python/examples/prompt_construction/litellm/with_directive_drafter.py index f8ac13b..9ff34fd 100644 --- a/python/examples/prompt_construction/litellm/with_directive_drafter.py +++ b/python/examples/prompt_construction/litellm/with_directive_drafter.py @@ -85,7 +85,9 @@ def _extract_response_content(response: object) -> str | None: ) -def _render_state_lines(premise: str | None, policies: Mapping[str, PolicyValue]) -> list[str]: +def _render_state_lines( + premise: str | None, policies: Mapping[str, PolicyValue] +) -> list[str]: use_items = sorted( key for key, value in policies.items() @@ -153,9 +155,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."] @@ -170,9 +170,7 @@ def _render_compiled_state_contract(engine: Engine) -> str: return "Host policy contract:\n" + "\n".join(f"- {line}" for line in lines) -def _build_messages( - user_input: str, engine: Engine -) -> list[dict[str, str]]: +def _build_messages(user_input: str, engine: Engine) -> list[dict[str, str]]: return [ { "role": "system", diff --git a/python/tests/test_litellm_with_directive_drafter.py b/python/tests/test_litellm_with_directive_drafter.py index d76a98b..287fc3a 100644 --- a/python/tests/test_litellm_with_directive_drafter.py +++ b/python/tests/test_litellm_with_directive_drafter.py @@ -81,7 +81,9 @@ def test_extract_drafted_text_observes_draft_result_behavior() -> None: assert module._extract_drafted_text(drafted_result) == "use docker" - no_directive_result = DraftResult(source="test", result=NoDirective("not a directive")) + no_directive_result = DraftResult( + source="test", result=NoDirective("not a directive") + ) assert module._extract_drafted_text(no_directive_result) is None