From 3079ef2d870b1af1a95908a5f0062b259cb293b4 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 22:51:14 -0400 Subject: [PATCH 1/5] refactor: use DirectiveDrafter API in LiteLLM hook --- ...ler_precall_hook_with_directive_drafter.py | 118 +++++++------- ...st_litellm_proxy_with_directive_drafter.py | 151 ++++++++++++++---- 2 files changed, 184 insertions(+), 85 deletions(-) diff --git a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py index 9d03aa6..d6761fb 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 @@ -15,8 +15,6 @@ import os from collections.abc import Callable, Mapping, Sequence from importlib import import_module -from importlib.resources import as_file, files -from importlib.resources.abc import Traversable from typing import Any, cast try: @@ -33,11 +31,16 @@ class CustomLogger: # type: ignore[no-redef] DecisionKind, create_engine, ) +from context_compiler.grammar import CanonicalDirective from context_compiler_directive_drafter import ( DRAFT_OUTCOME_DIRECTIVE, + DRAFT_OUTCOME_NO_DIRECTIVE, + DirectiveDrafter, + DraftResult, + NoDirective, + UnknownDirective, parse_preprocessor_output, - preprocess_heuristic, - render_prompt, + validate_preprocessor_output, ) from context_compiler_example_integrations.reference_integrations.litellm_proxy._checkpoint_support import ( MODE_PERSISTENT, @@ -49,7 +52,6 @@ class CustomLogger: # type: ignore[no-redef] resolve_session_context, ) from context_compiler_example_integrations.reference_integrations.litellm_proxy._litellm_support import ( - EngineSnapshot, extract_request_messages, render_compiled_state_contract, snapshot_engine_state, @@ -64,8 +66,16 @@ class CustomLogger: # type: ignore[no-redef] "achat_completion", } -_PROMPTS_DIR = files("context_compiler_directive_drafter").joinpath("prompts") CHECKPOINT_STORE: CheckpointStore = InMemoryCheckpointStore() +_FALLBACK_SYSTEM_PROMPT = ( + "Convert the latest user message into exactly one valid Context Compiler " + "directive, or output . Use only these directive forms: " + "set premise , change premise to , use , prohibit " + ", remove policy , use instead of , " + "clear premise, reset policies, clear state. If the message is ambiguous, " + "not a direct instruction to change compiler state, or could imply more " + "than one instruction, output . Do not explain." +) def _extract_response_content(response: object) -> str | None: @@ -91,43 +101,40 @@ def _extract_response_content(response: object) -> str | None: return None -def _prompt_file_path() -> Traversable: - profile = os.getenv("PREPROCESSOR_PROMPT_PROFILE", "default").strip().lower() - if profile == "llama": - return _PROMPTS_DIR.joinpath("llama.txt") - return _PROMPTS_DIR.joinpath("default.txt") - - def _get_litellm_completion() -> Callable[..., object]: litellm_module = import_module("litellm") return cast(Callable[..., object], litellm_module.completion) -def _llm_fallback_preprocess(message: str, state: EngineSnapshot) -> str | None: - with as_file(_prompt_file_path()) as prompt_path: - prompt = render_prompt(prompt_path, state["premise"], state["policies"]) - if prompt is None: - return None - +def _llm_fallback_draft(message: str) -> DraftResult: preprocessor_model = os.getenv("PREPROCESSOR_MODEL", "").strip() if not preprocessor_model: preprocessor_model = os.getenv("MODEL", "").strip() if not preprocessor_model: - return None + return DraftResult( + source="litellm_fallback", + result=UnknownDirective(reason="fallback_model_unconfigured"), + ) api_key = os.getenv("OPENAI_API_KEY") if not api_key: - return None + return DraftResult( + source="litellm_fallback", + result=UnknownDirective(reason="fallback_api_key_missing"), + ) try: completion = _get_litellm_completion() except ModuleNotFoundError: - return None + return DraftResult( + source="litellm_fallback", + result=UnknownDirective(reason="fallback_litellm_unavailable"), + ) kwargs: dict[str, object] = { "model": preprocessor_model, "messages": [ - {"role": "system", "content": prompt}, + {"role": "system", "content": _FALLBACK_SYSTEM_PROMPT}, {"role": "user", "content": message}, ], "api_key": api_key, @@ -141,37 +148,34 @@ def _llm_fallback_preprocess(message: str, state: EngineSnapshot) -> str | None: response = completion(**kwargs) raw_output = _extract_response_content(response) except Exception: - return None - - parsed = parse_preprocessor_output(raw_output) - if parsed is None: - return None - return parsed.text - + return DraftResult( + source="litellm_fallback", + result=UnknownDirective(reason="fallback_completion_failed"), + ) -def _preprocess_last_user_message( - message: str, state: EngineSnapshot | None -) -> str | None: - try: - 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 - except Exception: - logger.debug("litellm_proxy: heuristic_exception", exc_info=True) + validated = validate_preprocessor_output(raw_output) + if validated["classification"] == DRAFT_OUTCOME_DIRECTIVE: + parsed = parse_preprocessor_output(raw_output) + if parsed is not None: + return DraftResult(source="litellm_fallback", result=parsed) + return DraftResult( + source="litellm_fallback", + result=UnknownDirective(reason="invalid_canonical_directive"), + ) + if validated["classification"] == DRAFT_OUTCOME_NO_DIRECTIVE: + return DraftResult( + source="litellm_fallback", + result=NoDirective(reason="fallback_confident_non_directive"), + ) + return DraftResult( + source="litellm_fallback", + result=UnknownDirective(reason="fallback_unresolved"), + ) - if state is None: - return None - try: - return _llm_fallback_preprocess(message, state) - except Exception: - logger.debug("litellm_proxy: fallback_exception", exc_info=True) - return None +def _draft_last_user_message(message: str) -> DraftResult: + drafter = DirectiveDrafter(fallback=_llm_fallback_draft) + return drafter.draft_directive(message) class ContextCompilerPreCallHookWithPreprocessor(CustomLogger): @@ -219,15 +223,13 @@ async def async_pre_call_hook( "litellm_proxy: latest_user_text_present=%s", latest_user_text is not None ) engine_input = latest_user_text - drafted_input: str | None = None + drafted_result: DraftResult | None = None if latest_user_text is not None: - drafted_input = _preprocess_last_user_message( - latest_user_text, snapshot_engine_state(engine) - ) - logger.debug("litellm_proxy: drafted_input=%r", drafted_input) - if drafted_input is not None: - engine_input = drafted_input + 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) diff --git a/python/tests/test_litellm_proxy_with_directive_drafter.py b/python/tests/test_litellm_proxy_with_directive_drafter.py index 1f50696..5a9989c 100644 --- a/python/tests/test_litellm_proxy_with_directive_drafter.py +++ b/python/tests/test_litellm_proxy_with_directive_drafter.py @@ -6,6 +6,7 @@ from pathlib import Path import pytest +from context_compiler.grammar import CanonicalDirective REPO_ROOT = Path(__file__).resolve().parents[2] MODULE_PATH = ( @@ -44,11 +45,14 @@ def test_drafter_runs_only_for_current_turn(monkeypatch) -> None: hook = module.ContextCompilerPreCallHookWithPreprocessor() drafted_calls: list[tuple[str, dict[str, object]]] = [] - def fake_preprocess(message: str, state: dict[str, object] | None) -> str | None: - drafted_calls.append((message, {} if state is None else dict(state))) - return None + def fake_draft(message: str) -> object: + drafted_calls.append((message, {})) + return module.DraftResult( + source="test", + result=module.NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module, "_preprocess_last_user_message", fake_preprocess) + monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) data = { "model": "demo", "context_compiler_mode": "stateless", @@ -62,7 +66,7 @@ def fake_preprocess(message: str, state: dict[str, object] | None) -> str | None result = asyncio.run(hook.async_pre_call_hook(None, None, data, "completion")) assert result is data - assert drafted_calls == [("please use docker", {"premise": None, "policies": {}})] + assert drafted_calls == [("please use docker", {})] def test_drafter_output_applies_to_current_turn_only(monkeypatch) -> None: @@ -70,8 +74,15 @@ def test_drafter_output_applies_to_current_turn_only(monkeypatch) -> None: hook = module.ContextCompilerPreCallHookWithPreprocessor() monkeypatch.setattr( module, - "_preprocess_last_user_message", - lambda message, state: "prohibit docker", + "_draft_last_user_message", + lambda message: module.DraftResult( + source="test", + result=CanonicalDirective( + kind="set_policy", + operands=("docker", "prohibit"), + text="prohibit docker", + ), + ), ) data = { "model": "demo", @@ -97,11 +108,14 @@ def test_persistent_mode_with_drafter_rejects_failed_application_without_persist hook = module.ContextCompilerPreCallHookWithPreprocessor() drafted_inputs: list[str] = [] - def fake_preprocess(message: str, state: dict[str, object] | None) -> str | None: + def fake_draft(message: str) -> object: drafted_inputs.append(message) - return None + return module.DraftResult( + source="test", + result=module.NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module, "_preprocess_last_user_message", fake_preprocess) + monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) rejected_data = { "model": "demo", "context_compiler_mode": "persistent", @@ -138,7 +152,12 @@ def test_default_mode_is_stateless_and_requires_no_session_key(monkeypatch) -> N module = _load_module(monkeypatch, "litellm_proxy_with_drafter_default_stateless") hook = module.ContextCompilerPreCallHookWithPreprocessor() monkeypatch.setattr( - module, "_preprocess_last_user_message", lambda message, state: None + module, + "_draft_last_user_message", + lambda message: module.DraftResult( + source="test", + result=module.NoDirective(reason="reject.confident_non_directive"), + ), ) data = { "model": "demo", @@ -154,7 +173,12 @@ def test_stateless_mode_has_no_cross_call_continuity(monkeypatch) -> None: module = _load_module(monkeypatch, "litellm_proxy_with_drafter_stateless") hook = module.ContextCompilerPreCallHookWithPreprocessor() monkeypatch.setattr( - module, "_preprocess_last_user_message", lambda message, state: None + module, + "_draft_last_user_message", + lambda message: module.DraftResult( + source="test", + result=module.NoDirective(reason="reject.confident_non_directive"), + ), ) first = { "model": "demo", @@ -183,10 +207,13 @@ def test_persistent_mode_with_drafter_preserves_existing_checkpoint_on_failure( module.CHECKPOINT_STORE.clear() hook = module.ContextCompilerPreCallHookWithPreprocessor() - def fake_preprocess(message: str, state: dict[str, object] | None) -> str | None: - return None + def fake_draft(message: str) -> object: + return module.DraftResult( + source="test", + result=module.NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module, "_preprocess_last_user_message", fake_preprocess) + monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) seed_data = { "model": "demo", "context_compiler_mode": "persistent", @@ -235,8 +262,15 @@ def test_normal_update_explicitly_saves_checkpoint(monkeypatch) -> None: hook = module.ContextCompilerPreCallHookWithPreprocessor() monkeypatch.setattr( module, - "_preprocess_last_user_message", - lambda message, state: "prohibit peanuts", + "_draft_last_user_message", + lambda message: module.DraftResult( + source="test", + result=CanonicalDirective( + kind="set_policy", + operands=("peanuts", "prohibit"), + text="prohibit peanuts", + ), + ), ) data = { "model": "demo", @@ -261,14 +295,16 @@ def test_restore_happens_before_drafting(monkeypatch) -> None: "chat-restore-first", {"premise": None, "policies": {"peanuts": "prohibit"}, "version": 2}, ) - seen_states: list[dict[str, object]] = [] + seen_messages: list[str] = [] - def fake_preprocess(message: str, state: dict[str, object] | None) -> str | None: - assert state is not None - seen_states.append(dict(state)) - return None + def fake_draft(message: str) -> object: + seen_messages.append(message) + return module.DraftResult( + source="test", + result=module.NoDirective(reason="reject.confident_non_directive"), + ) - monkeypatch.setattr(module, "_preprocess_last_user_message", fake_preprocess) + monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) data = { "model": "demo", "context_compiler_mode": "persistent", @@ -278,7 +314,7 @@ def fake_preprocess(message: str, state: dict[str, object] | None) -> str | None asyncio.run(hook.async_pre_call_hook(None, None, data, "completion")) - assert seen_states == [{"premise": None, "policies": {"peanuts": "prohibit"}}] + assert seen_messages == ["please use docker"] def test_corrupt_checkpoint_fails_clearly(monkeypatch) -> None: @@ -307,7 +343,16 @@ def test_forwarded_messages_keep_original_user_prompt_text(monkeypatch) -> None: {"role": "user", "content": "please use docker"}, ] monkeypatch.setattr( - module, "_preprocess_last_user_message", lambda message, state: "use docker" + module, + "_draft_last_user_message", + lambda message: module.DraftResult( + source="test", + result=CanonicalDirective( + kind="set_policy", + operands=("docker", "use"), + text="use docker", + ), + ), ) data = { "model": "demo", @@ -329,8 +374,11 @@ def test_compound_directives_fall_through_to_normal_forwarding_when_not_applied( hook = module.ContextCompilerPreCallHookWithPreprocessor() monkeypatch.setattr( module, - "_preprocess_last_user_message", - lambda _message, _state: "use docker and prohibit peanuts", + "_draft_last_user_message", + lambda _message: module.DraftResult( + source="test", + result=module.UnknownDirective(reason="reject.multi_candidate_directive"), + ), ) data = { "model": "demo", @@ -349,6 +397,55 @@ def test_compound_directives_fall_through_to_normal_forwarding_when_not_applied( assert checkpoint["policies"] == {} +def test_fallback_returns_structured_canonical_draft(monkeypatch) -> None: + module = _load_module(monkeypatch, "litellm_proxy_with_drafter_fallback_directive") + monkeypatch.setenv("OPENAI_API_KEY", "dummy") + monkeypatch.setenv("MODEL", "openai/demo-model") + monkeypatch.setattr( + module, + "_get_litellm_completion", + lambda: lambda **_: {"choices": [{"message": {"content": "use docker"}}]}, + ) + + result = module._llm_fallback_draft("please use docker") + + assert isinstance(result.result, CanonicalDirective) + assert result.result.text == "use docker" + assert result.source == "litellm_fallback" + + +def test_fallback_returns_structured_no_directive(monkeypatch) -> None: + module = _load_module(monkeypatch, "litellm_proxy_with_drafter_fallback_none") + monkeypatch.setenv("OPENAI_API_KEY", "dummy") + monkeypatch.setenv("MODEL", "openai/demo-model") + monkeypatch.setattr( + module, + "_get_litellm_completion", + lambda: lambda **_: {"choices": [{"message": {"content": ""}}]}, + ) + + result = module._llm_fallback_draft("hello there") + + assert isinstance(result.result, module.NoDirective) + assert result.source == "litellm_fallback" + + +def test_fallback_returns_structured_unknown_directive(monkeypatch) -> None: + module = _load_module(monkeypatch, "litellm_proxy_with_drafter_fallback_unknown") + monkeypatch.setenv("OPENAI_API_KEY", "dummy") + monkeypatch.setenv("MODEL", "openai/demo-model") + monkeypatch.setattr( + module, + "_get_litellm_completion", + lambda: lambda **_: {"choices": [{"message": {"content": "clear everything"}}]}, + ) + + result = module._llm_fallback_draft("clear everything") + + assert isinstance(result.result, module.UnknownDirective) + assert result.source == "litellm_fallback" + + def test_no_removed_replay_api_remains(monkeypatch) -> None: module = _load_module(monkeypatch, "litellm_proxy_with_drafter_no_replay") From aeb40bea6439cbcdef21af7b9870cacb85c0b43a Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 22:59:27 -0400 Subject: [PATCH 2/5] chore: bump drafter version to 0.2.0dev1 --- pyproject.toml | 6 +++--- uv.lock | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6655d76..bbf8557 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,12 +47,12 @@ Issues = "https://github.com/rlippmann/context-compiler-example-integrations/iss [project.optional-dependencies] all = [ "chromadb", - "context-compiler-directive-drafter==0.2.0dev0", + "context-compiler-directive-drafter==0.2.0dev1", "fastapi", "litellm", ] drafter = [ - "context-compiler-directive-drafter==0.2.0dev0", + "context-compiler-directive-drafter==0.2.0dev1", ] fastapi = [ "fastapi", @@ -67,7 +67,7 @@ retrieval = [ [dependency-groups] dev = [ "chromadb", - "context-compiler-directive-drafter==0.2.0dev0", + "context-compiler-directive-drafter==0.2.0dev1", "fastapi", "httpx2>=2.5.0", "httpx>=0.28.1", diff --git a/uv.lock b/uv.lock index f0059b6..16a0b54 100644 --- a/uv.lock +++ b/uv.lock @@ -667,14 +667,14 @@ wheels = [ [[package]] name = "context-compiler-directive-drafter" -version = "0.2.0.dev0" +version = "0.2.0.dev1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "context-compiler" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ed/4a/4f594ecab9c89e771b566cbccc854c81fba89081d9dd736925d73be31fa7/context_compiler_directive_drafter-0.2.0.dev0.tar.gz", hash = "sha256:14bb1096f34b0c2cc48ba3c23504dd33a305fc3dd203213b142d7b914f01649e", size = 86452, upload-time = "2026-08-08T04:52:30.798Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6b/b1/990d8ea99be4386ca6be2633446a904fdc759269c45e45079c9ba0f0c06c/context_compiler_directive_drafter-0.2.0.dev1.tar.gz", hash = "sha256:cae1bd43e145e96d895178154c0c261b8ed28037b38d5232f4b00ee57ae88e4f", size = 85661, upload-time = "2026-08-10T02:49:52.538Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/e5/b69ddec62e5206c580e4b9f3dba95da925f366181065f3d9d9c360a47685/context_compiler_directive_drafter-0.2.0.dev0-py3-none-any.whl", hash = "sha256:505f41356fedd5757ee7df1801f1e35436e96afd5bc0920f93a814e9598c6433", size = 20887, upload-time = "2026-08-08T04:52:29.707Z" }, + { url = "https://files.pythonhosted.org/packages/34/b6/b2982194d61a7a9affbc2afefa2d36eb1d1c49918664c9c1fad6c366ea31/context_compiler_directive_drafter-0.2.0.dev1-py3-none-any.whl", hash = "sha256:16b294e8ab63e08fc6e494dfa7777bea62c2f1db9b19473f4a656968b9848ca6", size = 18724, upload-time = "2026-08-10T02:49:51.109Z" }, ] [[package]] @@ -727,8 +727,8 @@ requires-dist = [ { name = "chromadb", marker = "extra == 'all'" }, { name = "chromadb", marker = "extra == 'retrieval'" }, { name = "context-compiler", specifier = "==0.9.0.dev7" }, - { name = "context-compiler-directive-drafter", marker = "extra == 'all'", specifier = "==0.2.0.dev0" }, - { name = "context-compiler-directive-drafter", marker = "extra == 'drafter'", specifier = "==0.2.0.dev0" }, + { name = "context-compiler-directive-drafter", marker = "extra == 'all'", specifier = "==0.2.0.dev1" }, + { name = "context-compiler-directive-drafter", marker = "extra == 'drafter'", specifier = "==0.2.0.dev1" }, { name = "fastapi", marker = "extra == 'all'" }, { name = "fastapi", marker = "extra == 'fastapi'" }, { name = "litellm", marker = "extra == 'all'" }, @@ -739,7 +739,7 @@ provides-extras = ["all", "drafter", "fastapi", "litellm", "retrieval"] [package.metadata.requires-dev] dev = [ { name = "chromadb" }, - { name = "context-compiler-directive-drafter", specifier = "==0.2.0.dev0" }, + { name = "context-compiler-directive-drafter", specifier = "==0.2.0.dev1" }, { name = "fastapi" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "httpx2", specifier = ">=2.5.0" }, From 5f6cf240d92fe9085b8257feec46feb836cd67d5 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 23:08:15 -0400 Subject: [PATCH 3/5] refactor: migrate LiteLLM integration to DirectiveDrafter API --- ...ler_precall_hook_with_directive_drafter.py | 65 +++----------- ...st_litellm_proxy_with_directive_drafter.py | 85 +++++++++++-------- 2 files changed, 59 insertions(+), 91 deletions(-) diff --git a/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py b/python/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py index d6761fb..1c3736e 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 @@ -33,14 +33,9 @@ class CustomLogger: # type: ignore[no-redef] ) from context_compiler.grammar import CanonicalDirective from context_compiler_directive_drafter import ( - DRAFT_OUTCOME_DIRECTIVE, - DRAFT_OUTCOME_NO_DIRECTIVE, DirectiveDrafter, DraftResult, - NoDirective, - UnknownDirective, - parse_preprocessor_output, - validate_preprocessor_output, + get_converter_prompt, ) from context_compiler_example_integrations.reference_integrations.litellm_proxy._checkpoint_support import ( MODE_PERSISTENT, @@ -67,15 +62,6 @@ class CustomLogger: # type: ignore[no-redef] } CHECKPOINT_STORE: CheckpointStore = InMemoryCheckpointStore() -_FALLBACK_SYSTEM_PROMPT = ( - "Convert the latest user message into exactly one valid Context Compiler " - "directive, or output . Use only these directive forms: " - "set premise , change premise to , use , prohibit " - ", remove policy , use instead of , " - "clear premise, reset policies, clear state. If the message is ambiguous, " - "not a direct instruction to change compiler state, or could imply more " - "than one instruction, output . Do not explain." -) def _extract_response_content(response: object) -> str | None: @@ -106,35 +92,26 @@ def _get_litellm_completion() -> Callable[..., object]: return cast(Callable[..., object], litellm_module.completion) -def _llm_fallback_draft(message: str) -> DraftResult: +def _llm_fallback_candidate(message: str) -> str | None: preprocessor_model = os.getenv("PREPROCESSOR_MODEL", "").strip() if not preprocessor_model: preprocessor_model = os.getenv("MODEL", "").strip() if not preprocessor_model: - return DraftResult( - source="litellm_fallback", - result=UnknownDirective(reason="fallback_model_unconfigured"), - ) + return None api_key = os.getenv("OPENAI_API_KEY") if not api_key: - return DraftResult( - source="litellm_fallback", - result=UnknownDirective(reason="fallback_api_key_missing"), - ) + return None try: completion = _get_litellm_completion() except ModuleNotFoundError: - return DraftResult( - source="litellm_fallback", - result=UnknownDirective(reason="fallback_litellm_unavailable"), - ) + return None kwargs: dict[str, object] = { "model": preprocessor_model, "messages": [ - {"role": "system", "content": _FALLBACK_SYSTEM_PROMPT}, + {"role": "system", "content": get_converter_prompt()}, {"role": "user", "content": message}, ], "api_key": api_key, @@ -146,35 +123,15 @@ def _llm_fallback_draft(message: str) -> DraftResult: try: response = completion(**kwargs) - raw_output = _extract_response_content(response) + return _extract_response_content(response) except Exception: - return DraftResult( - source="litellm_fallback", - result=UnknownDirective(reason="fallback_completion_failed"), - ) - - validated = validate_preprocessor_output(raw_output) - if validated["classification"] == DRAFT_OUTCOME_DIRECTIVE: - parsed = parse_preprocessor_output(raw_output) - if parsed is not None: - return DraftResult(source="litellm_fallback", result=parsed) - return DraftResult( - source="litellm_fallback", - result=UnknownDirective(reason="invalid_canonical_directive"), - ) - if validated["classification"] == DRAFT_OUTCOME_NO_DIRECTIVE: - return DraftResult( - source="litellm_fallback", - result=NoDirective(reason="fallback_confident_non_directive"), - ) - return DraftResult( - source="litellm_fallback", - result=UnknownDirective(reason="fallback_unresolved"), - ) + return None def _draft_last_user_message(message: str) -> DraftResult: - drafter = DirectiveDrafter(fallback=_llm_fallback_draft) + drafter = DirectiveDrafter( + fallback=_llm_fallback_candidate, fallback_source="litellm_fallback" + ) return drafter.draft_directive(message) diff --git a/python/tests/test_litellm_proxy_with_directive_drafter.py b/python/tests/test_litellm_proxy_with_directive_drafter.py index 5a9989c..c4f3803 100644 --- a/python/tests/test_litellm_proxy_with_directive_drafter.py +++ b/python/tests/test_litellm_proxy_with_directive_drafter.py @@ -4,9 +4,11 @@ import types from copy import deepcopy from pathlib import Path +from typing import Any import pytest -from context_compiler.grammar import CanonicalDirective +from context_compiler.grammar import decompose_directive +from context_compiler_directive_drafter import NoDirective, UnknownDirective REPO_ROOT = Path(__file__).resolve().parents[2] MODULE_PATH = ( @@ -26,7 +28,7 @@ def _load_module(monkeypatch: pytest.MonkeyPatch, module_name: str): class _CustomLogger: pass - custom_logger_mod.CustomLogger = _CustomLogger + custom_logger_mod.CustomLogger = _CustomLogger # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, "litellm", litellm_mod) monkeypatch.setitem(sys.modules, "litellm.integrations", integrations_mod) monkeypatch.setitem( @@ -49,7 +51,7 @@ def fake_draft(message: str) -> object: drafted_calls.append((message, {})) return module.DraftResult( source="test", - result=module.NoDirective(reason="reject.confident_non_directive"), + result=NoDirective(reason="reject.confident_non_directive"), ) monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) @@ -72,16 +74,14 @@ def fake_draft(message: str) -> object: def test_drafter_output_applies_to_current_turn_only(monkeypatch) -> None: module = _load_module(monkeypatch, "litellm_proxy_with_drafter_applies") hook = module.ContextCompilerPreCallHookWithPreprocessor() + directive = decompose_directive("prohibit docker") + assert directive is not None monkeypatch.setattr( module, "_draft_last_user_message", lambda message: module.DraftResult( source="test", - result=CanonicalDirective( - kind="set_policy", - operands=("docker", "prohibit"), - text="prohibit docker", - ), + result=directive, ), ) data = { @@ -112,7 +112,7 @@ def fake_draft(message: str) -> object: drafted_inputs.append(message) return module.DraftResult( source="test", - result=module.NoDirective(reason="reject.confident_non_directive"), + result=NoDirective(reason="reject.confident_non_directive"), ) monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) @@ -156,7 +156,7 @@ def test_default_mode_is_stateless_and_requires_no_session_key(monkeypatch) -> N "_draft_last_user_message", lambda message: module.DraftResult( source="test", - result=module.NoDirective(reason="reject.confident_non_directive"), + result=NoDirective(reason="reject.confident_non_directive"), ), ) data = { @@ -177,7 +177,7 @@ def test_stateless_mode_has_no_cross_call_continuity(monkeypatch) -> None: "_draft_last_user_message", lambda message: module.DraftResult( source="test", - result=module.NoDirective(reason="reject.confident_non_directive"), + result=NoDirective(reason="reject.confident_non_directive"), ), ) first = { @@ -210,7 +210,7 @@ def test_persistent_mode_with_drafter_preserves_existing_checkpoint_on_failure( def fake_draft(message: str) -> object: return module.DraftResult( source="test", - result=module.NoDirective(reason="reject.confident_non_directive"), + result=NoDirective(reason="reject.confident_non_directive"), ) monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) @@ -260,16 +260,14 @@ def test_normal_update_explicitly_saves_checkpoint(monkeypatch) -> None: module = _load_module(monkeypatch, "litellm_proxy_with_drafter_save_after_update") module.CHECKPOINT_STORE.clear() hook = module.ContextCompilerPreCallHookWithPreprocessor() + directive = decompose_directive("prohibit peanuts") + assert directive is not None monkeypatch.setattr( module, "_draft_last_user_message", lambda message: module.DraftResult( source="test", - result=CanonicalDirective( - kind="set_policy", - operands=("peanuts", "prohibit"), - text="prohibit peanuts", - ), + result=directive, ), ) data = { @@ -301,7 +299,7 @@ def fake_draft(message: str) -> object: seen_messages.append(message) return module.DraftResult( source="test", - result=module.NoDirective(reason="reject.confident_non_directive"), + result=NoDirective(reason="reject.confident_non_directive"), ) monkeypatch.setattr(module, "_draft_last_user_message", fake_draft) @@ -338,6 +336,8 @@ def test_corrupt_checkpoint_fails_clearly(monkeypatch) -> None: def test_forwarded_messages_keep_original_user_prompt_text(monkeypatch) -> None: module = _load_module(monkeypatch, "litellm_proxy_with_drafter_forwarded_text") hook = module.ContextCompilerPreCallHookWithPreprocessor() + directive = decompose_directive("use docker") + assert directive is not None original_messages = [ {"role": "system", "content": "original system"}, {"role": "user", "content": "please use docker"}, @@ -347,11 +347,7 @@ def test_forwarded_messages_keep_original_user_prompt_text(monkeypatch) -> None: "_draft_last_user_message", lambda message: module.DraftResult( source="test", - result=CanonicalDirective( - kind="set_policy", - operands=("docker", "use"), - text="use docker", - ), + result=directive, ), ) data = { @@ -377,7 +373,7 @@ def test_compound_directives_fall_through_to_normal_forwarding_when_not_applied( "_draft_last_user_message", lambda _message: module.DraftResult( source="test", - result=module.UnknownDirective(reason="reject.multi_candidate_directive"), + result=UnknownDirective(reason="reject.multi_candidate_directive"), ), ) data = { @@ -397,7 +393,7 @@ def test_compound_directives_fall_through_to_normal_forwarding_when_not_applied( assert checkpoint["policies"] == {} -def test_fallback_returns_structured_canonical_draft(monkeypatch) -> None: +def test_fallback_returns_raw_candidate_text(monkeypatch) -> None: module = _load_module(monkeypatch, "litellm_proxy_with_drafter_fallback_directive") monkeypatch.setenv("OPENAI_API_KEY", "dummy") monkeypatch.setenv("MODEL", "openai/demo-model") @@ -407,14 +403,12 @@ def test_fallback_returns_structured_canonical_draft(monkeypatch) -> None: lambda: lambda **_: {"choices": [{"message": {"content": "use docker"}}]}, ) - result = module._llm_fallback_draft("please use docker") + result = module._llm_fallback_candidate("please use docker") - assert isinstance(result.result, CanonicalDirective) - assert result.result.text == "use docker" - assert result.source == "litellm_fallback" + assert result == "use docker" -def test_fallback_returns_structured_no_directive(monkeypatch) -> None: +def test_fallback_returns_raw_no_directive_sentinel(monkeypatch) -> None: module = _load_module(monkeypatch, "litellm_proxy_with_drafter_fallback_none") monkeypatch.setenv("OPENAI_API_KEY", "dummy") monkeypatch.setenv("MODEL", "openai/demo-model") @@ -424,13 +418,12 @@ def test_fallback_returns_structured_no_directive(monkeypatch) -> None: lambda: lambda **_: {"choices": [{"message": {"content": ""}}]}, ) - result = module._llm_fallback_draft("hello there") + result = module._llm_fallback_candidate("hello there") - assert isinstance(result.result, module.NoDirective) - assert result.source == "litellm_fallback" + assert result == "" -def test_fallback_returns_structured_unknown_directive(monkeypatch) -> None: +def test_fallback_returns_raw_unknown_candidate_text(monkeypatch) -> None: module = _load_module(monkeypatch, "litellm_proxy_with_drafter_fallback_unknown") monkeypatch.setenv("OPENAI_API_KEY", "dummy") monkeypatch.setenv("MODEL", "openai/demo-model") @@ -440,10 +433,28 @@ def test_fallback_returns_structured_unknown_directive(monkeypatch) -> None: lambda: lambda **_: {"choices": [{"message": {"content": "clear everything"}}]}, ) - result = module._llm_fallback_draft("clear everything") + result = module._llm_fallback_candidate("clear everything") - assert isinstance(result.result, module.UnknownDirective) - assert result.source == "litellm_fallback" + assert result == "clear everything" + + +def test_fallback_uses_shared_converter_prompt(monkeypatch) -> None: + module = _load_module(monkeypatch, "litellm_proxy_with_drafter_shared_prompt") + monkeypatch.setenv("OPENAI_API_KEY", "dummy") + monkeypatch.setenv("MODEL", "openai/demo-model") + seen: dict[str, Any] = {} + + def completion(**kwargs): + seen.update(kwargs) + return {"choices": [{"message": {"content": "use docker"}}]} + + monkeypatch.setattr(module, "_get_litellm_completion", lambda: completion) + monkeypatch.setattr(module, "get_converter_prompt", lambda: "shared prompt") + + result = module._llm_fallback_candidate("please use docker") + + assert result == "use docker" + assert seen["messages"][0] == {"role": "system", "content": "shared prompt"} def test_no_removed_replay_api_remains(monkeypatch) -> None: From fd522fd0e81685cec6eee3ef24994b69c928f0c7 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 23:12:53 -0400 Subject: [PATCH 4/5] fix: migrate OpenWebUI drafter prompt API --- .../open_webui_pipe_with_directive_drafter.py | 21 +++---------------- 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py b/python/reference_integrations/openwebui_pipe/open_webui_pipe_with_directive_drafter.py index 2f73a9b..81df95c 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 @@ -23,8 +23,6 @@ import logging import re from collections.abc import AsyncIterator -from importlib.resources import as_file, files -from importlib.resources.abc import Traversable from typing import Any, Literal, TypedDict, cast from fastapi import Request # type: ignore[import-not-found] @@ -60,16 +58,15 @@ def Field(*, default: Any, description: str = "") -> Any: # type: ignore[no-red from context_compiler.engine import Engine from context_compiler_directive_drafter import ( DRAFT_OUTCOME_DIRECTIVE, + get_converter_prompt, parse_preprocessor_output, preprocess_heuristic, - render_prompt, ) logger = logging.getLogger(__name__) _CC_MARKER = "[[cc_state]]" _ENGINES_BY_CHAT_KEY: dict[str, Engine] = {} -_PROMPTS_DIR = files("context_compiler_directive_drafter").joinpath("prompts") class _EngineSnapshot(TypedDict): @@ -90,15 +87,6 @@ def _is_directive_shaped_input(message: str) -> bool: ) -def _prompt_file_path(profile: str) -> Traversable: - # Runtime prompt selection for fallback drafting: - # - default: most instruction-following models - # - llama: models that need tighter prompt guidance - if profile == "llama": - return _PROMPTS_DIR.joinpath("llama.txt") - return _PROMPTS_DIR.joinpath("default.txt") - - def _resolve_chat_key( user: dict[str, Any], chat_id: str | None, @@ -708,19 +696,16 @@ async def _llm_fallback_preprocess( prompt_profile: str, model_id: str | None, ) -> tuple[str | None, str | None]: + del state, prompt_profile model_id = _normalize_model_id(model_id) if model_id is None: return None, None - with as_file(_prompt_file_path(prompt_profile)) as prompt_path: - prompt = render_prompt(prompt_path, state["premise"], state["policies"]) - if prompt is None: - return None, None payload: dict[str, Any] = { "model": model_id, "stream": False, "messages": [ - {"role": "system", "content": prompt}, + {"role": "system", "content": get_converter_prompt()}, {"role": "user", "content": message}, ], } From 9d9c6f63020c2dfe4469034d9c233549df614bc3 Mon Sep 17 00:00:00 2001 From: Robert Lippmann Date: Sun, 9 Aug 2026 23:22:11 -0400 Subject: [PATCH 5/5] fix: migrate LiteLLM example prompt API --- .../litellm/with_directive_drafter.py | 20 ++----------- .../test_litellm_with_directive_drafter.py | 28 ++++++++++--------- 2 files changed, 18 insertions(+), 30 deletions(-) diff --git a/python/examples/prompt_construction/litellm/with_directive_drafter.py b/python/examples/prompt_construction/litellm/with_directive_drafter.py index 109b73f..1c3393c 100644 --- a/python/examples/prompt_construction/litellm/with_directive_drafter.py +++ b/python/examples/prompt_construction/litellm/with_directive_drafter.py @@ -20,8 +20,6 @@ import re from collections.abc import Callable, Mapping, Sequence from importlib import import_module -from importlib.resources import as_file, files -from importlib.resources.abc import Traversable from typing import TypedDict, cast from context_compiler import ( @@ -35,9 +33,9 @@ from context_compiler.engine import Engine from context_compiler_directive_drafter import ( DRAFT_OUTCOME_DIRECTIVE, + get_converter_prompt, parse_preprocessor_output, preprocess_heuristic, - render_prompt, ) try: @@ -55,8 +53,6 @@ ) logger = logging.getLogger(__name__) - -_PROMPTS_DIR = files("context_compiler_directive_drafter").joinpath("prompts") SHOW_CONTEXT_COMPILER_TRACE = False @@ -236,18 +232,8 @@ def _call_litellm(messages: list[dict[str, str]]) -> str: return content -def _prompt_file_path() -> Traversable: - profile = os.getenv("PREPROCESSOR_PROMPT_PROFILE", "default").strip().lower() - if profile == "llama": - return _PROMPTS_DIR.joinpath("llama.txt") - return _PROMPTS_DIR.joinpath("default.txt") - - def _llm_fallback_preprocess(message: str, state: _EngineSnapshot) -> str | None: - with as_file(_prompt_file_path()) as prompt_path: - prompt = render_prompt(prompt_path, state["premise"], state["policies"]) - if prompt is None: - return None + del state try: completion = _get_litellm_completion() @@ -267,7 +253,7 @@ def _llm_fallback_preprocess(message: str, state: _EngineSnapshot) -> str | None kwargs: _LiteLLMCallKwargs = { "model": preprocessor_model, "messages": [ - {"role": "system", "content": prompt}, + {"role": "system", "content": get_converter_prompt()}, {"role": "user", "content": message}, ], "temperature": 0, diff --git a/python/tests/test_litellm_with_directive_drafter.py b/python/tests/test_litellm_with_directive_drafter.py index fffa331..8b0ecb3 100644 --- a/python/tests/test_litellm_with_directive_drafter.py +++ b/python/tests/test_litellm_with_directive_drafter.py @@ -82,11 +82,12 @@ def step_with_capture(user_input: str): lambda message: {"outcome": "no_directive", "directive": None}, ) monkeypatch.setattr(module, "_llm_fallback_preprocess", lambda message, state: None) - monkeypatch.setattr( - module, - "_call_litellm", - lambda messages: llm_calls.append(messages) or "stubbed reply", - ) + + def downstream(messages: list[dict[str, str]]) -> str: + llm_calls.append(messages) + return "stubbed reply" + + monkeypatch.setattr(module, "_call_litellm", downstream) result = module.handle_turn("hello there", engine) @@ -99,11 +100,12 @@ def test_local_update_and_clarify_responses_skip_downstream_litellm_call( monkeypatch, ) -> None: llm_calls: list[object] = [] - monkeypatch.setattr( - module, - "_call_litellm", - lambda messages: llm_calls.append(messages) or "should not be called", - ) + + def should_not_call(messages: list[dict[str, str]]) -> str: + llm_calls.append(messages) + return "should not be called" + + monkeypatch.setattr(module, "_call_litellm", should_not_call) monkeypatch.setattr( module, "preprocess_heuristic", @@ -232,7 +234,7 @@ def completion(**kwargs: Any) -> dict[str, object]: monkeypatch.setenv("MODEL", "openai/main-model") monkeypatch.delenv("PREPROCESSOR_MODEL", raising=False) monkeypatch.setattr(module, "_get_litellm_completion", lambda: completion) - monkeypatch.setattr(module, "render_prompt", lambda *_: "prompt") + monkeypatch.setattr(module, "get_converter_prompt", lambda: "prompt") monkeypatch.setattr( module, "parse_preprocessor_output", @@ -259,7 +261,7 @@ def completion(**kwargs: Any) -> dict[str, object]: monkeypatch.setenv("MODEL", "openai/main-model") monkeypatch.setenv("PREPROCESSOR_MODEL", "openai/preprocessor-model") monkeypatch.setattr(module, "_get_litellm_completion", lambda: completion) - monkeypatch.setattr(module, "render_prompt", lambda *_: "prompt") + monkeypatch.setattr(module, "get_converter_prompt", lambda: "prompt") monkeypatch.setattr( module, "parse_preprocessor_output", @@ -290,7 +292,7 @@ def test_fallback_accepts_structurally_valid_output_without_source_awareness( } ), ) - monkeypatch.setattr(module, "render_prompt", lambda *_: "prompt") + monkeypatch.setattr(module, "get_converter_prompt", lambda: "prompt") assert ( module._llm_fallback_preprocess(