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 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 text. Always fill with the assistant's reply to that full user turn. Do not leave 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 means the user has finished their turn; non-empty means the user is likely still speaking or typing. Output exactly these two XML tags and no other text: ... ... If the user's turn is already complete, leave empty. The only tag that may be empty is .""" @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=[""], ) 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=[""], ) 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 " " empty, and fill " "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_completion}\n" f"{assistant}" ) 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"" 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(" str: for stop_text in ("", "<|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()