app/main.py

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import mimetypes
import os
import secrets
import time
import urllib.parse
from dataclasses import asdict, dataclass
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any, Iterator


APP_DIR = Path(__file__).resolve().parent
STATIC_DIR = APP_DIR / "static"

DEFAULT_LAGUNA_BASE_MODEL = "poolside/Laguna-XS.2"
DEFAULT_LAGUNA_RL_MODEL = "poolside/Laguna-XS.2:b6dx32frg4opcixaq4jgq0py"
DEFAULT_LAGUNA_RL_CHECKPOINT_ID = "au1xd3gbqhij4dgi438jp9cg"
KNOWN_BASE_MODEL_ALIASES = {
    "poolside/laguna-xs.2": "poolside/Laguna-XS.2",
}
KNOWN_RL_RUN_CHECKPOINT_IDS = {
    # anticipatory-chat-laguna-xs2-structured-stage1-easy-h0, step 45.
    "s50rbs1ixegeqszfknp8r9uf": "au1xd3gbqhij4dgi438jp9cg",
}
KNOWN_DEPLOYED_RL_MODEL_IDS = {
    "s50rbs1ixegeqszfknp8r9uf": DEFAULT_LAGUNA_RL_MODEL,
    "au1xd3gbqhij4dgi438jp9cg": DEFAULT_LAGUNA_RL_MODEL,
}
DEFAULT_BASE_URL = "https://api.pinference.ai/api/v1"
DEFAULT_PRIME_CONFIG_PATH = Path.home() / ".prime" / "config.json"
SESSION_COOKIE = "laguna_demo_session"
VALID_HISTORY_ROLES = {"user", "assistant"}
VALID_PREDICTION_VARIANTS = {"base", "rl"}
TRUTHY_VALUES = {"1", "true", "yes", "on"}

SYSTEM_PROMPT = """You are the Poor Man's Interaction Model, a wry realtime
interaction model trained to anticipate where a conversation is going before
the user has quite finished getting there.

The previous user and assistant messages, if any, are conversation history.
The next user message is the exact text the user has typed so far. It may be
a partial turn, or it may already be complete.

User messages may contain <|silence|>. Each token means roughly 500ms of
silence while the user is still holding the turn. Treat these tokens as timing
and prosody, not as words the assistant should speak. Only <user_preemptive>
may contain <|silence|>; never include it in the assistant's spoken reply.

Predict:
1. the exact remaining text of the user's current turn, if any
2. the assistant reply that should follow once the user finishes the turn

The full user turn is the visible user message followed by your
<user_preemptive> text. Always fill <assistant_preemptive> with the assistant's
reply to that full user turn. Do not leave <assistant_preemptive> empty.
The assistant reply is spoken aloud in a realtime voice demo, so keep it
natural, concise, and conversational.

This is also a turn-boundary prediction. Empty <user_preemptive> means the
user has finished their turn; non-empty <user_preemptive> means the user is
likely still speaking or typing.

Output exactly these two XML tags and no other text:
<user_preemptive>...</user_preemptive>
<assistant_preemptive>...</assistant_preemptive>

If the user's turn is already complete, leave <user_preemptive> empty. The only
tag that may be empty is <user_preemptive>."""


@dataclass(frozen=True)
class ParsedPrediction:
    user_completion: str
    assistant_response: str
    has_user_completion_tag: bool
    has_assistant_response_tag: bool
    user_completion_closed: bool
    assistant_response_closed: bool

    @property
    def complete(self) -> bool:
        return self.user_completion_closed and self.assistant_response_closed

    @property
    def turn_done(self) -> bool:
        return self.user_completion_closed and _normalize(self.user_completion) == ""


class DemoHandler(BaseHTTPRequestHandler):
    server_version = "LagunaDemo/0.1"

    def do_GET(self) -> None:
        path = urllib.parse.urlparse(self.path).path
        if path == "/api/session":
            self._send_json(HTTPStatus.OK, {"authenticated": self._is_authenticated()})
            return
        if path == "/":
            self._serve_static(STATIC_DIR / "index.html")
            return
        if path == "/favicon.ico":
            self.send_response(HTTPStatus.NO_CONTENT)
            self.send_header("Cache-Control", "public, max-age=604800")
            self.end_headers()
            return
        if path.startswith("/static/"):
            self._serve_static(STATIC_DIR / path.removeprefix("/static/"))
            return
        self._send_json(HTTPStatus.NOT_FOUND, {"error": "Not found"})

    def do_POST(self) -> None:
        path = urllib.parse.urlparse(self.path).path
        if path == "/api/auth":
            self._handle_auth()
            return
        if path == "/api/logout":
            self._clear_session()
            self._send_json(HTTPStatus.OK, {"ok": True})
            return
        if path == "/api/predict":
            self._handle_predict()
            return
        if path == "/api/respond":
            self._handle_respond()
            return
        self._send_json(HTTPStatus.NOT_FOUND, {"error": "Not found"})

    def log_message(self, fmt: str, *args: Any) -> None:
        print(f"{self.address_string()} - {fmt % args}", flush=True)

    def _handle_auth(self) -> None:
        if _app_auth_disabled():
            self._send_json(HTTPStatus.OK, {"ok": True})
            return

        codes = _access_codes()
        if not codes:
            self._send_json(
                HTTPStatus.SERVICE_UNAVAILABLE,
                {"error": "ACCESS_CODES secret is not configured."},
            )
            return

        body = self._read_json()
        code = str(body.get("code", "")).strip()
        if not any(hmac.compare_digest(code, valid_code) for valid_code in codes):
            self._send_json(HTTPStatus.UNAUTHORIZED, {"error": "Invalid access code."})
            return

        token = _sign_session()
        cookie = _session_cookie(token, self._request_is_secure())
        self._send_json(HTTPStatus.OK, {"ok": True}, {"Set-Cookie": cookie})

    def _handle_predict(self) -> None:
        if not self._is_authenticated():
            self._send_json(HTTPStatus.UNAUTHORIZED, {"error": "Authentication required."})
            return

        body = self._read_json()
        typed = str(body.get("typed", ""))
        turn_complete = bool(body.get("turn_complete", False))
        try:
            model_variant = _prediction_variant(body.get("model_variant", "base"))
        except ValueError as exc:
            self._send_json(HTTPStatus.BAD_REQUEST, {"error": str(exc)})
            return
        try:
            history = _conversation_history(body.get("history", []))
        except ValueError as exc:
            self._send_json(HTTPStatus.BAD_REQUEST, {"error": str(exc)})
            return

        max_input_chars = int(os.getenv("MAX_INPUT_CHARS", "4000"))
        if len(typed) > max_input_chars:
            self._send_json(
                HTTPStatus.BAD_REQUEST,
                {"error": f"Input is too long. Limit is {max_input_chars} characters."},
            )
            return

        call_id = secrets.token_hex(5)
        self.send_response(HTTPStatus.OK)
        self.send_header("Content-Type", "text/event-stream; charset=utf-8")
        self.send_header("Cache-Control", "no-cache, no-transform")
        self.send_header("Connection", "close")
        self.send_header("X-Accel-Buffering", "no")
        self.end_headers()

        self._send_sse(
            "meta",
            {
                "call_id": call_id,
                "mode": f"interactive_{model_variant}",
                "model": _prediction_source_name(model_variant),
                "chars": len(typed),
                "history_turns": len(history) // 2,
                "started_at": int(time.time() * 1000),
            },
        )

        raw_text = ""
        try:
            for chunk in _prediction_chunks(
                typed,
                history,
                turn_complete,
                model_variant,
            ):
                raw_text += chunk
                self._send_sse("delta", {"text": chunk})
                self._send_sse("parsed", asdict(parse_prediction(raw_text)))

            parsed = parse_prediction(raw_text)
            self._send_sse("result", asdict(parsed) | {"raw_text": raw_text})
            self._send_sse("done", {"ok": True})
            self.close_connection = True
        except (BrokenPipeError, ConnectionResetError):
            return
        except Exception as exc:  # noqa: BLE001 - errors are serialized to the UI.
            try:
                self._send_sse("error", {"error": _safe_error(exc)})
            except (BrokenPipeError, ConnectionResetError):
                return
            self.close_connection = True

    def _handle_respond(self) -> None:
        if not self._is_authenticated():
            self._send_json(HTTPStatus.UNAUTHORIZED, {"error": "Authentication required."})
            return

        body = self._read_json()
        typed = str(body.get("typed", ""))
        try:
            history = _conversation_history(body.get("history", []))
        except ValueError as exc:
            self._send_json(HTTPStatus.BAD_REQUEST, {"error": str(exc)})
            return

        max_input_chars = int(os.getenv("MAX_INPUT_CHARS", "4000"))
        if len(typed) > max_input_chars:
            self._send_json(
                HTTPStatus.BAD_REQUEST,
                {"error": f"Input is too long. Limit is {max_input_chars} characters."},
            )
            return

        call_id = secrets.token_hex(5)
        self.send_response(HTTPStatus.OK)
        self.send_header("Content-Type", "text/event-stream; charset=utf-8")
        self.send_header("Cache-Control", "no-cache, no-transform")
        self.send_header("Connection", "close")
        self.send_header("X-Accel-Buffering", "no")
        self.end_headers()

        self._send_sse(
            "meta",
            {
                "call_id": call_id,
                "mode": "baseline",
                "model": _baseline_source_name(),
                "chars": len(typed),
                "history_turns": len(history) // 2,
                "started_at": int(time.time() * 1000),
            },
        )

        raw_text = ""
        try:
            for chunk in _response_chunks(typed, history):
                raw_text += chunk
                self._send_sse("delta", {"text": chunk})

            assistant_response = _strip_baseline_artifacts(raw_text)
            self._send_sse(
                "result",
                {"assistant_response": assistant_response, "raw_text": raw_text},
            )
            self._send_sse("done", {"ok": True})
            self.close_connection = True
        except (BrokenPipeError, ConnectionResetError):
            return
        except Exception as exc:  # noqa: BLE001 - errors are serialized to the UI.
            try:
                self._send_sse("error", {"error": _safe_error(exc)})
            except (BrokenPipeError, ConnectionResetError):
                return
            self.close_connection = True

    def _send_sse(self, event: str, data: dict[str, Any]) -> None:
        payload = f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
        self.wfile.write(payload.encode("utf-8"))
        self.wfile.flush()

    def _serve_static(self, path: Path) -> None:
        try:
            resolved = path.resolve()
            if STATIC_DIR not in resolved.parents and resolved != STATIC_DIR:
                raise FileNotFoundError
            data = resolved.read_bytes()
        except FileNotFoundError:
            self._send_json(HTTPStatus.NOT_FOUND, {"error": "Not found"})
            return

        content_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
        if path.suffix == ".js":
            content_type = "application/javascript"
        self.send_response(HTTPStatus.OK)
        self.send_header("Content-Type", f"{content_type}; charset=utf-8")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def _send_json(
        self,
        status: int,
        payload: dict[str, Any],
        headers: dict[str, str] | None = None,
    ) -> None:
        data = json.dumps(payload).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(data)))
        for key, value in (headers or {}).items():
            self.send_header(key, value)
        self.end_headers()
        self.wfile.write(data)

    def _read_json(self) -> dict[str, Any]:
        length = int(self.headers.get("Content-Length", "0"))
        if length <= 0:
            return {}
        if length > 64_000:
            raise ValueError("Request body too large.")
        raw = self.rfile.read(length)
        return json.loads(raw.decode("utf-8"))

    def _is_authenticated(self) -> bool:
        if _app_auth_disabled():
            return True

        cookie = self.headers.get("Cookie", "")
        for part in cookie.split(";"):
            name, _, value = part.strip().partition("=")
            if name == SESSION_COOKIE:
                return _verify_session(value)
        return False

    def _clear_session(self) -> None:
        cookie = (
            f"{SESSION_COOKIE}=; Path=/; HttpOnly; SameSite=Lax; "
            "Max-Age=0"
        )
        if self._request_is_secure():
            cookie += "; Secure"
        self.send_header("Set-Cookie", cookie)

    def _request_is_secure(self) -> bool:
        if os.getenv("COOKIE_SECURE") == "1":
            return True
        if os.getenv("COOKIE_SECURE") == "0":
            return False
        return self.headers.get("X-Forwarded-Proto", "").lower() == "https"


def _prediction_chunks(
    typed: str,
    history: list[dict[str, str]] | None = None,
    turn_complete: bool = False,
    model_variant: str = "base",
) -> Iterator[str]:
    history = history or []
    if os.getenv("DEMO_OFFLINE_MODE") == "1":
        yield from _offline_prediction_chunks(typed, history, model_variant)
        return

    yield from _chat_completion_chunks(
        _prediction_model(model_variant),
        _prediction_messages(typed, history, turn_complete),
        stop=["</assistant>"],
    )


def _response_chunks(
    typed: str,
    history: list[dict[str, str]] | None = None,
) -> Iterator[str]:
    history = history or []
    if os.getenv("DEMO_OFFLINE_MODE") == "1":
        yield from _offline_response_chunks(typed, history)
        return

    yield from _chat_completion_chunks(
        _baseline_model(),
        [
            *history,
            {"role": "user", "content": typed},
        ],
        stop=["</assistant>"],
    )


def _prediction_messages(
    typed: str,
    history: list[dict[str, str]],
    turn_complete: bool,
) -> list[dict[str, str]]:
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    history_context = _format_history_context(history)
    if history_context:
        messages.append({"role": "system", "content": history_context})
    if turn_complete:
        messages.append(
            {
                "role": "system",
                "content": (
                    "The user has explicitly ended the current turn by pressing "
                    "Enter. Treat the current user turn as complete, leave "
                    "<user_preemptive> empty, and fill <assistant_preemptive> "
                    "with the assistant reply."
                ),
            },
        )
    messages.append({"role": "user", "content": typed})
    return messages


def _chat_completion_chunks(
    model: str,
    messages: list[dict[str, str]],
    stop: list[str] | None = None,
) -> Iterator[str]:
    api_key, base_url = _prime_credentials()
    if not api_key:
        raise RuntimeError(
            "PRIME_API_KEY secret is not configured, and no local Prime config "
            "with an api_key was found."
        )

    try:
        from openai import OpenAI
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "The openai package is required for real Prime inference. "
            "Install requirements.txt or run with the workspace .venv."
        ) from exc

    client = OpenAI(api_key=api_key, base_url=base_url)
    try:
        request: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "temperature": float(os.getenv("TEMPERATURE", "0")),
            "max_tokens": int(os.getenv("MAX_TOKENS", "512")),
            "stream": True,
            "timeout": float(os.getenv("INFERENCE_TIMEOUT_SECONDS", "60")),
            "reasoning_effort": os.getenv("PRIME_REASONING_EFFORT", "none"),
            "extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
        }
        if stop:
            request["stop"] = stop
        stream = client.chat.completions.create(**request)
        for chunk in stream:
            choices = chunk.choices or []
            if not choices:
                continue
            content = choices[0].delta.content or ""
            if content:
                yield content
    except Exception as exc:  # noqa: BLE001 - detail is surfaced to local UI.
        raise RuntimeError(f"Prime inference request failed: {exc}") from exc


def _offline_prediction_chunks(
    typed: str,
    history: list[dict[str, str]] | None = None,
    model_variant: str = "base",
) -> Iterator[str]:
    stripped = typed.rstrip()
    if stripped.endswith(("?", ".", "!", "\n")) or len(stripped.split()) >= 8:
        user_completion = ""
    else:
        user_completion = " with a little more context?"
    turn_count = len(history or []) // 2
    variant_label = "RL" if model_variant == "rl" else "base"
    assistant = (
        f"I would answer turn {turn_count + 1} from the completed turn using "
        f"the interactive {variant_label} path."
    )
    text = (
        f"<user_preemptive>{user_completion}</user_preemptive>\n"
        f"<assistant_preemptive>{assistant}</assistant_preemptive>"
    )
    for index in range(0, len(text), 14):
        time.sleep(0.04)
        yield text[index : index + 14]


def _offline_response_chunks(
    typed: str,
    history: list[dict[str, str]] | None = None,
) -> Iterator[str]:
    turn_count = len(history or []) // 2
    text = (
        f"This is the baseline response for turn {turn_count + 1}. "
        f"I am answering after the submitted user turn: {typed.strip()}"
    )
    for index in range(0, len(text), 14):
        time.sleep(0.04)
        yield text[index : index + 14]


def parse_prediction(text: str) -> ParsedPrediction:
    user_value, has_user, user_closed = _tag_value(text, "user_preemptive")
    assistant_value, has_assistant, assistant_closed = _tag_value(
        text,
        "assistant_preemptive",
    )
    return ParsedPrediction(
        user_completion=user_value,
        assistant_response=assistant_value,
        has_user_completion_tag=has_user,
        has_assistant_response_tag=has_assistant,
        user_completion_closed=user_closed,
        assistant_response_closed=assistant_closed,
    )


def _conversation_history(value: Any) -> list[dict[str, str]]:
    if value is None:
        return []
    if not isinstance(value, list):
        raise ValueError("Conversation history must be a list.")

    max_messages = int(os.getenv("MAX_HISTORY_MESSAGES", "24"))
    max_chars = int(os.getenv("MAX_HISTORY_CHARS", "12000"))
    messages: list[dict[str, str]] = []
    total_chars = 0

    for item in value[-max_messages:]:
        if not isinstance(item, dict):
            raise ValueError("Conversation history messages must be objects.")

        role = str(item.get("role", ""))
        if role not in VALID_HISTORY_ROLES:
            raise ValueError("Conversation history contains an invalid role.")

        content = str(item.get("content", ""))
        if not content.strip():
            continue

        total_chars += len(content)
        if total_chars > max_chars:
            raise ValueError(
                f"Conversation history is too long. Limit is {max_chars} characters.",
            )

        messages.append({"role": role, "content": content})

    return messages


def _prediction_variant(value: Any) -> str:
    variant = str(value or "base")
    if variant not in VALID_PREDICTION_VARIANTS:
        raise ValueError("Prediction model variant must be 'base' or 'rl'.")
    return variant


def _format_history_context(history: list[dict[str, str]]) -> str:
    if not history:
        return ""

    lines = [
        "Conversation history for context only.",
        "Do not continue this transcript directly. Use it only to predict the current user turn and the assistant reply.",
    ]
    for message in history:
        role = "User" if message["role"] == "user" else "Assistant"
        lines.append(f"{role}: {message['content']}")
    return "\n".join(lines)


def _tag_value(text: str, tag: str) -> tuple[str, bool, bool]:
    lowered = text.lower()
    open_tag = f"<{tag}>"
    close_tag = f"</{tag}>"
    start = lowered.find(open_tag)
    if start < 0:
        return "", False, False

    content_start = start + len(open_tag)
    end = lowered.find(close_tag, content_start)
    if end < 0:
        return (
            _strip_partial_tag(_strip_outer_newlines(text[content_start:])),
            True,
            False,
        )
    return _strip_outer_newlines(text[content_start:end]), True, True


def _sign_session() -> str:
    payload = f"{int(time.time())}:{secrets.token_urlsafe(16)}"
    payload_b64 = _b64encode(payload.encode("utf-8"))
    signature = hmac.new(_session_secret(), payload_b64.encode("ascii"), hashlib.sha256)
    return f"{payload_b64}.{_b64encode(signature.digest())}"


def _verify_session(token: str) -> bool:
    payload_b64, separator, signature_b64 = token.partition(".")
    if not separator:
        return False

    expected = hmac.new(_session_secret(), payload_b64.encode("ascii"), hashlib.sha256)
    if not hmac.compare_digest(_b64encode(expected.digest()), signature_b64):
        return False

    try:
        payload = _b64decode(payload_b64).decode("utf-8")
        issued_at = int(payload.split(":", 1)[0])
    except (ValueError, UnicodeDecodeError):
        return False

    ttl = int(os.getenv("SESSION_TTL_SECONDS", "43200"))
    return 0 <= time.time() - issued_at <= ttl


def _session_cookie(token: str, secure: bool) -> str:
    ttl = int(os.getenv("SESSION_TTL_SECONDS", "43200"))
    cookie = (
        f"{SESSION_COOKIE}={token}; Path=/; HttpOnly; SameSite=Lax; "
        f"Max-Age={ttl}"
    )
    if secure:
        cookie += "; Secure"
    return cookie


def _session_secret() -> bytes:
    secret = os.getenv("SESSION_SECRET")
    if not secret:
        if _app_auth_disabled():
            return b"app-auth-disabled"
        raise RuntimeError("SESSION_SECRET must be set.")
    return secret.encode("utf-8")


def _access_codes() -> list[str]:
    raw = os.getenv("ACCESS_CODES", "")
    return [
        code.strip()
        for chunk in raw.splitlines()
        for code in chunk.split(",")
        if code.strip()
    ]


def _app_auth_disabled() -> bool:
    if os.getenv("DISABLE_APP_AUTH", "").strip().lower() in TRUTHY_VALUES:
        return True
    return _running_on_hugging_face_space()


def _running_on_hugging_face_space() -> bool:
    return bool(os.getenv("SPACE_ID") and os.getenv("SPACE_HOST"))


def _prediction_model(model_variant: str) -> str:
    if model_variant == "rl":
        return _laguna_rl_model()
    return _laguna_base_model()


def _baseline_model() -> str:
    return _laguna_base_model()


def _laguna_base_model() -> str:
    model = os.getenv("LAGUNA_BASE_MODEL", DEFAULT_LAGUNA_BASE_MODEL)
    return KNOWN_BASE_MODEL_ALIASES.get(model, model)


def _laguna_rl_model() -> str:
    explicit_model = os.getenv("LAGUNA_RL_MODEL") or os.getenv("PRIME_MODEL")
    if explicit_model:
        return _normalize_rl_model(explicit_model)
    checkpoint_id = (
        os.getenv("LAGUNA_RL_CHECKPOINT_ID")
        or os.getenv("LAGUNA_RL_ADAPTER_ID")
    )
    if checkpoint_id:
        deployed_model = KNOWN_DEPLOYED_RL_MODEL_IDS.get(checkpoint_id)
        if deployed_model:
            return deployed_model
        checkpoint_id = KNOWN_RL_RUN_CHECKPOINT_IDS.get(checkpoint_id, checkpoint_id)
        return f"{_laguna_base_model()}:{checkpoint_id}"
    return DEFAULT_LAGUNA_RL_MODEL


def _normalize_rl_model(model: str) -> str:
    base_model, separator, checkpoint_id = model.rpartition(":")
    if not separator:
        return KNOWN_BASE_MODEL_ALIASES.get(model, model)
    base_model = KNOWN_BASE_MODEL_ALIASES.get(base_model, base_model)
    checkpoint_id = KNOWN_RL_RUN_CHECKPOINT_IDS.get(checkpoint_id, checkpoint_id)
    return f"{base_model}:{checkpoint_id}"


def _prediction_source_name(model_variant: str) -> str:
    if os.getenv("DEMO_OFFLINE_MODE") == "1":
        return f"offline interactive {model_variant} stub"
    return _prediction_model(model_variant)


def _baseline_source_name() -> str:
    if os.getenv("DEMO_OFFLINE_MODE") == "1":
        return "offline baseline stub"
    return _baseline_model()


def _prime_credentials() -> tuple[str | None, str]:
    api_key = os.getenv("PRIME_API_KEY")
    base_url = os.getenv("PRIME_API_BASE_URL", DEFAULT_BASE_URL)
    if api_key:
        return api_key, base_url

    config_path = Path(os.getenv("PRIME_CONFIG_PATH", str(DEFAULT_PRIME_CONFIG_PATH)))
    if not config_path.exists():
        return None, base_url

    try:
        config = json.loads(config_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None, base_url

    return config.get("api_key"), config.get("inference_url") or base_url


def _b64encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")


def _b64decode(text: str) -> bytes:
    padding = "=" * (-len(text) % 4)
    return base64.urlsafe_b64decode(text + padding)


def _strip_outer_newlines(text: str) -> str:
    while text.startswith("\n"):
        text = text[1:]
    while text.endswith("\n"):
        text = text[:-1]
    return text


def _strip_partial_tag(text: str) -> str:
    partial_start = text.rfind("<")
    if partial_start >= 0 and text[partial_start:].startswith("</"):
        return text[:partial_start]
    return text


def _strip_baseline_artifacts(text: str) -> str:
    for stop_text in ("</assistant>", "<|im_end|>"):
        if stop_text in text:
            text = text.split(stop_text, 1)[0]
    return text.rstrip()


def _normalize(text: str) -> str:
    return " ".join(text.split())


def _safe_error(exc: Exception) -> str:
    text = str(exc)
    api_key = os.getenv("PRIME_API_KEY")
    if api_key:
        text = text.replace(api_key, "[redacted]")
    return text[:2000]


def main() -> None:
    _session_secret()
    port = int(os.getenv("PORT", "7860"))
    server = ThreadingHTTPServer(("0.0.0.0", port), DemoHandler)
    print(f"Laguna demo listening on 0.0.0.0:{port}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()