context-compiler-example-integrations 0.1.0__py3-none-any.whl

context_compiler_example_integrations/reference_integrations/litellm_proxy/_checkpoint_support.py

@@ -0,0 +1,146 @@
+"""Narrow checkpoint/session support for LiteLLM Proxy reference hooks.
+
+This module is intentionally repo-local and small in scope:
+- explicit persistent vs stateless mode selection
+- explicit session-key resolution
+- storage-neutral checkpoint store contract
+- latest-user-turn extraction helpers
+
+The in-memory store is suitable only for tests and single-process examples.
+It does not provide durability, multi-worker coordination, or atomic updates.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Protocol
+
+MODE_PERSISTENT = "persistent"
+MODE_STATELESS = "stateless"
+SESSION_MODE_ENV_VAR = "CONTEXT_COMPILER_SESSION_MODE"
+
+
+class CheckpointStore(Protocol):
+    def load(self, session_key: str) -> Mapping[str, object] | None: ...
+
+    def save(self, session_key: str, checkpoint: Mapping[str, object]) -> None: ...
+
+
+class InMemoryCheckpointStore:
+    """Single-process example checkpoint store for tests and local reference use."""
+
+    def __init__(self) -> None:
+        self._checkpoints: dict[str, dict[str, object]] = {}
+
+    def load(self, session_key: str) -> Mapping[str, object] | None:
+        checkpoint = self._checkpoints.get(session_key)
+        if checkpoint is None:
+            return None
+        return dict(checkpoint)
+
+    def save(self, session_key: str, checkpoint: Mapping[str, object]) -> None:
+        self._checkpoints[session_key] = dict(checkpoint)
+
+    def clear(self) -> None:
+        self._checkpoints.clear()
+
+
+@dataclass(frozen=True)
+class SessionContext:
+    mode: str
+    session_key: str | None
+    source: str | None
+
+
+def resolve_session_context(data: Mapping[str, object]) -> SessionContext:
+    mode = _resolve_mode(data)
+    session_key, source = _resolve_session_key(data)
+    return SessionContext(mode=mode, session_key=session_key, source=source)
+
+
+def extract_latest_user_text(messages: list[dict[str, object]]) -> str | None:
+    for message in reversed(messages):
+        if message.get("role") != "user":
+            continue
+        content = extract_text_content(message.get("content"))
+        if content is not None:
+            return content
+        return None
+    return None
+
+
+def extract_text_content(content: object) -> str | None:
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        text_parts: list[str] = []
+        for item in content:
+            if not isinstance(item, dict):
+                continue
+            if item.get("type") != "text":
+                continue
+            text = item.get("text")
+            if isinstance(text, str):
+                text_parts.append(text)
+        if text_parts:
+            return " ".join(text_parts)
+    return None
+
+
+def checkpoint_to_jsonable(checkpoint_json: str) -> dict[str, object]:
+    import json
+
+    raw = json.loads(checkpoint_json)
+    if not isinstance(raw, dict):
+        raise ValueError("Checkpoint JSON must decode to an object.")
+    return raw
+
+
+def checkpoint_from_jsonable(checkpoint: Mapping[str, object]) -> str:
+    import json
+
+    return json.dumps(checkpoint, separators=(",", ":"), sort_keys=True)
+
+
+def _resolve_mode(data: Mapping[str, object]) -> str:
+    candidates = [
+        data.get("context_compiler_mode"),
+        os.getenv(SESSION_MODE_ENV_VAR),
+    ]
+    for candidate in candidates:
+        if not isinstance(candidate, str):
+            continue
+        normalized = candidate.strip().lower()
+        if normalized in {MODE_PERSISTENT, MODE_STATELESS}:
+            return normalized
+    return MODE_STATELESS
+
+
+def _resolve_session_key(data: Mapping[str, object]) -> tuple[str | None, str | None]:
+    candidates: list[tuple[str, object]] = [
+        ("context_compiler_session_key", data.get("context_compiler_session_key")),
+        (
+            "metadata.context_compiler_session_key",
+            _nested_lookup(data.get("metadata"), "context_compiler_session_key"),
+        ),
+    ]
+    for source, value in candidates:
+        normalized = _normalize_key(value)
+        if normalized is not None:
+            return normalized, source
+    return None, None
+
+
+def _nested_lookup(value: object, key: str) -> object:
+    if isinstance(value, Mapping):
+        return value.get(key)
+    return None
+
+
+def _normalize_key(value: object) -> str | None:
+    if not isinstance(value, str):
+        return None
+    normalized = value.strip()
+    return normalized or None

context_compiler_example_integrations/reference_integrations/litellm_proxy/config.example.yaml

@@ -0,0 +1,23 @@
+model_list:
+  # `model_name` is the client-facing alias sent to the proxy.
+  # `litellm_params.model` is the upstream provider/model LiteLLM calls.
+  - model_name: gpt-4o-mini
+    litellm_params:
+      model: openai/gpt-4o-mini
+      api_key: os.environ/OPENAI_API_KEY
+
+  - model_name: llama3.1
+    litellm_params:
+      model: ollama/llama3.1:8b
+      api_base: http://localhost:11434
+
+litellm_settings:
+  callbacks:
+    # Basic checkpoint-backed hook.
+    # Default mode is stateless unless the request or environment sets
+    # CONTEXT_COMPILER_SESSION_MODE=persistent.
+    - context_compiler_precall_hook.proxy_handler_instance
+    # Directive-drafter-enabled checkpoint-backed hook (use this instead).
+    # Persistent mode requires context_compiler_session_key or
+    # metadata.context_compiler_session_key on the request.
+    # - context_compiler_precall_hook_with_directive_drafter.proxy_handler_instance

context_compiler_example_integrations/reference_integrations/litellm_proxy/context_compiler_precall_hook.py

@@ -0,0 +1,163 @@
+"""Minimal LiteLLM Proxy pre-call hook example.
+
+Architecture:
+- Resolve explicit persistent or stateless mode for the current request.
+- In persistent mode, restore compiler checkpoint by session key.
+- Process only the latest user turn exactly once.
+- Save checkpoint after each decision, including clarify.
+- If clarification is required, block upstream model call.
+- Otherwise inject compiled state guidance into a system message.
+"""
+
+import logging
+from typing import Any
+
+try:
+    from litellm.integrations.custom_logger import CustomLogger
+except ModuleNotFoundError:
+    # Keep this import path optional: CI/tests run without integration extras.
+    # A tiny fallback base class keeps module imports deterministic so coverage
+    # validates behavior instead of failing or silently skipping on missing litellm.
+    class CustomLogger:  # type: ignore[no-redef]
+        pass
+
+
+from context_compiler import (
+    POLICY_PROHIBIT,
+    create_engine,
+    get_clarify_prompt,
+    State,
+    get_policy_items,
+    get_premise_value,
+    is_clarify,
+)
+from context_compiler.engine import DecisionKind
+from context_compiler_example_integrations.reference_integrations.litellm_proxy._checkpoint_support import (
+    MODE_PERSISTENT,
+    CheckpointStore,
+    InMemoryCheckpointStore,
+    checkpoint_from_jsonable,
+    checkpoint_to_jsonable,
+    extract_latest_user_text,
+    resolve_session_context,
+)
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_CALL_TYPES = {
+    "completion",
+    "acompletion",
+    "chat_completion",
+    "achat_completion",
+}
+CHECKPOINT_STORE: CheckpointStore = InMemoryCheckpointStore()
+
+
+def _render_compiled_state_contract(compiled_state: State) -> str:
+    prohibited = get_policy_items(compiled_state, POLICY_PROHIBIT)
+    premise = get_premise_value(compiled_state)
+
+    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:
+        lines.append(
+            "When the answer depends on user preference/style, "
+            f"treat the current premise as: {premise}."
+        )
+    lines.append(
+        "If the user message conflicts with these constraints, follow them exactly."
+    )
+
+    return "Host policy contract:\n" + "\n".join(f"- {line}" for line in lines)
+
+
+def _extract_request_messages(data: dict[str, object]) -> list[dict[str, object]]:
+    raw_messages = data.get("messages")
+    if not isinstance(raw_messages, list):
+        return []
+    return [msg for msg in raw_messages if isinstance(msg, dict)]
+
+
+class ContextCompilerPreCallHook(CustomLogger):
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: Any,
+        cache: Any,
+        data: dict[str, object],
+        call_type: str,
+    ) -> dict[str, object] | str:
+        del user_api_key_dict, cache
+        logger.debug("litellm_proxy: call_type=%s", call_type)
+        if call_type not in _SUPPORTED_CALL_TYPES:
+            return data
+
+        request_messages = _extract_request_messages(data)
+        logger.debug("litellm_proxy: message_count=%d", len(request_messages))
+        session = resolve_session_context(data)
+        logger.debug(
+            "litellm_proxy: mode=%s session_key_source=%s",
+            session.mode,
+            session.source,
+        )
+        if session.mode == MODE_PERSISTENT and session.session_key is None:
+            return (
+                "Context Compiler persistent mode requires a stable session key. "
+                "Set context_compiler_session_key or "
+                "metadata.context_compiler_session_key."
+            )
+
+        latest_user_text = extract_latest_user_text(request_messages)
+        logger.debug(
+            "litellm_proxy: latest_user_text_present=%s", latest_user_text is not None
+        )
+
+        engine = create_engine()
+        if session.mode == MODE_PERSISTENT and session.session_key is not None:
+            checkpoint = CHECKPOINT_STORE.load(session.session_key)
+            if checkpoint is not None:
+                try:
+                    engine.import_checkpoint_json(checkpoint_from_jsonable(checkpoint))
+                except Exception as exc:
+                    return (
+                        "Context Compiler checkpoint load failed for session "
+                        f"{session.session_key!r}: {exc}"
+                    )
+
+        if latest_user_text is not None:
+            decision = engine.step(latest_user_text)
+        else:
+            decision = {
+                "kind": DecisionKind.PASSTHROUGH,
+                "state": engine.state,
+                "prompt_to_user": None,
+            }
+
+        if session.mode == MODE_PERSISTENT and session.session_key is not None:
+            CHECKPOINT_STORE.save(
+                session.session_key,
+                checkpoint_to_jsonable(engine.export_checkpoint_json()),
+            )
+
+        logger.debug("litellm_proxy: decision_kind=%s", decision["kind"])
+
+        if is_clarify(decision):
+            logger.debug("litellm_proxy: blocking_on_clarify=true")
+            return get_clarify_prompt(decision) or "Confirmation required."
+
+        compiled_state = engine.state
+        # 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),
+        }
+        # Prepend one compiler contract system message, then forward the original
+        # request messages unchanged. Existing system messages are preserved.
+        logger.debug("litellm_proxy: inject_system_message=true")
+        data["messages"] = [system_message, *request_messages]
+        return data
+
+
+proxy_handler_instance = ContextCompilerPreCallHook()

context_compiler_example_integrations/reference_integrations/litellm_proxy/context_compiler_precall_hook_with_directive_drafter.py

@@ -0,0 +1,289 @@
+"""LiteLLM Proxy pre-call hook with optional directive drafter on latest user message.
+
+Architecture:
+- Resolve explicit persistent or stateless mode for the current request.
+- In persistent mode, restore compiler checkpoint by session key.
+- Draft only the latest user message after restore.
+- Call ``engine.step(...)`` exactly once for the current turn.
+- Save checkpoint after each decision, including clarify.
+- If clarification is required, block upstream model call.
+- Otherwise inject compiled state guidance into a system message.
+"""
+
+import logging
+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:
+    from litellm.integrations.custom_logger import CustomLogger
+except ModuleNotFoundError:
+    # Keep this import path optional: CI/tests run without integration extras.
+    # A tiny fallback base class keeps module imports deterministic so coverage
+    # validates behavior instead of failing or silently skipping on missing litellm.
+    class CustomLogger:  # type: ignore[no-redef]
+        pass
+
+
+from context_compiler import (
+    POLICY_PROHIBIT,
+    State,
+    create_engine,
+    get_clarify_prompt,
+    get_policy_items,
+    get_premise_value,
+    is_clarify,
+)
+from context_compiler.engine import DecisionKind
+from context_compiler_directive_drafter import (
+    PREPROCESS_OUTCOME_DIRECTIVE,
+    parse_preprocessor_output,
+    preprocess_heuristic,
+    render_prompt,
+)
+from context_compiler_example_integrations.reference_integrations.litellm_proxy._checkpoint_support import (
+    MODE_PERSISTENT,
+    CheckpointStore,
+    InMemoryCheckpointStore,
+    checkpoint_from_jsonable,
+    checkpoint_to_jsonable,
+    extract_latest_user_text,
+    resolve_session_context,
+)
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_CALL_TYPES = {
+    "completion",
+    "acompletion",
+    "chat_completion",
+    "achat_completion",
+}
+
+_PROMPTS_DIR = files("context_compiler_directive_drafter").joinpath("prompts")
+CHECKPOINT_STORE: CheckpointStore = InMemoryCheckpointStore()
+
+
+def _render_compiled_state_contract(compiled_state: State) -> str:
+    prohibited = get_policy_items(compiled_state, POLICY_PROHIBIT)
+    premise = get_premise_value(compiled_state)
+
+    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:
+        lines.append(
+            "When the answer depends on user preference/style, "
+            f"treat the current premise as: {premise}."
+        )
+    lines.append(
+        "If the user message conflicts with these constraints, follow them exactly."
+    )
+
+    return "Host policy contract:\n" + "\n".join(f"- {line}" for line in lines)
+
+
+def _extract_request_messages(data: dict[str, object]) -> list[dict[str, object]]:
+    raw_messages = data.get("messages")
+    if not isinstance(raw_messages, list):
+        return []
+    return [msg for msg in raw_messages if isinstance(msg, dict)]
+
+
+def _extract_response_content(response: object) -> str | None:
+    if isinstance(response, Mapping):
+        choices = response.get("choices")
+        if isinstance(choices, Sequence) and choices:
+            first = choices[0]
+            if isinstance(first, Mapping):
+                message = first.get("message")
+                if isinstance(message, Mapping):
+                    content = message.get("content")
+                    if isinstance(content, str):
+                        return content
+
+    choices_attr = getattr(response, "choices", None)
+    if isinstance(choices_attr, Sequence) and choices_attr:
+        first = choices_attr[0]
+        message_attr = getattr(first, "message", None)
+        content_attr = getattr(message_attr, "content", None)
+        if isinstance(content_attr, str):
+            return content_attr
+
+    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: State) -> str | None:
+    with as_file(_prompt_file_path()) as prompt_path:
+        prompt = render_prompt(prompt_path, state)
+    if prompt is None:
+        return None
+
+    preprocessor_model = os.getenv("PREPROCESSOR_MODEL", "").strip()
+    if not preprocessor_model:
+        preprocessor_model = os.getenv("MODEL", "").strip()
+    if not preprocessor_model:
+        return None
+
+    api_key = os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        return None
+
+    try:
+        completion = _get_litellm_completion()
+    except ModuleNotFoundError:
+        return None
+
+    kwargs: dict[str, object] = {
+        "model": preprocessor_model,
+        "messages": [
+            {"role": "system", "content": prompt},
+            {"role": "user", "content": message},
+        ],
+        "api_key": api_key,
+        "temperature": 0,
+    }
+    api_base = os.getenv("OPENAI_BASE_URL")
+    if api_base:
+        kwargs["api_base"] = api_base
+
+    try:
+        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
+
+
+def _preprocess_last_user_message(message: str, state: State | None) -> str | None:
+    try:
+        heuristic_result = preprocess_heuristic(message)
+        if (
+            heuristic_result["outcome"] == PREPROCESS_OUTCOME_DIRECTIVE
+            and heuristic_result["directive"]
+        ):
+            parsed = parse_preprocessor_output(heuristic_result["directive"])
+            if parsed is not None:
+                return parsed
+    except Exception:
+        logger.debug("litellm_proxy: heuristic_exception", exc_info=True)
+
+    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
+
+
+class ContextCompilerPreCallHookWithPreprocessor(CustomLogger):
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: Any,
+        cache: Any,
+        data: dict[str, object],
+        call_type: str,
+    ) -> dict[str, object] | str:
+        del user_api_key_dict, cache
+        logger.debug("litellm_proxy: call_type=%s", call_type)
+        if call_type not in _SUPPORTED_CALL_TYPES:
+            return data
+
+        request_messages = _extract_request_messages(data)
+        logger.debug("litellm_proxy: message_count=%d", len(request_messages))
+        session = resolve_session_context(data)
+        logger.debug(
+            "litellm_proxy: mode=%s session_key_source=%s",
+            session.mode,
+            session.source,
+        )
+        if session.mode == MODE_PERSISTENT and session.session_key is None:
+            return (
+                "Context Compiler persistent mode requires a stable session key. "
+                "Set context_compiler_session_key or "
+                "metadata.context_compiler_session_key."
+            )
+
+        engine = create_engine()
+        if session.mode == MODE_PERSISTENT and session.session_key is not None:
+            checkpoint = CHECKPOINT_STORE.load(session.session_key)
+            if checkpoint is not None:
+                try:
+                    engine.import_checkpoint_json(checkpoint_from_jsonable(checkpoint))
+                except Exception as exc:
+                    return (
+                        "Context Compiler checkpoint load failed for session "
+                        f"{session.session_key!r}: {exc}"
+                    )
+
+        latest_user_text = extract_latest_user_text(request_messages)
+        logger.debug(
+            "litellm_proxy: latest_user_text_present=%s", latest_user_text is not None
+        )
+        engine_input = latest_user_text
+        drafted_input: str | None = None
+
+        if latest_user_text is not None and not engine.has_pending_clarification():
+            drafted_input = _preprocess_last_user_message(
+                latest_user_text, engine.state
+            )
+            logger.debug("litellm_proxy: drafted_input=%r", drafted_input)
+            if drafted_input is not None:
+                engine_input = drafted_input
+
+        if engine_input is not None:
+            decision = engine.step(engine_input)
+        else:
+            decision = {
+                "kind": DecisionKind.PASSTHROUGH,
+                "state": engine.state,
+                "prompt_to_user": None,
+            }
+
+        if session.mode == MODE_PERSISTENT and session.session_key is not None:
+            CHECKPOINT_STORE.save(
+                session.session_key,
+                checkpoint_to_jsonable(engine.export_checkpoint_json()),
+            )
+
+        logger.debug("litellm_proxy: decision_kind=%s", decision["kind"])
+
+        if is_clarify(decision):
+            logger.debug("litellm_proxy: blocking_on_clarify=true")
+            return get_clarify_prompt(decision) or "Confirmation required."
+
+        compiled_state = engine.state
+        system_message: dict[str, object] = {
+            "role": "system",
+            "content": "You are a helpful assistant.\n"
+            + _render_compiled_state_contract(compiled_state),
+        }
+        logger.debug("litellm_proxy: inject_system_message=true")
+        # Preserve original request messages; drafting changes only compiler input.
+        data["messages"] = [system_message, *request_messages]
+        return data
+
+
+proxy_handler_instance = ContextCompilerPreCallHookWithPreprocessor()