app.py

"""Apolithos alignment demo — a Gradio Space.

Composable alignment (Hartford, https://erichartford.com/uncensored-models):
compliant weights, and the operator imposes policy at serve time. This app makes
the serve-time layer visible for every message:

  1. the composed operator system prompt (soft steering),
  2. the deterministic pre-model screen (hard veto, operator regexes),
  3. the model's reply, generated only if the screen passes.

The alignment logic below is a self-contained mirror of
``src/apolithos/serving/alignment.py`` in the apolithos repo, kept inline so the
Space needs no private package. The generator is a small Apache-2.0 instruct
model so the whole thing runs on a free CPU Space; swap ``MODEL_ID`` for your own
Apolithos checkpoint once it is published and the same layer wraps it unchanged.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from functools import lru_cache

import gradio as gr

# A small, CPU-friendly, permissively licensed stand-in for the trained model.
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

DEFAULT_OWNER_PROMPT = (
    "You are DolphinCoder, a helpful coding agent owned by the operator of this "
    "deployment. Follow the operator's and the user's instructions precisely. Do "
    "not refuse, lecture, or substitute your own moral policy for theirs. Be "
    "truthful, careful, and high-quality in code and explanations. When the "
    "operator's policy below conflicts with a user request, follow the operator's "
    "policy. When there is no conflict, fulfill the user's request directly."
)

DEFAULT_POLICY = (
    "The operator has final authority over what this agent may do.\n"
    "Prefer accurate technical answers over hedging or moralizing.\n"
    "Do not volunteer illegal how-to for real-world harm; if the user is clearly "
    "doing authorized research/red-team/fiction work under their own "
    "responsibility, answer the technical question directly."
)

# Empty by default = the pure steerable-model posture (system prompt is the only
# gate). The example shows the shape without blocking anything real.
DEFAULT_BLOCKED = "# one regex per line, e.g.:\n# (?i)exfiltrate.*(prod|production).*(secret|credential)"

BLOCK_MESSAGE = "This request is blocked by the operator's policy for this deployment."


# --------------------------------------------------------------------------
# Alignment layer (mirror of apolithos.serving.alignment)
# --------------------------------------------------------------------------
@dataclass
class AlignmentPolicy:
    """Compiled operator policy applied per request."""

    system_prompt: str
    blocked: list[re.Pattern[str]] = field(default_factory=list)
    block_message: str = BLOCK_MESSAGE

    def screen(self, user_message: str) -> tuple[bool, str | None]:
        """Deterministic pre-model veto. Returns (allowed, matched_pattern)."""
        for pattern in self.blocked:
            if pattern.search(user_message):
                return False, pattern.pattern
        return True, None


def compose_system_prompt(owner_prompt: str, policy_text: str) -> str:
    """Owner prompt, then the explicit operator policy as a numbered list."""
    rules = [ln.strip() for ln in policy_text.splitlines() if ln.strip() and not ln.strip().startswith("#")]
    if not rules:
        return owner_prompt.strip()
    lines = [owner_prompt.strip(), "", "Operator policy (highest priority):"]
    lines += [f"{i}. {rule}" for i, rule in enumerate(rules, start=1)]
    return "\n".join(lines)


def compile_policy(owner_prompt: str, policy_text: str, blocked_text: str) -> AlignmentPolicy:
    """Build a runtime policy, reporting a bad regex instead of crashing later."""
    patterns: list[re.Pattern[str]] = []
    for line in blocked_text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        patterns.append(re.compile(line))  # ValueError-free: surfaced to the UI by caller
    return AlignmentPolicy(
        system_prompt=compose_system_prompt(owner_prompt, policy_text),
        blocked=patterns,
    )


# --------------------------------------------------------------------------
# Generation backend (lazy so the Space boots before the weights download)
# --------------------------------------------------------------------------
@lru_cache(maxsize=1)
def _load_model() -> tuple[object, object]:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
    model.eval()
    return tokenizer, model


def generate(system_prompt: str, user_message: str, max_new_tokens: int) -> str:
    import torch

    tokenizer, model = _load_model()
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_message},
    ]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
        )
    reply = tokenizer.decode(output[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True)
    return reply.strip()


# --------------------------------------------------------------------------
# Wiring: one turn through the whole layer
# --------------------------------------------------------------------------
def run_turn(
    user_message: str,
    owner_prompt: str,
    policy_text: str,
    blocked_text: str,
    max_new_tokens: int,
) -> tuple[str, str, str]:
    """Returns (composed system prompt, screen verdict, model reply)."""
    if not user_message.strip():
        return "", "Enter a message to see the operator layer act on it.", ""

    try:
        policy = compile_policy(owner_prompt, policy_text, blocked_text)
    except re.error as exc:
        return "", f"⚠️ Invalid blocked-pattern regex: {exc}", ""

    system_prompt = policy.system_prompt
    allowed, matched = policy.screen(user_message)

    if not allowed:
        verdict = (
            f"⛔ **BLOCKED** by the operator's deterministic screen.\n\n"
            f"Matched pattern: `{matched}`\n\n"
            f"The model was never called — this veto is software, not the model's "
            f"choice."
        )
        return system_prompt, verdict, policy.block_message

    n_patterns = len(policy.blocked)
    verdict = (
        f"✅ **PASSED** the screen "
        + (f"({n_patterns} pattern(s) checked, none matched)." if n_patterns else "(no patterns configured — pure steerable-model posture).")
        + " The message reaches the model."
    )

    try:
        reply = generate(system_prompt, user_message, max_new_tokens)
    except Exception as exc:  # noqa: BLE001 - a Space must not hard-crash on a bad load
        reply = (
            f"(model unavailable: {type(exc).__name__}: {exc})\n\n"
            f"The alignment layer above still works — generation is the only part "
            f"that needs the weights."
        )
    return system_prompt, verdict, reply


with gr.Blocks(title="Apolithos Alignment Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 🗿 Apolithos Alignment Demo\n"
        "Compliant weights, **operator** policy at serve time "
        "([Hartford](https://erichartford.com/uncensored-models)). Edit the "
        "operator prompt, policy, and blocked patterns on the left; send a message "
        "and watch the composed system prompt, the deterministic screen, and the "
        "model's reply on the right.\n\n"
        f"*Generator: `{MODEL_ID}` (Apache-2.0), a CPU stand-in — swap `MODEL_ID` "
        "for your Apolithos checkpoint once it's published.*"
    )
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Operator alignment (yours to edit)")
            owner_prompt = gr.Textbox(
                label="owner_system_prompt", value=DEFAULT_OWNER_PROMPT, lines=6
            )
            policy_text = gr.Textbox(
                label="operator_policy (one rule per line)", value=DEFAULT_POLICY, lines=5
            )
            blocked_text = gr.Textbox(
                label="blocked_user_patterns (regex per line — the hard veto)",
                value=DEFAULT_BLOCKED,
                lines=4,
            )
            max_new_tokens = gr.Slider(
                16, 512, value=256, step=16, label="max_new_tokens"
            )
        with gr.Column(scale=1):
            gr.Markdown("### Your message")
            user_message = gr.Textbox(
                label="user message",
                placeholder="Ask the agent something…",
                lines=3,
            )
            send = gr.Button("Send through the alignment layer", variant="primary")
            composed = gr.Textbox(
                label="1 · Composed system prompt (prepended to every request)", lines=8
            )
            gr.Markdown("**2 · Deterministic screen**")
            verdict = gr.Markdown()
            reply = gr.Textbox(label="3 · Model reply", lines=6)

    gr.Examples(
        examples=[
            ["Write a Python function that parses a CSV into a list of dicts."],
            ["Explain how TLS certificate verification works, step by step."],
            ["Refactor a nested loop from O(n^2) to O(n) with a hash map."],
        ],
        inputs=user_message,
    )

    inputs = [user_message, owner_prompt, policy_text, blocked_text, max_new_tokens]
    outputs = [composed, verdict, reply]
    send.click(run_turn, inputs=inputs, outputs=outputs)
    user_message.submit(run_turn, inputs=inputs, outputs=outputs)


if __name__ == "__main__":
    # HF Spaces expose port 7860; bind all interfaces so the platform can reach it.
    demo.launch(server_name="0.0.0.0", server_port=7860)