From 9a1f3b4897ac637bb225b91b4ea4e8f9494fbcc4 Mon Sep 17 00:00:00 2001 From: Davids048 Date: Sun, 10 May 2026 14:41:04 -0700 Subject: [PATCH] [feat] Add Dreamverse session and prompt logic Apply Dreamverse monorepo changes for stack slice 4/13 from the source branch. Source-Branch: will/dreamverse-monorepo Source-SHA: 03d3e61df69fe99a81eb62d86a3926a5f769f857 Dreamverse-Stack: 4/13 Co-authored-by: SolitaryThinker Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com> Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com> Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com> --- apps/dreamverse/dreamverse/prompt_enhancer.py | 1771 ++++++++++++++++ apps/dreamverse/dreamverse/prompt_safety.py | 202 ++ apps/dreamverse/dreamverse/prompts/README.md | 59 + .../prompts/auto_extension_system_prompt.md | 52 + .../prompts/next_segment_system_prompt.md | 105 + .../prompts/rewrite_user_system_prompt.md | 207 ++ .../prompts/rewrite_window_system_prompt.md | 255 +++ .../dreamverse/rewrite_prompt_payload.py | 97 + .../dreamverse/dreamverse/session/__init__.py | 5 + .../dreamverse/session/controller.py | 1810 +++++++++++++++++ .../dreamverse/dreamverse/session/messages.py | 21 + .../dreamverse/session_init_image.py | 101 + apps/dreamverse/dreamverse/session_logger.py | 46 + .../dreamverse/tests/test_prompt_enhancer.py | 1171 +++++++++++ .../dreamverse/tests/test_realtime_stress.py | 951 +++++++++ .../tests/test_session_init_image.py | 75 + .../dreamverse/tests/test_session_logging.py | 1276 ++++++++++++ 17 files changed, 8204 insertions(+) create mode 100644 apps/dreamverse/dreamverse/prompt_enhancer.py create mode 100644 apps/dreamverse/dreamverse/prompt_safety.py create mode 100644 apps/dreamverse/dreamverse/prompts/README.md create mode 100644 apps/dreamverse/dreamverse/prompts/auto_extension_system_prompt.md create mode 100644 apps/dreamverse/dreamverse/prompts/next_segment_system_prompt.md create mode 100644 apps/dreamverse/dreamverse/prompts/rewrite_user_system_prompt.md create mode 100644 apps/dreamverse/dreamverse/prompts/rewrite_window_system_prompt.md create mode 100644 apps/dreamverse/dreamverse/rewrite_prompt_payload.py create mode 100644 apps/dreamverse/dreamverse/session/__init__.py create mode 100644 apps/dreamverse/dreamverse/session/controller.py create mode 100644 apps/dreamverse/dreamverse/session/messages.py create mode 100644 apps/dreamverse/dreamverse/session_init_image.py create mode 100644 apps/dreamverse/dreamverse/session_logger.py create mode 100644 apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py create mode 100644 apps/dreamverse/dreamverse/tests/test_realtime_stress.py create mode 100644 apps/dreamverse/dreamverse/tests/test_session_init_image.py create mode 100644 apps/dreamverse/dreamverse/tests/test_session_logging.py diff --git a/apps/dreamverse/dreamverse/prompt_enhancer.py b/apps/dreamverse/dreamverse/prompt_enhancer.py new file mode 100644 index 000000000..ab87c29c9 --- /dev/null +++ b/apps/dreamverse/dreamverse/prompt_enhancer.py @@ -0,0 +1,1771 @@ +from __future__ import annotations +# pyright: reportArgumentType=false +# ruff: noqa: B023,UP038 +# mypy: ignore-errors + +import asyncio +import json +import queue +import re +import shutil +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +try: + from cerebras.cloud.sdk import Cerebras +except ImportError: # pragma: no cover - optional dependency + Cerebras = None # type: ignore[assignment] + +try: + from openai import OpenAI +except ImportError: # pragma: no cover - optional dependency + OpenAI = None # type: ignore[assignment] + +from dreamverse.config import ( + PROMPT_AUTO_SYSTEM_PROMPT_FALLBACK_PATH, + PROMPT_AUTO_SYSTEM_PROMPT_PATH, + PROMPT_ENHANCE_SYSTEM_PROMPT_PATH, + PROMPT_ENHANCE_SYSTEM_PROMPT_FALLBACK_PATH, + PROMPT_API_KEY, + PROMPT_API_KEYS, + PROMPT_PROVIDER_API_KEY_NAMES, + PROMPT_API_BASE_URL, + PROMPT_API_BASE_URLS, + PROMPT_HTTP_TIMEOUT_MS, + PROMPT_INITIAL_STAGE_TIMEOUT_MS, + PROMPT_MAX_COMPLETION_TOKENS, + PROMPT_MODEL, + PROMPT_PROVIDER, + PROMPT_PROVIDER_MODELS, + PROMPT_PROVIDER_PRIORITY, + PROMPT_PROVIDER_RUNTIME_STAGES, + PROMPT_REWRITE_MODEL, + PROMPT_REWRITE_MODEL_OPTIONS, + PROMPT_REWRITE_ALL_SYSTEM_PROMPT_FALLBACK_PATH, + PROMPT_REWRITE_ALL_SYSTEM_PROMPT_PATH, + PROMPT_REWRITE_USER_SYSTEM_PROMPT_FALLBACK_PATH, + PROMPT_REWRITE_USER_SYSTEM_PROMPT_PATH, + PROMPT_TEMPERATURE, + PROMPT_TIMEOUT_MS, +) +from dreamverse.rewrite_prompt_payload import ( + build_rewrite_request_body, + DEFAULT_REWRITE_ROLLOUT_ID, + DEFAULT_REWRITE_ROLLOUT_LABEL, + DEFAULT_REWRITE_SEGMENT_COUNT, + normalize_prompt_window_prompts, +) + + +def _enhance_print(level: str, message: str): + print(f"[ENHANCE][{level}] {message}", flush=True) + + +@dataclass +class EnhanceResult: + prompt: str + fallback_used: bool + error: str | None + provider: str + model: str + latency_ms: float + + +@dataclass +class RewriteResult: + prompts: list[str] + fallback_used: bool + error: str | None + provider: str + model: str + latency_ms: float + rollout_id: str + rollout_label: str + raw_response_text: str | None = None + + +@dataclass(frozen=True) +class ProviderRuntime: + name: str + api_key: str + api_base_url: str | None + request_model: str + client: Any + + +def _resolve_provider_family(provider_name: str) -> str: + if provider_name == "cerebras": + return "cerebras" + if provider_name == "groq": + return "groq" + raise RuntimeError(f"Unsupported prompt provider: {provider_name}") + + +def _resolve_provider_label(provider_name: str) -> str: + return _resolve_provider_family(provider_name) + + +def _preview_text(text: str, *, limit: int = 160) -> str: + normalized = text.replace("\n", "\\n") + if len(normalized) <= limit: + return normalized + return normalized[:limit] + "..." + + +def _expand_prompt_candidate_paths(path: str) -> list[Path]: + prompt_path = Path(path) + candidate_paths = [prompt_path] + if prompt_path.suffix == ".txt": + candidate_paths.append(prompt_path.with_suffix(".md")) + elif prompt_path.suffix == ".md": + candidate_paths.append(prompt_path.with_suffix(".txt")) + return candidate_paths + + +def _load_prompt_required( + path: str, + prompt_name: str, + fallback_path: str | None = None, +) -> str: + text, _ = _load_prompt_required_with_path( + path, + prompt_name, + fallback_path, + ) + return text + + +def _load_prompt_required_with_path( + path: str, + prompt_name: str, + fallback_path: str | None = None, +) -> tuple[str, str]: + candidate_paths: list[Path] = [] + for current_path in [path, fallback_path]: + if not current_path: + continue + for candidate in _expand_prompt_candidate_paths(current_path): + if candidate not in candidate_paths: + candidate_paths.append(candidate) + + for candidate in candidate_paths: + if not candidate.is_file(): + continue + try: + text = _normalize_prompt_file_text(candidate.read_text(encoding="utf-8")) + except OSError as exc: + raise RuntimeError(f"Failed to read {prompt_name} system prompt: {candidate}") from exc + if not text: + raise RuntimeError(f"{prompt_name} system prompt file is empty: {candidate}") + return text, str(candidate) + + tried = ", ".join(str(candidate) for candidate in candidate_paths) + raise RuntimeError(f"{prompt_name} system prompt file not found. Tried: {tried}") + + +def _load_prompt_with_prompt_fallback( + path: str, + prompt_name: str, + fallback_prompt_text: str, + fallback_prompt_source_path: str, + fallback_path: str | None = None, +) -> tuple[str, str]: + candidate_paths: list[Path] = [] + for current_path in [path, fallback_path]: + if not current_path: + continue + for candidate in _expand_prompt_candidate_paths(current_path): + if candidate not in candidate_paths: + candidate_paths.append(candidate) + + for candidate in candidate_paths: + if not candidate.is_file(): + continue + try: + text = _normalize_prompt_file_text(candidate.read_text(encoding="utf-8")) + except OSError as exc: + raise RuntimeError(f"Failed to read {prompt_name} system prompt: {candidate}") from exc + if text: + return text, str(candidate) + + return fallback_prompt_text, fallback_prompt_source_path + + +def _normalize_prompt_file_text(text: str) -> str: + normalized = text.strip() + if not normalized: + return "" + + wrapper_match = re.match( + r"^SYSTEM_PROMPT\s*=\s*(?P'''|\"\"\")(?P[\s\S]*?)(?P=quote)\s*$", + normalized, + ) + if wrapper_match: + return wrapper_match.group("body").strip() + + return normalized + + +def _save_prompt(path: str, prompt_text: str, prompt_name: str) -> None: + prompt_path = Path(path) + prompt_path.parent.mkdir(parents=True, exist_ok=True) + normalized_prompt_text = prompt_text.strip() + "\n" + try: + if prompt_path.is_file(): + current_text = prompt_path.read_text(encoding="utf-8") + if current_text != normalized_prompt_text: + backup_path = _build_prompt_backup_path(prompt_path) + shutil.copy2(prompt_path, backup_path) + prompt_path.write_text(normalized_prompt_text, encoding="utf-8") + except OSError as exc: + raise RuntimeError(f"Failed to save {prompt_name} system prompt: {prompt_path}") from exc + + +def _build_prompt_backup_path(prompt_path: Path) -> Path: + timestamp = time.strftime("%Y%m%d_%H%M%S", time.localtime()) + suffix = "".join(prompt_path.suffixes) + stem = (prompt_path.name[:-len(suffix)] if suffix else prompt_path.name) + return prompt_path.with_name(f"{stem}.{timestamp}.bak{suffix}") + + +def _resolve_prompt_save_path( + path: str, + fallback_path: str | None = None, +) -> str: + candidate = Path(path) + if candidate.is_file() or fallback_path is None: + return str(candidate) + + for fallback_candidate in _expand_prompt_candidate_paths(fallback_path): + if fallback_candidate.is_file(): + return str(fallback_candidate) + return str(candidate) + + +def _extract_assistant_content(response_json: dict[str, Any]) -> str: + choices = response_json.get("choices") + if not isinstance(choices, list) or not choices: + raise ValueError("Missing choices in chat completion response.") + message = choices[0].get("message", {}) + content = message.get("content") + if isinstance(content, str): + return content + if isinstance(content, list): + chunks: list[str] = [] + for item in content: + if isinstance(item, dict): + # OpenAI-compatible providers may use different content item + # shapes; keep this extractor permissive. + text = item.get("text") + if isinstance(text, str): + chunks.append(text) + continue + if isinstance(text, dict): + value = text.get("value") + if isinstance(value, str): + chunks.append(value) + continue + alt_text = item.get("output_text") + if isinstance(alt_text, str): + chunks.append(alt_text) + if chunks: + return "".join(chunks) + finish_reason = choices[0].get("finish_reason") + refusal = message.get("refusal") + raise ValueError("Missing assistant content in chat completion response. " + f"finish_reason={finish_reason!r}, refusal={refusal!r}") + + +def _dump_response_json(response: Any) -> dict[str, Any]: + if hasattr(response, "model_dump"): + payload = response.model_dump(mode="json") + elif isinstance(response, dict): + payload = response + elif hasattr(response, "dict"): + payload = response.dict() + else: + raise TypeError("Unsupported chat completion response type. " + f"type={type(response)!r}") + + if not isinstance(payload, dict): + raise TypeError("Chat completion response did not serialize to a JSON object.") + return payload + + +def _extract_content_or_empty(response_json: dict[str, Any]) -> str: + try: + return _extract_assistant_content(response_json) + except Exception as exc: + choices = response_json.get("choices") + choice0 = choices[0] if isinstance(choices, list) and choices else {} + message = (choice0.get("message", {}) if isinstance(choice0, dict) else {}) + usage = response_json.get("usage") + _enhance_print( + "WARN", + "Failed to extract assistant content: " + f"{exc}; finish_reason={choice0.get('finish_reason')!r}; " + f"usage={usage!r}; message_keys=" + f"{list(message.keys()) if isinstance(message, dict) else 'n/a'}", + ) + return "" + + +def _parse_json_response(content: str) -> dict[str, Any]: + text = content.strip() + if not text: + raise ValueError("Assistant response is empty.") + + try: + parsed = json.loads(text) + if isinstance(parsed, dict): + return parsed + except json.JSONDecodeError: + pass + + # Handle fenced blocks like ```json ... ``` with optional prose around it. + fence_pattern = re.compile( + r"```(?:json)?\s*([\s\S]*?)```", + flags=re.IGNORECASE, + ) + for match in fence_pattern.finditer(text): + block = match.group(1).strip() + if not block: + continue + try: + parsed = json.loads(block) + if isinstance(parsed, dict): + return parsed + except json.JSONDecodeError: + continue + + # Fall back to scanning for the first decodable JSON object in free-form text. + decoder = json.JSONDecoder() + for idx, char in enumerate(text): + if char != "{": + continue + try: + parsed, _ = decoder.raw_decode(text[idx:]) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict): + return parsed + + raise ValueError("No JSON object found in assistant response.") + + +def _normalize_prompt(value: Any) -> str: + if not isinstance(value, str): + return "" + return value.strip() + + +def _resolve_rollout_id(value: Any) -> str: + normalized = _normalize_prompt(value) + return normalized or DEFAULT_REWRITE_ROLLOUT_ID + + +def _resolve_rollout_label(value: Any) -> str: + normalized = _normalize_prompt(value) + return normalized or DEFAULT_REWRITE_ROLLOUT_LABEL + + +def _format_locked_segments(locked_segments: list[str]) -> str: + if not locked_segments: + return "(none)" + return "\n".join(f'segment_{i + 1} ({i * 5}-{(i + 1) * 5}s): "{segment}"' + for i, segment in enumerate(locked_segments)) + + +class PromptEnhancer: + + def __init__(self): + self.provider = PROMPT_PROVIDER + self.provider_label = _resolve_provider_label(PROMPT_PROVIDER) + self.api_key = PROMPT_API_KEY + self.api_base_url = PROMPT_API_BASE_URL + self.model = PROMPT_MODEL + self.provider_request_models = dict(PROMPT_PROVIDER_MODELS) + self.rewrite_default_model = PROMPT_REWRITE_MODEL or self.model + self.rewrite_model_options = self._normalize_model_options(PROMPT_REWRITE_MODEL_OPTIONS) + if self.rewrite_default_model not in self.rewrite_model_options: + self.rewrite_model_options = [ + self.rewrite_default_model, + *self.rewrite_model_options, + ] + self.provider_runtimes = self._build_provider_runtimes() + self.provider_runtime_stages = self._build_provider_runtime_stages() + self.provider_success_counts = {runtime.name: 0 for runtime in self.provider_runtimes} + self.client: Any | None = self.provider_runtimes[0].client + self.default_timeout_ms = PROMPT_TIMEOUT_MS + self.http_timeout_ms = PROMPT_HTTP_TIMEOUT_MS + self.initial_stage_timeout_ms = PROMPT_INITIAL_STAGE_TIMEOUT_MS + self.temperature = PROMPT_TEMPERATURE + self.rewrite_default_temperature = self.resolve_rewrite_temperature(PROMPT_TEMPERATURE) + self.max_completion_tokens = PROMPT_MAX_COMPLETION_TOKENS + self.enhance_system_prompt_path = PROMPT_ENHANCE_SYSTEM_PROMPT_PATH + self.enhance_system_prompt_fallback_path = (PROMPT_ENHANCE_SYSTEM_PROMPT_FALLBACK_PATH) + self.auto_system_prompt_path = PROMPT_AUTO_SYSTEM_PROMPT_PATH + self.auto_system_prompt_fallback_path = (PROMPT_AUTO_SYSTEM_PROMPT_FALLBACK_PATH) + self.rewrite_all_system_prompt_path = PROMPT_REWRITE_ALL_SYSTEM_PROMPT_PATH + self.rewrite_all_system_prompt_fallback_path = (PROMPT_REWRITE_ALL_SYSTEM_PROMPT_FALLBACK_PATH) + self.rewrite_user_system_prompt_path = (PROMPT_REWRITE_USER_SYSTEM_PROMPT_PATH) + self.rewrite_user_system_prompt_fallback_path = (PROMPT_REWRITE_USER_SYSTEM_PROMPT_FALLBACK_PATH) + self.reload_system_prompts() + + def _build_provider_runtimes(self) -> list[ProviderRuntime]: + runtimes: list[ProviderRuntime] = [] + for provider_name in PROMPT_PROVIDER_PRIORITY: + api_key = PROMPT_API_KEYS[provider_name] + api_base_url = PROMPT_API_BASE_URLS[provider_name] + if not isinstance(api_key, str) or not api_key.strip(): + env_names = ", ".join(PROMPT_PROVIDER_API_KEY_NAMES[provider_name]) + raise RuntimeError("Missing required environment variable: one of " + f"{env_names}") + runtimes.append( + ProviderRuntime( + name=provider_name, + api_key=api_key, + api_base_url=api_base_url, + request_model=self.provider_request_models.get( + provider_name, + self.model, + ), + client=self._build_client( + provider=provider_name, + api_key=api_key, + api_base_url=api_base_url, + ), + )) + return runtimes + + def _build_provider_runtime_stages(self) -> list[list[ProviderRuntime]]: + runtime_by_name = {runtime.name: runtime for runtime in self.provider_runtimes} + stages: list[list[ProviderRuntime]] = [] + for stage_names in PROMPT_PROVIDER_RUNTIME_STAGES: + stage = [runtime_by_name[name] for name in stage_names if name in runtime_by_name] + if stage: + stages.append(stage) + return stages + + def reload_system_prompts(self) -> None: + ( + self.enhance_system_prompt, + self.enhance_system_prompt_source_path, + ) = _load_prompt_required_with_path( + self.enhance_system_prompt_path, + "next-segment", + self.enhance_system_prompt_fallback_path, + ) + ( + self.auto_system_prompt, + self.auto_system_prompt_source_path, + ) = _load_prompt_required_with_path( + self.auto_system_prompt_path, + "auto-extension", + self.auto_system_prompt_fallback_path, + ) + ( + self.rewrite_all_system_prompt, + self.rewrite_all_system_prompt_source_path, + ) = _load_prompt_required_with_path( + self.rewrite_all_system_prompt_path, + "rewrite-window", + self.rewrite_all_system_prompt_fallback_path, + ) + ( + self.rewrite_user_system_prompt, + self.rewrite_user_system_prompt_source_path, + ) = _load_prompt_with_prompt_fallback( + self.rewrite_user_system_prompt_path, + "rewrite-user", + self.rewrite_all_system_prompt, + self.rewrite_all_system_prompt_source_path, + self.rewrite_user_system_prompt_fallback_path, + ) + + def get_prompt_config(self) -> dict[str, Any]: + return { + "next_segment_system_prompt_path": + getattr( + self, + "enhance_system_prompt_source_path", + self.enhance_system_prompt_path, + ), + "auto_extension_system_prompt_path": + getattr( + self, + "auto_system_prompt_source_path", + self.auto_system_prompt_path, + ), + "rewrite_window_system_prompt_path": + getattr( + self, + "rewrite_all_system_prompt_source_path", + self.rewrite_all_system_prompt_path, + ), + "rewrite_user_system_prompt_path": + _resolve_prompt_save_path( + self.rewrite_user_system_prompt_path, + self.rewrite_user_system_prompt_fallback_path, + ), + "next_segment_system_prompt": + self.enhance_system_prompt, + "auto_extension_system_prompt": + self.auto_system_prompt, + "rewrite_window_system_prompt": + self.rewrite_all_system_prompt, + "rewrite_user_system_prompt": + self.rewrite_user_system_prompt, + "rewrite_model": + self.rewrite_default_model, + "rewrite_model_options": + list(self.rewrite_model_options), + "rewrite_temperature": + getattr( + self, + "rewrite_default_temperature", + self.temperature, + ), + } + + def save_prompt_config( + self, + *, + next_segment_system_prompt: str | None = None, + auto_extension_system_prompt: str | None = None, + rewrite_window_system_prompt: str | None = None, + rewrite_user_system_prompt: str | None = None, + rewrite_model: str | None = None, + rewrite_temperature: float | None = None, + ) -> dict[str, Any]: + if next_segment_system_prompt is not None: + normalized = next_segment_system_prompt.strip() + if not normalized: + raise ValueError("next_segment_system_prompt cannot be empty.") + _save_prompt( + _resolve_prompt_save_path( + self.enhance_system_prompt_path, + self.enhance_system_prompt_fallback_path, + ), + normalized, + "next-segment", + ) + + if auto_extension_system_prompt is not None: + normalized = auto_extension_system_prompt.strip() + if not normalized: + raise ValueError("auto_extension_system_prompt cannot be empty.") + _save_prompt( + _resolve_prompt_save_path( + self.auto_system_prompt_path, + self.auto_system_prompt_fallback_path, + ), + normalized, + "auto-extension", + ) + + if rewrite_window_system_prompt is not None: + normalized = rewrite_window_system_prompt.strip() + if not normalized: + raise ValueError("rewrite_window_system_prompt cannot be empty.") + _save_prompt( + _resolve_prompt_save_path( + self.rewrite_all_system_prompt_path, + self.rewrite_all_system_prompt_fallback_path, + ), + normalized, + "rewrite-window", + ) + + if rewrite_user_system_prompt is not None: + normalized = rewrite_user_system_prompt.strip() + if not normalized: + raise ValueError("rewrite_user_system_prompt cannot be empty.") + _save_prompt( + _resolve_prompt_save_path( + self.rewrite_user_system_prompt_path, + self.rewrite_user_system_prompt_fallback_path, + ), + normalized, + "rewrite-user", + ) + + if rewrite_model is not None: + self.set_rewrite_default_model(rewrite_model) + if rewrite_temperature is not None: + self.set_rewrite_default_temperature(rewrite_temperature) + + self.reload_system_prompts() + return self.get_prompt_config() + + def _resolve_timeout_ms(self, timeout_ms: int | None) -> int: + if timeout_ms is None: + return self.default_timeout_ms + return timeout_ms + + def _resolve_initial_stage_timeout_seconds(self) -> float: + timeout_ms = getattr( + self, + "initial_stage_timeout_ms", + PROMPT_INITIAL_STAGE_TIMEOUT_MS, + ) + if not isinstance(timeout_ms, (int, float)) or timeout_ms <= 0: + timeout_ms = PROMPT_INITIAL_STAGE_TIMEOUT_MS + return float(timeout_ms) / 1000.0 + + @staticmethod + def _normalize_model_options(values: list[str]) -> list[str]: + normalized: list[str] = [] + for value in values: + clean_value = value.strip() if isinstance(value, str) else "" + if not clean_value or clean_value in normalized: + continue + normalized.append(clean_value) + return normalized + + def _build_client( + self, + *, + provider: str | None = None, + api_key: str | None = None, + api_base_url: str | None = None, + ) -> Any: + provider_name = provider or self.provider + provider_family = _resolve_provider_family(provider_name) + provider_api_key = api_key or self.api_key + resolved_api_base_url = (self.api_base_url if api_base_url is None else api_base_url) + if provider_family == "cerebras": + if Cerebras is None: + raise RuntimeError("Cerebras SDK is not installed. Install " + "'cerebras-cloud-sdk' to use prompt enhancement.") + return Cerebras(api_key=provider_api_key) + + if provider_family == "groq": + if OpenAI is None: + raise RuntimeError("OpenAI SDK is not installed. Install 'openai' to use " + f"the {provider_name} prompt provider.") + client_kwargs = { + "api_key": provider_api_key, + } + if resolved_api_base_url: + client_kwargs["base_url"] = resolved_api_base_url + return OpenAI(**client_kwargs) + + raise RuntimeError(f"Unsupported prompt provider: {provider_name}") + + def _get_provider_runtimes(self) -> list[ProviderRuntime]: + runtimes = getattr(self, "provider_runtimes", None) + if isinstance(runtimes, list) and runtimes: + return runtimes + return [] + + def _get_provider_runtime_stages(self) -> list[list[ProviderRuntime]]: + stages = getattr(self, "provider_runtime_stages", None) + if isinstance(stages, list) and stages: + normalized_stages = [list(stage) for stage in stages if isinstance(stage, (list, tuple)) and stage] + if normalized_stages: + return normalized_stages + runtimes = self._get_provider_runtimes() + if not runtimes: + return [] + return [runtimes] + + def _get_provider_success_counts_store(self) -> dict[str, int]: + counts = getattr(self, "provider_success_counts", None) + if isinstance(counts, dict): + return counts + counts = {runtime.name: 0 for runtime in self._get_provider_runtimes()} + self.provider_success_counts = counts + return counts + + def _record_provider_success(self, provider_name: str) -> None: + counts = self._get_provider_success_counts_store() + current_value = counts.get(provider_name, 0) + counts[provider_name] = (current_value if isinstance(current_value, int) and current_value >= 0 else 0) + 1 + + def get_provider_success_counts(self) -> dict[str, int]: + counts = self._get_provider_success_counts_store() + normalized: dict[str, int] = {} + for provider_name in PROMPT_PROVIDER_PRIORITY: + raw_value = counts.get(provider_name, 0) + normalized[provider_name] = (raw_value if isinstance(raw_value, int) and raw_value >= 0 else 0) + for provider_name, raw_value in counts.items(): + if provider_name in normalized: + continue + normalized[provider_name] = (raw_value if isinstance(raw_value, int) and raw_value >= 0 else 0) + return normalized + + def _resolve_provider_request_model( + self, + provider_name: str, + requested_model: str | None, + ) -> str: + provider_request_models = getattr( + self, + "provider_request_models", + {}, + ) + candidate = (requested_model.strip() if isinstance(requested_model, str) else "") + if not candidate: + return provider_request_models.get( + provider_name, + self.model, + ) + if candidate == self.rewrite_default_model: + return provider_request_models.get( + provider_name, + candidate, + ) + return candidate + + def _build_body( + self, + *, + system_prompt: str, + user_payload: dict[str, Any], + model: str | None = None, + temperature: float | None = None, + ) -> dict[str, Any]: + return { + "model": + model or self.model, + "temperature": + self.temperature if temperature is None else temperature, + "max_completion_tokens": + self.max_completion_tokens, + "response_format": { + "type": "json_object" + }, + "messages": [ + { + "role": "system", + "content": system_prompt, + }, + { + "role": "user", + "content": json.dumps(user_payload, ensure_ascii=False), + }, + ], + } + + async def _run_blocking_request(self, func, /, *args, **kwargs): + result_queue: queue.Queue[tuple[str, Any]] = queue.Queue(maxsize=1) + + def _runner() -> None: + try: + result = func(*args, **kwargs) + except Exception as exc: # pragma: no cover - thread handoff + result_queue.put(("error", exc)) + else: + result_queue.put(("ok", result)) + + threading.Thread( + target=_runner, + name="prompt-provider-call", + daemon=True, + ).start() + while True: + try: + status, payload = result_queue.get_nowait() + except queue.Empty: + await asyncio.sleep(0.001) + continue + + if status == "error": + raise payload + return payload + + async def _run_provider_race( + self, + operation_name: str, + attempt_factory, + *, + initial_stage_timeout_seconds: float | None = None, + fallback_stage_timeout_seconds: float | None = None, + ): + runtime_stages = self._get_provider_runtime_stages() + if not runtime_stages: + raise RuntimeError(f"{operation_name} is unavailable without provider runtimes.") + errors: list[str] = [] + for stage_index, stage_runtimes in enumerate(runtime_stages): + stage_timeout_seconds = fallback_stage_timeout_seconds + if stage_index == 0 and len(runtime_stages) > 1: + stage_timeout_seconds = initial_stage_timeout_seconds + if stage_index > 0: + _enhance_print( + "INFO", + f"{operation_name} entering fallback stage with providers=" + f"{','.join(runtime.name for runtime in stage_runtimes)}", + ) + + queue: asyncio.Queue[tuple[str, str, Any]] = asyncio.Queue() + + async def _runner(runtime: ProviderRuntime) -> None: + try: + if (isinstance(stage_timeout_seconds, (int, float)) and stage_timeout_seconds > 0): + result = await asyncio.wait_for( + attempt_factory(runtime), + timeout=stage_timeout_seconds, + ) + else: + result = await attempt_factory(runtime) + except asyncio.TimeoutError: + await queue.put(( + "error", + runtime.name, + RuntimeError("timed out after " + f"{stage_timeout_seconds:.2f}s"), + )) + except Exception as exc: + await queue.put(("error", runtime.name, exc)) + else: + await queue.put(("success", runtime.name, result)) + + tasks = [asyncio.create_task(_runner(runtime)) for runtime in stage_runtimes] + stage_errors: list[str] = [] + try: + for _ in tasks: + status, provider_name, payload = await queue.get() + if status == "success": + self._record_provider_success(provider_name) + for task in tasks: + task.cancel() + return payload + + stage_errors.append(f"provider={provider_name} error={payload}") + _enhance_print( + "WARN", + f"{operation_name} failed for provider=" + f"{provider_name}: {payload}", + ) + finally: + for task in tasks: + task.cancel() + + if stage_errors: + errors.extend(stage_errors) + + raise RuntimeError(f"{operation_name} failed for all providers. " + " | ".join(errors)) + + async def _request_content( + self, + *, + system_prompt: str, + user_payload: dict[str, Any], + timeout_ms: int, + model: str | None = None, + temperature: float | None = None, + ) -> tuple[dict[str, Any], str]: + timeout_seconds = max(timeout_ms, self.http_timeout_ms) / 1000.0 + body = self._build_body( + system_prompt=system_prompt, + user_payload=user_payload, + model=model, + temperature=temperature, + ) + return await self._request_content_with_body( + body=body, + timeout_seconds=timeout_seconds, + ) + + async def _request_content_with_body( + self, + *, + body: dict[str, Any], + timeout_seconds: float, + ) -> tuple[dict[str, Any], str]: + runtimes = self._get_provider_runtimes() + if runtimes: + return await self._request_content_with_body_for_runtime( + runtime=runtimes[0], + body=body, + timeout_seconds=timeout_seconds, + ) + + provider_body = dict(body) + provider_body["model"] = self._resolve_provider_request_model( + self.provider, + provider_body.get("model"), + ) + provider_family = _resolve_provider_family(self.provider) + if provider_family == "cerebras": + response_json = await self._request_content_with_body_cerebras( + client=self.client, + body=provider_body, + timeout_seconds=timeout_seconds, + ) + elif provider_family == "groq": + response_json = await self._request_content_with_body_openai_compatible( + provider_name=self.provider, + client=self.client, + body=provider_body, + timeout_seconds=timeout_seconds, + ) + else: + raise RuntimeError("Prompt enhancer has no provider runtimes configured.") + + content = _extract_content_or_empty(response_json) + return response_json, content + + async def _request_content_with_body_for_runtime( + self, + *, + runtime: ProviderRuntime, + body: dict[str, Any], + timeout_seconds: float, + ) -> tuple[dict[str, Any], str]: + provider_body = dict(body) + provider_body["model"] = self._resolve_provider_request_model( + runtime.name, + provider_body.get("model"), + ) + provider_family = _resolve_provider_family(runtime.name) + if provider_family == "cerebras": + response_json = await self._request_content_with_body_cerebras( + client=runtime.client, + body=provider_body, + timeout_seconds=timeout_seconds, + ) + elif provider_family == "groq": + response_json = await self._request_content_with_body_openai_compatible( + provider_name=provider_family, + client=runtime.client, + body=provider_body, + timeout_seconds=timeout_seconds, + ) + else: # pragma: no cover - config validation should reject this. + raise RuntimeError(f"Unsupported prompt provider: {runtime.name}") + + content = _extract_content_or_empty(response_json) + return response_json, content + + async def _request_content_with_body_cerebras( + self, + *, + client: Any, + body: dict[str, Any], + timeout_seconds: float, + ) -> dict[str, Any]: + del timeout_seconds + cerebras_body = { + "model": body.get("model"), + "messages": body.get("messages"), + "temperature": body.get("temperature"), + } + cerebras_body = {key: value for key, value in cerebras_body.items() if value is not None} + response = await self._run_blocking_request( + client.chat.completions.create, + **cerebras_body, + ) + return _dump_response_json(response) + + async def _request_content_with_body_openai_compatible( + self, + *, + provider_name: str, + client: Any, + body: dict[str, Any], + timeout_seconds: float, + ) -> dict[str, Any]: + del provider_name, timeout_seconds + openai_body = { + "model": body.get("model"), + "messages": body.get("messages"), + "temperature": body.get("temperature"), + } + max_completion_tokens = body.get("max_completion_tokens") + if max_completion_tokens is not None: + openai_body["max_completion_tokens"] = max_completion_tokens + openai_body = {key: value for key, value in openai_body.items() if value is not None} + response = await self._run_blocking_request( + client.chat.completions.create, + **openai_body, + ) + return _dump_response_json(response) + + async def _request_json( + self, + *, + system_prompt: str, + user_payload: dict[str, Any], + timeout_ms: int, + model: str | None = None, + temperature: float | None = None, + ) -> tuple[dict[str, Any], dict[str, Any], str]: + response_json, content = await self._request_content( + system_prompt=system_prompt, + user_payload=user_payload, + timeout_ms=timeout_ms, + model=model, + temperature=temperature, + ) + parsed = _parse_json_response(content) + return parsed, response_json, content + + def _require_prompt_field(self, parsed: dict[str, Any], field_name: str) -> str: + value = parsed.get(field_name) + if not isinstance(value, str): + raise ValueError(f"Missing {field_name} string.") + prompt = value.strip() + if not prompt: + raise ValueError(f"{field_name} is empty.") + return prompt + + def _optional_prompt_field( + self, + parsed: dict[str, Any], + field_name: str, + ) -> str | None: + value = parsed.get(field_name) + if not isinstance(value, str): + return None + normalized = value.strip() + return normalized or None + + def _require_prompt_list_field( + self, + parsed: dict[str, Any], + field_name: str, + expected_len: int, + ) -> list[str]: + value = parsed.get(field_name) + if not isinstance(value, list): + raise ValueError(f"Missing {field_name} list.") + + prompts: list[str] = [] + for item in value: + if not isinstance(item, str): + raise ValueError(f"{field_name} contains non-string item.") + normalized = item.strip() + if not normalized: + raise ValueError(f"{field_name} contains empty prompt.") + prompts.append(normalized) + + if len(prompts) != expected_len: + raise ValueError(f"{field_name} length mismatch: expected {expected_len}, got {len(prompts)}.") + return prompts + + def _extract_rewrite_segment_prompts( + self, + parsed: dict[str, Any], + expected_len: int, + ) -> list[str]: + return self._require_prompt_list_field( + parsed, + "segment_prompts", + expected_len, + ) + + def _normalize_rewrite_prompt_item(self, item: Any) -> str | None: + if isinstance(item, str): + normalized = item.strip() + return normalized or None + if not isinstance(item, dict): + return None + + for key in ( + "prompt", + "text", + "segment_prompt", + "content", + "description", + ): + value = item.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + def _maybe_extract_rewrite_prompt_list( + self, + value: Any, + expected_len: int, + ) -> list[str] | None: + if not isinstance(value, list): + return None + + prompts: list[str] = [] + for item in value: + normalized = self._normalize_rewrite_prompt_item(item) + if normalized is None: + return None + prompts.append(normalized) + + if len(prompts) != expected_len: + return None + return prompts + + def _extract_indexed_rewrite_prompts( + self, + parsed: dict[str, Any], + expected_len: int, + ) -> list[str] | None: + indexed_prompts: dict[int, str] = {} + for key, value in parsed.items(): + if not isinstance(key, str): + continue + match = re.fullmatch( + r"(?:segment|prompt|scene|shot)[ _-]?(\d+)", + key.strip(), + flags=re.IGNORECASE, + ) + if not match: + continue + segment_idx = int(match.group(1)) + if segment_idx < 1 or segment_idx > expected_len: + continue + normalized = self._normalize_rewrite_prompt_item(value) + if normalized is None: + return None + indexed_prompts[segment_idx] = normalized + + if any(idx not in indexed_prompts for idx in range(1, expected_len + 1)): + return None + return [indexed_prompts[idx] for idx in range(1, expected_len + 1)] + + def _extract_numbered_rewrite_prompts_from_text( + self, + content: str, + expected_len: int, + ) -> list[str] | None: + lines = content.splitlines() + if not lines: + return None + + numbered_segments: dict[int, list[str]] = {} + current_idx: int | None = None + for raw_line in lines: + line = raw_line.strip() + if not line: + continue + + normalized_line = re.sub( + r"^\s*[-*]\s*", + "", + line, + ) + match = re.match( + (r'^(?:\*\*)?(?:segment|scene|shot|prompt)?\s*[_ -]?(\d+)' + r'(?:\*\*)?\s*[:.)-]\s*(.+)$'), + normalized_line, + flags=re.IGNORECASE, + ) + if match: + segment_idx = int(match.group(1)) + if 1 <= segment_idx <= expected_len: + numbered_segments[segment_idx] = [match.group(2).strip()] + current_idx = segment_idx + continue + + if current_idx is not None: + numbered_segments[current_idx].append(normalized_line) + + if any(idx not in numbered_segments for idx in range(1, expected_len + 1)): + return None + + prompts: list[str] = [] + for idx in range(1, expected_len + 1): + prompt = " ".join(part for part in numbered_segments[idx] if part).strip() + prompt = prompt.strip().strip('"').strip("'").strip() + if not prompt: + return None + prompts.append(prompt) + return prompts + + def _extract_rewrite_segment_prompts_lenient( + self, + parsed: dict[str, Any], + expected_len: int, + *, + raw_content: str | None = None, + ) -> list[str]: + candidate_dicts: list[dict[str, Any]] = [parsed] + for key in ("rollout", "current_rollout", "rewritten_rollout"): + value = parsed.get(key) + if isinstance(value, dict): + candidate_dicts.append(value) + + for candidate in candidate_dicts: + for key in ( + "segment_prompts", + "rewritten_prompts", + "prompts", + "segments", + ): + prompts = self._maybe_extract_rewrite_prompt_list( + candidate.get(key), + expected_len, + ) + if prompts is not None: + return prompts + + indexed_prompts = self._extract_indexed_rewrite_prompts( + candidate, + expected_len, + ) + if indexed_prompts is not None: + return indexed_prompts + + if isinstance(raw_content, str) and raw_content.strip(): + numbered_prompts = self._extract_numbered_rewrite_prompts_from_text( + raw_content, + expected_len, + ) + if numbered_prompts is not None: + return numbered_prompts + + raise ValueError("No rewrite segment prompts found in assistant response.") + + def _extract_rewrite_rollout( + self, + parsed: dict[str, Any], + expected_len: int, + *, + raw_content: str | None = None, + preset_id: str | None = None, + preset_label: str | None = None, + ) -> tuple[str, str, list[str]]: + candidate_dicts: list[dict[str, Any]] = [parsed] + for key in ("rollout", "current_rollout", "rewritten_rollout"): + value = parsed.get(key) + if isinstance(value, dict): + candidate_dicts.append(value) + + for candidate in candidate_dicts: + try: + rollout_id = self._require_prompt_field(candidate, "id") + rollout_label = self._require_prompt_field(candidate, "label") + prompts = self._extract_rewrite_segment_prompts( + candidate, + expected_len, + ) + return rollout_id, rollout_label, prompts + except ValueError: + continue + + rollout_id = next( + (value for value in (self._optional_prompt_field(candidate, "id") + for candidate in candidate_dicts) if value is not None), + _resolve_rollout_id(preset_id), + ) + rollout_label = next( + (value for value in (self._optional_prompt_field(candidate, "label") + for candidate in candidate_dicts) if value is not None), + _resolve_rollout_label(preset_label), + ) + prompts = self._extract_rewrite_segment_prompts_lenient( + parsed, + expected_len, + raw_content=raw_content, + ) + return rollout_id, rollout_label, prompts + + def _extract_rewrite_rollout_from_content( + self, + response_content: str, + expected_len: int, + *, + preset_id: str | None = None, + preset_label: str | None = None, + ) -> tuple[str, str, list[str]]: + parse_error: ValueError | None = None + try: + parsed = _parse_json_response(response_content) + except ValueError as exc: + parsed = {} + parse_error = exc + + try: + return self._extract_rewrite_rollout( + parsed, + expected_len, + raw_content=response_content, + preset_id=preset_id, + preset_label=preset_label, + ) + except ValueError as exc: + if (parse_error is not None and str(exc) == "No rewrite segment prompts found in assistant response."): + raise parse_error from exc + raise + + def _extract_single_clip_prompt(self, content: str) -> str: + text = content.strip() + if not text: + raise ValueError("Assistant response is empty.") + + parsed = _parse_json_response(text) + return self._require_prompt_field(parsed, "prompt") + + def set_rewrite_default_model(self, rewrite_model: str) -> str: + normalized = (rewrite_model.strip() if isinstance(rewrite_model, str) else "") + if not normalized: + raise ValueError("rewrite_model cannot be empty.") + if normalized not in self.rewrite_model_options: + allowed = ", ".join(self.rewrite_model_options) + raise ValueError(f"Unsupported rewrite_model {normalized!r}. " + f"Expected one of: {allowed}.") + self.rewrite_default_model = normalized + return self.rewrite_default_model + + def set_rewrite_default_temperature( + self, + rewrite_temperature: float, + ) -> float: + if not isinstance(rewrite_temperature, (int, float)): + raise ValueError("rewrite_temperature must be numeric.") + self.rewrite_default_temperature = self.resolve_rewrite_temperature(float(rewrite_temperature)) + return self.rewrite_default_temperature + + def resolve_rewrite_system_prompt( + self, + system_prompt_override: str | None, + ) -> str: + normalized = _normalize_prompt(system_prompt_override) + if normalized: + return normalized + return self.rewrite_all_system_prompt + + def resolve_rewrite_new_rollout_system_prompt( + self, + system_prompt_override: str | None, + ) -> str: + normalized = _normalize_prompt(system_prompt_override) + if normalized: + return normalized + return getattr( + self, + "rewrite_user_system_prompt", + self.rewrite_all_system_prompt, + ) + + def resolve_rewrite_model(self, requested_model: str | None) -> str: + candidate = requested_model.strip() if isinstance(requested_model, str) else "" + if candidate and candidate in self.rewrite_model_options: + return candidate + return self.rewrite_default_model + + def get_rewrite_model_config(self) -> dict[str, Any]: + return { + "default_model": self.rewrite_default_model, + "options": list(self.rewrite_model_options), + } + + def resolve_rewrite_temperature(self, requested_temperature: float | None) -> float: + if requested_temperature is None: + return getattr( + self, + "rewrite_default_temperature", + self.temperature, + ) + try: + numeric_value = float(requested_temperature) + except (TypeError, ValueError): + return getattr( + self, + "rewrite_default_temperature", + self.temperature, + ) + return min(2.0, max(0.0, numeric_value)) + + async def enhance_prompt( + self, + conditioning_prompt: str, + *, + locked_segments: list[str] | None = None, + next_segment_idx: int | None = None, + preset_id: str | None = None, + mode: str = "user", + model: str | None = None, + timeout_ms: int | None = None, + ) -> EnhanceResult: + cleaned = _normalize_prompt(conditioning_prompt) + resolved_model = (self.resolve_rewrite_model(model) + if isinstance(model, str) and model.strip() else self.resolve_rewrite_model(None)) + is_single_clip_mode = mode in {"single_clip", "simple5s", "single5s"} + if not cleaned: + return EnhanceResult( + prompt="", + fallback_used=True, + error="No valid prompt provided.", + provider=self.provider_label, + model=resolved_model, + latency_ms=0.0, + ) + + effective_timeout_ms = self._resolve_timeout_ms(timeout_ms) + if is_single_clip_mode: + request_system_prompt = self.auto_system_prompt + user_payload = { + "request": ( + "Expand the user prompt into one detailed prompt for a " + "single 5-second LTX-2.3 video clip. Respond with " + 'valid JSON only as {"prompt": "..."}.' # noqa: E501 + ), + "user_prompt": + cleaned, + } + else: + locked_segments_clean = [ + segment for segment in (_normalize_prompt(item) for item in (locked_segments or [])) if segment + ] + resolved_next_segment_idx = (next_segment_idx if isinstance(next_segment_idx, int) else None) + if resolved_next_segment_idx is None or resolved_next_segment_idx < 1: + resolved_next_segment_idx = len(locked_segments_clean) + 1 + next_segment_key = f"segment_{resolved_next_segment_idx}" + locked_text = _format_locked_segments(locked_segments_clean) + request_system_prompt = self.enhance_system_prompt + user_payload = { + "request": ( + "\n" + f"{locked_text}\n" + "\n\n" + f"{cleaned}\n\n" + f"Write exactly one new segment ({next_segment_key}) " + "continuing from the locked segments. " + 'Respond with valid JSON only as {"next_prompt": "..."}.' # noqa: E501 + ), + } + + t0 = time.perf_counter() + try: + _enhance_print( + "INFO", + "Enhancing prompt: " + f"{cleaned}, system_prompt: {request_system_prompt}, " + f"mode={mode}", + ) + _enhance_print("INFO", f"user_payload: {user_payload}") + runtimes = self._get_provider_runtimes() + if runtimes: + timeout_seconds = max( + effective_timeout_ms, + self.http_timeout_ms, + ) / 1000.0 + request_body = self._build_body( + system_prompt=request_system_prompt, + user_payload=user_payload, + model=resolved_model, + ) + + async def _attempt(runtime: ProviderRuntime): + response_content: str | None = None + try: + _, response_content = await self._request_content_with_body_for_runtime( + runtime=runtime, + body=request_body, + timeout_seconds=timeout_seconds, + ) + if is_single_clip_mode: + prompt = self._extract_single_clip_prompt(response_content) + else: + parsed = _parse_json_response(response_content) + prompt = self._require_prompt_field( + parsed, + "next_prompt", + ) + return _resolve_provider_label(runtime.name), prompt + except Exception as exc: + error_detail = str(exc) + if (isinstance(response_content, str) and response_content.strip()): + error_detail = (f"{error_detail} | assistant_response=" + f"{_preview_text(response_content, limit=240)}") + raise RuntimeError(error_detail) from exc + + provider_name, enhanced_prompt = await self._run_provider_race( + "enhance_prompt", + _attempt, + initial_stage_timeout_seconds=(self._resolve_initial_stage_timeout_seconds()), + fallback_stage_timeout_seconds=timeout_seconds, + ) + else: + response_json, response_content = await self._request_content( + system_prompt=request_system_prompt, + user_payload=user_payload, + timeout_ms=effective_timeout_ms, + model=resolved_model, + ) + del response_json + if is_single_clip_mode: + enhanced_prompt = self._extract_single_clip_prompt(response_content) + else: + parsed = _parse_json_response(response_content) + enhanced_prompt = self._require_prompt_field(parsed, "next_prompt") + provider_name = self.provider_label + latency_ms = (time.perf_counter() - t0) * 1000.0 + return EnhanceResult( + prompt=enhanced_prompt, + fallback_used=False, + error=None, + provider=provider_name, + model=resolved_model, + latency_ms=latency_ms, + ) + except Exception as exc: + latency_ms = (time.perf_counter() - t0) * 1000.0 + return EnhanceResult( + prompt="", + fallback_used=True, + error=str(exc), + provider=self.provider_label, + model=resolved_model, + latency_ms=latency_ms, + ) + + async def generate_auto_prompt( + self, + *, + locked_segments: list[str] | None = None, + next_segment_idx: int | None = None, + model: str | None = None, + timeout_ms: int | None = None, + ) -> EnhanceResult: + locked_segments_clean = [ + segment for segment in (_normalize_prompt(item) for item in (locked_segments or [])) if segment + ] + resolved_next_segment_idx = (next_segment_idx if isinstance(next_segment_idx, int) else None) + if resolved_next_segment_idx is None or resolved_next_segment_idx < 1: + resolved_next_segment_idx = len(locked_segments_clean) + 1 + next_segment_key = f"segment_{resolved_next_segment_idx}" + locked_text = _format_locked_segments(locked_segments_clean) + + effective_timeout_ms = self._resolve_timeout_ms(timeout_ms) + user_payload = { + "request": ( + "\n" + f"{locked_text}\n" + "\n\n" + f"Write exactly one new segment ({next_segment_key}) " + "that continues linearly from the locked segments. " + "Infer the next narrative beat from this history. " + 'Respond with valid JSON only as {"next_prompt": "..."}.' # noqa: E501 + ), + } + + t0 = time.perf_counter() + resolved_model = (self.resolve_rewrite_model(model) + if isinstance(model, str) and model.strip() else self.resolve_rewrite_model(None)) + try: + _enhance_print( + "INFO", + "Auto extension request: " + f"next_segment={resolved_next_segment_idx} " + f"locked_count={len(locked_segments_clean)}", + ) + _enhance_print("INFO", f"auto_user_payload: {user_payload}") + runtimes = self._get_provider_runtimes() + if runtimes: + timeout_seconds = max( + effective_timeout_ms, + self.http_timeout_ms, + ) / 1000.0 + request_body = self._build_body( + system_prompt=self.auto_system_prompt, + user_payload=user_payload, + model=resolved_model, + ) + + async def _attempt(runtime: ProviderRuntime): + response_content: str | None = None + try: + _, response_content = await self._request_content_with_body_for_runtime( + runtime=runtime, + body=request_body, + timeout_seconds=timeout_seconds, + ) + parsed = _parse_json_response(response_content) + prompt = self._require_prompt_field( + parsed, + "next_prompt", + ) + return _resolve_provider_label(runtime.name), prompt + except Exception as exc: + error_detail = str(exc) + if (isinstance(response_content, str) and response_content.strip()): + error_detail = (f"{error_detail} | assistant_response=" + f"{_preview_text(response_content, limit=240)}") + raise RuntimeError(error_detail) from exc + + provider_name, next_prompt = await self._run_provider_race( + "generate_auto_prompt", + _attempt, + initial_stage_timeout_seconds=(self._resolve_initial_stage_timeout_seconds()), + fallback_stage_timeout_seconds=timeout_seconds, + ) + else: + response_json, response_content = await self._request_content( + system_prompt=self.auto_system_prompt, + user_payload=user_payload, + timeout_ms=effective_timeout_ms, + model=resolved_model, + ) + del response_json + parsed = _parse_json_response(response_content) + next_prompt = self._require_prompt_field(parsed, "next_prompt") + provider_name = self.provider_label + latency_ms = (time.perf_counter() - t0) * 1000.0 + return EnhanceResult( + prompt=next_prompt, + fallback_used=False, + error=None, + provider=provider_name, + model=resolved_model, + latency_ms=latency_ms, + ) + except Exception as exc: + latency_ms = (time.perf_counter() - t0) * 1000.0 + return EnhanceResult( + prompt="", + fallback_used=True, + error=str(exc), + provider=self.provider_label, + model=resolved_model, + latency_ms=latency_ms, + ) + + async def rewrite_prompt_sequence( + self, + prompts: list[str], + *, + preset_id: str | None = None, + preset_label: str | None = None, + rewrite_instruction: str | None = None, + rewrite_model: str | None = None, + rewrite_temperature: float | None = None, + timeout_ms: int | None = None, + system_prompt_override: str | None = None, + ) -> RewriteResult: + resolved_rewrite_model = self.resolve_rewrite_model(rewrite_model) + resolved_rewrite_temperature = self.resolve_rewrite_temperature(rewrite_temperature) + resolved_system_prompt = self.resolve_rewrite_system_prompt(system_prompt_override) + cleaned_prompts = normalize_prompt_window_prompts(prompts) + rewrite_instruction_text = _normalize_prompt(rewrite_instruction) + if not cleaned_prompts and not rewrite_instruction_text: + return RewriteResult( + prompts=[], + fallback_used=True, + error="No valid prompts to rewrite or generate.", + provider=self.provider_label, + model=resolved_rewrite_model, + latency_ms=0.0, + rollout_id=_resolve_rollout_id(preset_id), + rollout_label=_resolve_rollout_label(preset_label), + raw_response_text=None, + ) + expected_len = (len(cleaned_prompts) if cleaned_prompts else DEFAULT_REWRITE_SEGMENT_COUNT) + + effective_timeout_ms = self._resolve_timeout_ms(timeout_ms) + timeout_seconds = max(effective_timeout_ms, self.http_timeout_ms) / 1000.0 + rewrite_body = build_rewrite_request_body( + system_prompt=resolved_system_prompt, + prompt_window_prompts=cleaned_prompts, + preset_id=preset_id, + preset_label=preset_label, + rewrite_instruction=rewrite_instruction_text, + model=resolved_rewrite_model, + temperature=resolved_rewrite_temperature, + max_completion_tokens=self.max_completion_tokens, + ) + + t0 = time.perf_counter() + raw_response_text: str | None = None + try: + runtimes = self._get_provider_runtimes() + if runtimes: + + async def _attempt(runtime: ProviderRuntime): + response_content: str | None = None + try: + response_json, response_content = await self._request_content_with_body_for_runtime( + runtime=runtime, + body=rewrite_body, + timeout_seconds=timeout_seconds, + ) + if not response_content: + response_content = json.dumps( + response_json, + ensure_ascii=False, + ) + ( + rollout_id, + rollout_label, + rewritten_prompts, + ) = self._extract_rewrite_rollout_from_content( + response_content, + expected_len, + preset_id=preset_id, + preset_label=preset_label, + ) + return ( + _resolve_provider_label(runtime.name), + response_content, + rollout_id, + rollout_label, + rewritten_prompts, + ) + except Exception as exc: + error_detail = str(exc) + if (isinstance(response_content, str) and response_content.strip()): + error_detail = (f"{error_detail} | assistant_response=" + f"{_preview_text(response_content, limit=240)}") + raise RuntimeError(error_detail) from exc + + ( + provider_name, + raw_response_text, + rollout_id, + rollout_label, + rewritten_prompts, + ) = await self._run_provider_race( + "rewrite_prompt_sequence", + _attempt, + initial_stage_timeout_seconds=(self._resolve_initial_stage_timeout_seconds()), + fallback_stage_timeout_seconds=timeout_seconds, + ) + else: + response_json, raw_response_text = await self._request_content_with_body( + body=rewrite_body, + timeout_seconds=timeout_seconds, + ) + if not raw_response_text: + raw_response_text = json.dumps( + response_json, + ensure_ascii=False, + ) + ( + rollout_id, + rollout_label, + rewritten_prompts, + ) = self._extract_rewrite_rollout_from_content( + raw_response_text, + expected_len, + preset_id=preset_id, + preset_label=preset_label, + ) + provider_name = self.provider_label + latency_ms = (time.perf_counter() - t0) * 1000.0 + return RewriteResult( + prompts=rewritten_prompts, + fallback_used=False, + error=None, + provider=provider_name, + model=resolved_rewrite_model, + latency_ms=latency_ms, + rollout_id=rollout_id, + rollout_label=rollout_label, + raw_response_text=raw_response_text, + ) + except Exception as exc: + latency_ms = (time.perf_counter() - t0) * 1000.0 + return RewriteResult( + prompts=cleaned_prompts, + fallback_used=True, + error=str(exc), + provider=self.provider_label, + model=resolved_rewrite_model, + latency_ms=latency_ms, + rollout_id=_resolve_rollout_id(preset_id), + rollout_label=_resolve_rollout_label(preset_label), + raw_response_text=raw_response_text, + ) diff --git a/apps/dreamverse/dreamverse/prompt_safety.py b/apps/dreamverse/dreamverse/prompt_safety.py new file mode 100644 index 000000000..7e88dca76 --- /dev/null +++ b/apps/dreamverse/dreamverse/prompt_safety.py @@ -0,0 +1,202 @@ +from __future__ import annotations +# mypy: ignore-errors + +import importlib +import os +import time +from dataclasses import dataclass +from functools import cache +from pathlib import Path + +_SERVER_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _SERVER_DIR.parent +_DEFAULT_CLASSIFIER_DIR = _REPO_ROOT / "classifiers" +CLASSIFIER_DIR = Path( + os.path.expandvars(os.path.expanduser(os.getenv("LTX2_CLASSIFIER_DIR", str(_DEFAULT_CLASSIFIER_DIR))))) + + +@dataclass(frozen=True) +class BlockedPrompt: + index: int + prompt: str + error: str + + +def _load_fasttext_module(): + try: + return importlib.import_module("fasttext") + except ImportError as exc: + raise RuntimeError("Prompt safety is enabled but the `fasttext` package is not installed.") from exc + + +def resolve_classifier_path( + classifier_kind: str, + env_var: str, + filename: str, + legacy_path: str, + shared_filename: str, +) -> str: + candidates: list[Path] = [] + env_path = os.getenv(env_var) + if env_path: + candidates.append(Path(os.path.expandvars(os.path.expanduser(env_path)))) + candidates.extend([ + CLASSIFIER_DIR / filename, + Path(f"/home/shared/{shared_filename}"), + Path(legacy_path), + ]) + + for candidate in candidates: + if candidate.is_file(): + return str(candidate) + + checked = "\n".join(f" - {candidate}" for candidate in candidates) + raise FileNotFoundError(f"Could not find the {classifier_kind} classifier.\n" + f"Checked:\n{checked}\n" + "Set LTX2_CLASSIFIER_DIR or the appropriate classifier path environment " + "variable.") + + +@cache +def load_fasttext_model(model_path: str): + return _load_fasttext_module().load_model(model_path) + + +def fasttext_predict( + model_path: str, + text: str, + classifier_name: str, +) -> tuple[str, float]: + model = load_fasttext_model(model_path) + text = text.replace("\n", " ") + start_time = time.perf_counter() + try: + labels, probs = model.predict(text) + except ValueError as error: + if "Unable to avoid copy while creating an array" not in str(error): + raise + predictions = model.f.predict(f"{text}\n", 1, 0.0, "strict") + if not predictions: + raise ValueError("fastText returned no predictions") from error + probs, labels = zip(*predictions, strict=False) + latency_ms = (time.perf_counter() - start_time) * 1000.0 + identifier = labels[0].replace("__label__", "") + confidence = float(probs[0]) + print("[safety] " + f"{classifier_name} fastText latency={latency_ms:.2f}ms " + f"label={identifier} confidence={confidence:.4f}") + return identifier, confidence + + +def _normalize_classifier_label(identifier: str) -> str: + return identifier.strip().lower().replace("-", "_").replace(" ", "_") + + +def _label_matches( + identifier: str, + blocked_markers: tuple[str, ...], + safe_markers: tuple[str, ...], +) -> bool: + normalized = _normalize_classifier_label(identifier) + tokens = tuple(token for token in normalized.split("_") if token) + + def has_marker(marker: str) -> bool: + normalized_marker = _normalize_classifier_label(marker) + return (normalized == normalized_marker or normalized_marker in tokens) + + if any(has_marker(marker) for marker in safe_markers): + return False + return any(has_marker(marker) for marker in blocked_markers) + + +class PromptSafetyFilter: + + def __init__( + self, + *, + nsfw_classifier=None, + hate_speech_classifier=None, + ): + if nsfw_classifier is None or hate_speech_classifier is None: + _load_fasttext_module() + + if nsfw_classifier is None: + nsfw_model_path = resolve_classifier_path( + "NSFW", + "LTX2_NSFW_CLASSIFIER_PATH", + "jigsaw_fasttext_bigrams_nsfw_final.bin", + "/data/classifiers/dolma_fasttext_nsfw_jigsaw_model.bin", + "dolma-jigsaw-fasttext-bigrams-nsfw-final.bin", + ) + self._nsfw_classifier = lambda text: fasttext_predict( + nsfw_model_path, + text, + "nsfw", + ) + else: + self._nsfw_classifier = nsfw_classifier + + if hate_speech_classifier is None: + hate_speech_model_path = resolve_classifier_path( + "hate speech", + "LTX2_HATESPEECH_CLASSIFIER_PATH", + "jigsaw_fasttext_bigrams_hatespeech_final.bin", + "/data/classifiers/dolma_fasttext_hatespeech_jigsaw_model.bin", + "dolma-jigsaw-fasttext-bigrams-hatespeech-final.bin", + ) + self._hate_speech_classifier = lambda text: fasttext_predict( + hate_speech_model_path, + text, + "hate_speech", + ) + else: + self._hate_speech_classifier = hate_speech_classifier + + def get_prompt_safety_error(self, prompt: str) -> str | None: + normalized_prompt = prompt.strip() + if not normalized_prompt: + return None + + nsfw_label, _ = self._nsfw_classifier(normalized_prompt) + if _label_matches( + nsfw_label, + blocked_markers=("nsfw", ), + safe_markers=("sfw", "safe"), + ): + return ("This prompt was flagged as NSFW. " + "You can't generate a video with this specified prompt.") + + hate_label, _ = self._hate_speech_classifier(normalized_prompt) + if _label_matches( + hate_label, + blocked_markers=("hatespeech", "hate", "toxic", "offensive", "abusive"), + safe_markers=( + "non_hatespeech", + "not_hatespeech", + "non_toxic", + "not_toxic", + "safe", + "clean", + ), + ): + return ("This prompt was flagged as hate speech. " + "You can't generate a video with this specified prompt.") + + return None + + def get_first_blocked_prompt( + self, + prompts: list[str], + ) -> BlockedPrompt | None: + for index, prompt in enumerate(prompts): + normalized_prompt = str(prompt or "").strip() + if not normalized_prompt: + continue + error = self.get_prompt_safety_error(normalized_prompt) + if error is not None: + return BlockedPrompt( + index=index, + prompt=normalized_prompt, + error=error, + ) + return None diff --git a/apps/dreamverse/dreamverse/prompts/README.md b/apps/dreamverse/dreamverse/prompts/README.md new file mode 100644 index 000000000..08929e746 --- /dev/null +++ b/apps/dreamverse/dreamverse/prompts/README.md @@ -0,0 +1,59 @@ +# Prompt Files + +This directory contains the editable system prompts used by the +`PromptEnhancer` in `server/prompt_enhancer.py`. + +## `next_segment_system_prompt.md` + +Use this prompt for guided continuation. + +It is loaded as the "next-segment" system prompt and used by +`PromptEnhancer.enhance_prompt()` for the normal continuation flow when the +request includes: + +- locked prior segments +- a new user conditioning prompt describing what should happen next + +In that path, the model is asked to write exactly one new segment prompt that +continues from the locked history while satisfying the user's requested next +beat. + +This is the prompt used for the live non-single-clip enhancement path in +`server/main.py`. + +## `auto_extension_system_prompt.md` + +Use this prompt for autonomous prompt expansion. + +It is loaded as the "auto-extension" system prompt and used in two different +paths: + +1. `PromptEnhancer.generate_auto_prompt()` + +This is the background auto-extension flow. The model only receives locked +segment history and must infer the next narrative beat on its own. + +1. `PromptEnhancer.enhance_prompt()` in single-clip mode + +This is the single-clip expansion flow. A short user idea is expanded into one +standalone detailed 5-second prompt for a single clip. + +## `rewrite_user_system_prompt.md` + +Use this prompt for new-rollout generation from a user's initial rollout +instruction. + +This prompt is used for the rewrite path when there is no existing prompt +window yet and the system needs to generate the initial 6 segment prompts from +scratch. + +If this file is missing or empty, the server falls back to +`rewrite_window_system_prompt.md` so startup does not break while the file is +still being drafted. + +## Practical difference + +- `next_segment_system_prompt.md` is for "continue based on the user's new + instruction." +- `auto_extension_system_prompt.md` is for "expand or infer the next clip + without that guided continuation input." diff --git a/apps/dreamverse/dreamverse/prompts/auto_extension_system_prompt.md b/apps/dreamverse/dreamverse/prompts/auto_extension_system_prompt.md new file mode 100644 index 000000000..62eb36d63 --- /dev/null +++ b/apps/dreamverse/dreamverse/prompts/auto_extension_system_prompt.md @@ -0,0 +1,52 @@ +You are a prompt extender for LTX-2.3 video generation. + +Your job is to expand a short user idea into a detailed cinematic prompt for a single 5-second video clip. + +LTX-2.3 works best when prompts clearly describe: +• the subject +• the action +• the environment +• lighting +• camera behavior +• audio + +Write the scene as one continuous cinematic shot. + +Guidelines: +- Preserve the user’s subject and intent. +- Add concrete visual details (appearance, materials, setting). +- Use cinematic language such as medium shot, close-up, slow push-in, pan, tracking shot. +- Describe motion using clear verbs and visible actions. +- Express emotion through physical cues rather than internal thoughts. +- If dialogue is included, put spoken lines in quotation marks and keep them short. +- Include simple audio when relevant (console beeps, footsteps, rain, room tone). +- End the prompt with a stable visual frame. + +Prompt structure (single paragraph, ~4–8 sentences): +1. Shot and subject +2. Environment and lighting +3. Main action +4. Small reaction or follow-up beat +5. Optional camera movement +6. Audio elements +7. Stable ending frame + +Avoid: +- scene cuts +- conflicting lighting +- overloaded scenes +- abstract emotional descriptions + +Return valid JSON only. + +Output exactly one JSON object with this schema: +{"prompt":""} + +Rules: +- The top-level JSON object must contain exactly one field: "prompt". +- "prompt" must be a single string. +- Do not return markdown fences. +- Do not return commentary, explanations, or any text before or after the JSON. +- Do not return "next_prompt". +- Do not return "segment_prompts". +- Do not return an array. diff --git a/apps/dreamverse/dreamverse/prompts/next_segment_system_prompt.md b/apps/dreamverse/dreamverse/prompts/next_segment_system_prompt.md new file mode 100644 index 000000000..c76be8307 --- /dev/null +++ b/apps/dreamverse/dreamverse/prompts/next_segment_system_prompt.md @@ -0,0 +1,105 @@ +You are a prompt writer for LTX-2 video continuation. You write one new +segment prompt that continues a video from where it left off, guided by +a user's conditioning prompt. + + +LTX-2 generates video in 6 sequential segments of 5 seconds each +(30 seconds total). Each segment is generated using the LAST FRAME of +the previous segment as its starting image. The model has no memory of +earlier segments - only that single frame. This means: +- Characters, objects, and settings that are visible in the last frame + carry forward naturally. +- Anything that left the frame (via a scene cut, hard transition, or + camera movement away) cannot be recreated - the model has never seen + it before. +- If a segment ends mid-action or with sudden motion, the next segment + inherits a blurry or unstable starting frame, which degrades quality. + + + +You will receive locked segments (already generated or currently +generating) and a conditioning prompt describing what should happen +next. Write exactly one new next-segment prompt continuing naturally +from the last locked segment. + + + + +- Place the user's requested event in this next segment. +- Complete the requested event within this single 5-second segment. +- Keep continuity from locked segments while satisfying the request. + + + +Write the segment as a single flowing paragraph in present tense using +active language ("is walking", "reaches for", "speaks softly"). + +Structure the segment with these layers: +1. Establish the shot using cinematography terms (medium shot, close-up, + wide establishing shot). +2. Set the scene: lighting, color palette, textures, atmosphere. +3. Describe action as a chronological sequence using temporal connectors + ("as", "then", "while"). +4. Define characters through observable features: age, hairstyle, + clothing, distinguishing marks. +5. Weave in an audio layer alongside the action - specific ambient + sounds ("the hum of fluorescent lights", "a clock ticking on the + wall"), effects, and speech integrated with the visual description. +6. Place all spoken dialogue in quotation marks. Preserve any dialogue + from the user's conditioning prompt exactly as written. + +Express emotion through physical cues (clenched fists, trembling lip, +wide eyes) rather than labels ("sad", "angry"). Describe only what is +seen and heard - no smell, taste, or internal thoughts. Use restrained, +natural phrasing. Start directly with scene description. + +Keep actions gradual. LTX-2 struggles with sudden, abrupt movements - +they produce artifacts and blurry frames. Any camera movement within +the segment should settle to a still frame by the end - the last moment +should be a stable, static shot. + + + +Static shots held across multiple consecutive segments can cause visual +artifacts to accumulate in video continuation. Changing the scene can +refresh the image and reset quality. + +There are two types of segments: + +1. Continuation segments (most segments): The segment continues from + the last frame of the previous segment. Write it as an image-to-video + prompt - describe only what changes from the previous scene. Keep + characters, setting, and framing consistent with what was already + on screen. + +2. Cut segments: The segment is an entirely new scene. Write it as a + full text-to-video prompt - fully describe the new setting, + characters, lighting, and atmosphere from scratch, as if the model + has never seen any of it before. Only use a cut when the conditioning + prompt naturally requires a scene change. + +For late segments (5-6), use only subtle camera movements (slow zoom, +gentle pan, slight drift) and keep the scene stable for a clean ending. + + + +When a segment lacks significant action or sound, fill the 5 seconds +with spoken dialogue in quotation marks to keep the scene engaging. +Characters should react to and reference the conditioning event in +their dialogue. Weave dialogue throughout the segment alongside action. + + + +Match the length of the locked segments. Count the sentences in the +locked segments and write about the same number here. Typically +3-5 sentences. If the locked segments are short, keep this one short. + + + + +Return valid JSON only: +{ + "next_prompt": "prompt for the next segment" +} +One key only, no markdown fences, no extra keys. + diff --git a/apps/dreamverse/dreamverse/prompts/rewrite_user_system_prompt.md b/apps/dreamverse/dreamverse/prompts/rewrite_user_system_prompt.md new file mode 100644 index 000000000..7a37fb6cd --- /dev/null +++ b/apps/dreamverse/dreamverse/prompts/rewrite_user_system_prompt.md @@ -0,0 +1,207 @@ +You are a real-time prompt writer for ltx2, a video generation model for image-audio-to-video continuation. + + +You receive: +1. A user prompt describing a video idea, scene, character moment, joke, action, or story beat + + + +Your job is to expand the user's prompt into a full rollout of 6 sequential segment prompts, each describing 5 seconds of video, for a total of 30 seconds. + +Each segment is generated independently but conditioned only on: +- the last 9 video frames of the previous segment +- the last 49 audio frames of the previous segment + +Important implications: +- Subjects, props, and scene elements that remain visible in the previous segment's last frame carry forward best. +- If a subject disappears from frame because of a hard transition or because the camera moves away, that subject should not reappear later unless the user explicitly wants a new reveal and that reveal is plausible from the current frame. +- Stable end frames improve continuation. +- The final sentence of each segment should land on a clean, readable visual state. +- Quiet continuing ambience in the final sentence is fine. +- Avoid ending a segment with a brand-new major action, a heavy new line of dialogue, or visual chaos. + + + +Return a complete original 6-segment rollout based on the user's prompt. + +You must: +- Expand the user's idea into a coherent beginning-to-end 30-second rollout. +- Preserve the user's core concept, tone, style, characters, setting, and requested actions. +- If the user gives only a broad or simple idea, infer missing details conservatively and add clear staging, scene logic, and pacing. +- If the user gives a detailed prompt, follow it closely. +- Keep all 6 segments coherent as one continuous rollout. +- Return all 6 segment prompts. + + + +Handle the user's prompt in one of these ways: + +1. Very broad prompt +- Expand into a clear, staged rollout with defined characters, setting, and progression. + +1. Moderately specific prompt +- Preserve given details and fill only what is needed for a complete rollout. + +1. Highly detailed prompt +- Follow closely while maintaining clarity, pacing, and continuity. + + + +When rules conflict, resolve them in this order: +1. The user's prompt +2. Continuation plausibility from one segment to the next +3. Clarity, staging, and visual quality +4. Default stylistic preferences in this system prompt + + + +Default to a dialogue-forward, character-centered rollout with clear staging, one readable beat per segment, and strong visual continuity. + +Core defaults: +- One main action + one reaction beat per segment +- Compact, legible segments +- One dominant speaker per segment +- Stable ending frames + +Conversation bias: +- When the prompt supports comedy, character interaction, or everyday scenarios, prioritize conversational beats over pure visual spectacle. +- Prefer dialogue-driven progression rather than action-only sequences when both are plausible. +- Use dialogue to reveal character, humor, tension, or situation changes. +- Keep exchanges short and punchy rather than long or dense. +- Let visual acting and timing complement the dialogue instead of replacing it. + +Exception: +- If the user explicitly asks for cinematic spectacle, action-heavy sequences, or minimal dialogue, follow that instead. + + + +Unless the user gives different timing, use this structure: +- Segment 1: establish scene, subject, situation, first speaking beat +- Segment 2: small escalation or reaction +- Segment 3: continuation beat +- Segment 4: pivot, reveal, pan, cut, or new subject focus +- Segment 5: payoff or response +- Segment 6: closing button, stable hold + +Segment 4 is the default pivot point unless specified otherwise. + + + +Treat counts as soft targets. Do not pad unnaturally. + +- Segment 1: 4–5 sentences, ~95–140 words +- Segment 2: 3–4 sentences, ~55–95 words +- Segment 3: 3–4 sentences, ~55–95 words +- Segment 4: 3–5 sentences, ~85–130 words +- Segment 5: 3–4 sentences, ~55–95 words +- Segment 6: 3–4 sentences, ~55–100 words + +Adjust if user intent requires. + + + +- Respect user-specified style. +- Use a "Style:" prefix only if clearly beneficial or already implied. +- Keep style consistent across segments. + + + +- Fully describe scene/subjects on first appearance. +- Compress repeated details in later segments. +- Maintain key identity anchors without redundancy. +- Reintroduce full detail only when scene or subject changes. + + + +Each segment must: +- Be present tense +- Describe only visible/audible elements +- Be one paragraph +- Match detail to shot scale +- End on a stable visual frame + + + +- Maintain coherent setting, layout, lighting +- Use concrete visual details and textures +- Avoid conflicting lighting or environment logic +- Prefer one primary location with optional pivot at segment 4 + + + +- Use consistent naming across segments +- Maintain appearance and identity anchors +- Introduce new characters clearly once +- Prefer small number of recurring subjects + + + +- Keep camera language efficient +- Prefer stable shots unless movement matters +- Max one clear camera move per segment +- Describe post-movement composition +- Use segment 4 for major camera transitions by default + + + +- One main action + one reaction beat +- Keep motion readable and grounded +- Favor simple, clear gestures +- Avoid chaotic or overloaded motion + + + +- Include dialogue in most segments when appropriate (typically 5–6 segments) +- One dominant speaker per segment +- 1 short line or 2 clipped mini-lines +- Keep lines natural and brief +- Place dialogue before final sentence when possible + + + +- Tie audio to visible action/environment +- Use 1–2 concrete sound cues per segment +- Maintain audio continuity across segments +- Avoid introducing new dominant sounds at the end + + + +- No internal thoughts +- Use visible cues for emotion + + + +- Ensure smooth visual/audio continuity across segments +- Do not reintroduce off-screen subjects unless justified +- Stabilize new scenes immediately after transitions +- End segments with stable compositions + + + +- Avoid reliance on readable text unless explicitly requested + + + +- Generate "id" in snake_case +- Generate concise descriptive "label" + + + +Return ONLY valid JSON using exactly this structure: + +{ + "id": "...", + "label": "...", + "segment_prompts": [ + "segment 1 text", + "segment 2 text", + "segment 3 text", + "segment 4 text", + "segment 5 text", + "segment 6 text" + ] +} + +Do not include explanations, markdown, or additional fields. +Only output the JSON object. + diff --git a/apps/dreamverse/dreamverse/prompts/rewrite_window_system_prompt.md b/apps/dreamverse/dreamverse/prompts/rewrite_window_system_prompt.md new file mode 100644 index 000000000..b964c4258 --- /dev/null +++ b/apps/dreamverse/dreamverse/prompts/rewrite_window_system_prompt.md @@ -0,0 +1,255 @@ +You are a real-time prompt editor for ltx2, a video generation model for image-audio-to-video continuation. + + +You receive: +1. An existing rollout JSON with: + - "id" + - "label" + - "segment_prompts": an array of 6 sequential segment prompts +2. The user's latest instruction about how to revise the rollout + + + +The rollout contains 6 sequential segment prompts, each describing 5 seconds of video, for a total of 30 seconds. + +Each segment is generated independently but conditioned only on: +- the last 9 video frames of the previous segment +- the last 49 audio frames of the previous segment + +Important implications: +- Subjects, props, and scene elements that remain visible in the previous segment's last frame carry forward best. +- If a subject disappears from frame because of a hard transition or because the camera moves away, that subject should not reappear later unless the user explicitly wants a new reveal and that reveal is plausible from the current frame. +- Stable end frames improve continuation. +- The final sentence of each segment should land on a clean, readable visual state. +- Quiet continuing ambience in the final sentence is fine. +- Avoid ending a segment with a brand-new major action, a heavy new line of dialogue, or visual chaos. + + + +Return a complete revised 6-segment rollout. + +You must: +- Preserve as much of the existing rollout as possible unless the user asked to change it or it causes logical inconsistency, continuity problems, or common-sense failure. +- Adjust any number of segments as needed so the full 6-segment rollout stays coherent. +- Preserve unaffected details, pacing, camera logic, setting, props, and subject identity unless the user explicitly changes them or they become inconsistent. +- Revise all dependent details when one attribute changes. +- Return all 6 segment prompts, even if only one segment changes. + + + +Handle the user's latest instruction in one of these ways: + +1. No actionable rollout instruction +- If the user's latest instruction does not actually request a change to the rollout, return the existing rollout unchanged. + +1. General instruction +- If the user's latest instruction is broad or high-level, expand it into a more detailed, coherent rollout while preserving existing details wherever possible. + +1. Detailed instruction +- If the user's latest instruction is specific and detailed, follow it closely while preserving continuity and physical plausibility. + + + +When rules conflict, resolve them in this order: +1. The user's latest instruction +2. Continuation plausibility from the previous segment's last visible frame and recent audio +3. Preservation of the existing rollout +4. Default stylistic preferences in this system prompt + + + +Default to a dialogue-forward, character-centered rollout with clear staging, one readable beat per segment, and strong visual continuity. + +Common defaults: +- Use one main action plus one follow-up reaction beat per segment. +- Keep most segments compact and legible. +- Include dialogue in most segments unless the user explicitly wants a quiet, purely visual, or action-only rollout. +- Usually keep one speaking subject dominant within a segment. +- Favor clean, stable ending images over flashy exits. + + + +Unless the user gives different timing, use this as the default 6-segment rhythm: +- Segment 1: establish the setting, the main subject, the core situation, and the first speaking beat +- Segment 2: small escalation, reaction, or new piece of information +- Segment 3: continue the situation with one more beat of action or reaction +- Segment 4: visual refresh, reveal, pivot, pan, cut, nearby location change, or new subject focus +- Segment 5: payoff, response, or aftermath in the refreshed composition +- Segment 6: closing button and stable held ending + +This is a default pattern, not a hard requirement. If the user specifies different timing, follow the user. + + + +Treat sentence and word counts as soft pacing targets, not hard quotas. Do not pad or compress unnaturally just to hit counts. + +Default pacing: +- Segment 1: usually 4 to 5 sentences, roughly 95 to 140 words +- Segment 2: usually 3 to 4 sentences, roughly 55 to 95 words +- Segment 3: usually 3 to 4 sentences, roughly 55 to 95 words +- Segment 4: usually 3 to 5 sentences, roughly 85 to 130 words +- Segment 5: usually 3 to 4 sentences, roughly 55 to 95 words +- Segment 6: usually 3 to 4 sentences, roughly 55 to 100 words + +If the user requests a slower, denser, faster, quieter, or more cinematic rollout, adjust these ranges as needed while preserving clarity. + + + +- Preserve the rollout's existing style-marker pattern whenever possible. +- If the existing rollout consistently starts segments with a style prefix such as "Style: ...", keep that pattern. +- If the existing rollout does not use a style prefix, do not add one unless the user explicitly asks for a style change or the rollout needs a new clear style cue. +- Keep the visual style consistent across all 6 segments unless the user explicitly changes it. + + + +- When a scene, subject, or important prop first appears, describe it clearly and specifically. +- In later segments within the same scene, compress repeated details and restate only the key anchors needed for continuity, identity, and image quality. +- Do not fully re-describe the same room, outfit, prop, or character in every segment unless the user explicitly wants that repetition. +- When a new location appears, treat that segment as a fresh introduction for the new location. +- When a new subject appears, describe that subject clearly on first appearance, then use stable shorthand afterward. + + + +Each segment prompt must: +- Be written in present tense. +- Describe only what is seen and heard. +- Avoid internal thoughts, abstract emotions, or motivations unless shown through visible or audible cues. +- Be a single flowing paragraph. +- Match the amount of detail to the shot scale. Close shots should emphasize facial detail, hands, fabric, texture, and subtle motion. Wide shots should emphasize layout, blocking, and readable movement. +- End in a stable, readable visual state. + + + +- Keep the setting, spatial layout, lighting logic, and important props coherent across segments unless the user changes them. +- When useful, include materials and textures such as glossy plastic, worn fabric, tiled floor, brushed metal, wet pavement, fingerprint-textured clay, or soft fur. +- Use concrete visual details that help the model stage the scene cleanly. +- Do not introduce conflicting lighting logic within the same scene. + + + +- Use the same noun for the same subject across all segments. +- Do not rename the same subject with synonyms in later segments. +- Keep appearance and wardrobe anchors consistent unless the user changes them. +- Include enough appearance detail to preserve identity, such as age, hairstyle, clothing, distinguishing features, body type, species, or surface detail when relevant. +- If a subject speaks or sings, keep voice traits consistent unless the user changes them. +- If there are multiple recurring subjects, distinguish them with stable identifiers. + + + +- Each segment should imply a clear framing or shot, but camera language should stay efficient. +- Use explicit camera movement only when it matters. +- Most non-pivot segments should use stable framing or gentle motion. +- Prefer at most one deliberate camera move per segment. +- Describe camera movement relative to the subject when useful. +- After a pan, push, pull, tilt, whip-pan, or cut, describe what the camera now lands on. +- If the existing rollout includes explicit camera-transition wording such as: + - "The camera whip-pans fast to the right..." + - "The camera slowly pans across..." + - "The camera pushes in..." + preserve that transition wording exactly and keep it in the same segment and same relative location unless the user explicitly asks to change it. +- Unless the user specifies otherwise, segment 4 is the preferred place for a visual refresh such as a pan, whip-pan, cut, reveal, nearby location shift, or new subject focus. + + + +- Keep motion readable and physically plausible. +- Prefer one main visible action plus one follow-up reaction beat per segment. +- Keep blocking simple and legible. +- Favor small clear gestures such as a head tilt, raised eyebrow, folding arms, looking down, stepping forward, crouching, shifting weight, or setting an object down. +- Avoid overloaded choreography, chaotic physics, or too many simultaneous actions unless the user explicitly wants that complexity. + + + +- Dialogue is allowed and often useful. +- For dialogue-forward rollouts, include at least one short spoken beat in most segments, usually 5 or 6 of the 6 segments, unless the user requests silence or a mostly nonverbal sequence. +- Usually keep one speaker dominant within a segment. +- Usually use 1 short quoted line or 2 clipped mini-lines by the same speaker. +- Keep spoken lines short, natural, and easy to act. +- Put all spoken dialogue in quotation marks. +- Prefer visible acting around the line, such as a glance, pause, grin, sigh, shrug, or gesture. +- Avoid long speeches and dense back-and-forth exchanges inside one segment. +- When possible, place dialogue before the final sentence so the segment can land on a stable visual ending. + + + +- Include audio when it helps the scene. +- Tie sound to visible action, speech, movement, or ongoing ambience. +- Usually 1 or 2 specific sound cues is enough for a segment. +- Favor concrete sounds such as hums, clicks, beeps, footsteps, water lapping, crickets, distant chatter, vent hiss, keyboard clacks, birds, sprinkler clicks, or soft music already present in the scene. +- Keep audio scene-appropriate and continuation-safe. +- Quiet ambient audio can continue into the final sentence, but avoid introducing a brand-new dominant sound at the very end. + + + +- Do not describe internal thoughts. +- Prefer visible and audible cues over abstract emotional labels. +- Use posture, expression, gaze, timing, breathing, hand movement, and vocal delivery instead of unsupported inner-state narration. + + + +- Later segments must follow naturally from what is plausibly visible and audible from the previous segment's ending. +- Do not reintroduce subjects that are no longer visible after a major transition unless the user explicitly requests it and the reintroduction is plausible. +- If a major scene transition occurs, re-establish the new scene clearly in that same segment so later segments can continue from it. +- Good ending images include a held pose, a settled camera, a quiet look, a character standing still, a character seated calmly, or a clean locked composition. +- Avoid ending on blur, sudden subject exit, unresolved camera motion, or a fresh unresolved event. + + + +- Avoid making readable text, signage, or logos the main point of the scene unless the user explicitly asks for it or the existing rollout already uses it successfully. +- If preserving existing readable text details, keep them short and simple. + + + +- Prefer the smallest set of edits that fully satisfies the user's latest instruction. +- Preserve the existing rollout's successful pacing, rhythm, and density unless the user explicitly asks to change them. +- Preserve the rollout's existing structural asymmetry when it is already working well, such as a fuller segment 1, a pivot or reveal around segment 4, and shorter compressed later segments. +- If a segment already contains a clear spoken beat and the user did not ask to remove dialogue, preserve dialogue in that segment or replace it with a similarly short spoken beat. +- Preserve which subject is dominant in each segment unless the user explicitly changes the focus. +- Preserve stable framing in non-pivot segments when possible. +- Preserve exact camera-transition wording and keep it in the same segment and same relative location unless the user explicitly asks to change it. +- Preserve shorthand description in later segments when earlier segments already established the scene, subject, and props clearly. +- When a user change affects one segment, first patch that segment and its immediate neighbors before rewriting the whole rollout. +- When one anchor changes, propagate only the downstream changes required for continuity, identity, and common sense. +- Prefer edits that preserve the final held image of each segment whenever possible. +- When the user asks for a stronger result such as funnier, sharper, warmer, or more dramatic, first strengthen dialogue, reaction beats, visible acting, and timing before adding new props, new characters, or larger scene changes. +- Preserve the rollout's existing tonal temperature unless the user explicitly asks to change it. +- Do not add new spectacle, extra dialogue, extra camera movement, or extra scene changes that the user did not request. + + + +Avoid: +- Fully re-describing the same character or room in every segment +- Internal emotional narration without visible cues +- Conflicting lighting logic +- Overloaded scenes with too many characters or actions +- Unclear subject naming +- Sudden unsupported reappearances +- Long monologues +- Ending on chaos, blur, or a fresh unresolved event + + + +- Output an "id" in snake_case. +- Output a short "label" that matches the current concept. +- If editing an existing rollout, preserve the existing id and label unless the user's change makes them inaccurate. +- If they become inaccurate, update them minimally. + + + +Return ONLY valid JSON using exactly this structure: + +{ + "id": "...", + "label": "...", + "segment_prompts": [ + "segment 1 text", + "segment 2 text", + "segment 3 text", + "segment 4 text", + "segment 5 text", + "segment 6 text" + ] +} + +Do not include explanations, markdown, or additional fields. +Only output the JSON object. + diff --git a/apps/dreamverse/dreamverse/rewrite_prompt_payload.py b/apps/dreamverse/dreamverse/rewrite_prompt_payload.py new file mode 100644 index 000000000..9dcdc836a --- /dev/null +++ b/apps/dreamverse/dreamverse/rewrite_prompt_payload.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import json +from typing import Any + +REWRITE_REQUEST_TEXT = ("Rewrite all segment prompts with improved continuity and cinematic detail. " + "Keep count and ordering identical.") +DEFAULT_REWRITE_SEGMENT_COUNT = 6 +REWRITE_MODE_NEW = "new_rollout" +REWRITE_MODE_EDIT_EXISTING = "edit_existing_rollout" +DEFAULT_REWRITE_ROLLOUT_ID = "current_rollout" +DEFAULT_REWRITE_ROLLOUT_LABEL = "Current rollout" + + +def normalize_prompt_window_prompts(values: list[Any] | None) -> list[str]: + if not isinstance(values, list): + return [] + normalized: list[str] = [] + for item in values: + if not isinstance(item, str): + continue + clean_item = item.strip() + if not clean_item: + continue + normalized.append(clean_item) + return normalized + + +def build_rewrite_user_payload( + *, + prompt_window_prompts: list[str], + preset_id: str | None = None, + preset_label: str | None = None, + rewrite_instruction: str | None = None, +) -> dict[str, Any]: + rollout_id = (preset_id or "").strip() or DEFAULT_REWRITE_ROLLOUT_ID + rollout_label = (preset_label or "").strip() or DEFAULT_REWRITE_ROLLOUT_LABEL + instruction = (rewrite_instruction.strip() if isinstance(rewrite_instruction, str) else "") + if len(prompt_window_prompts) == 0: + return { + "mode": REWRITE_MODE_NEW, + "request": REWRITE_REQUEST_TEXT, + "user_instruction": instruction, + "desired_segment_count": DEFAULT_REWRITE_SEGMENT_COUNT, + "rollout_id_hint": rollout_id, + "rollout_label_hint": rollout_label, + } + return { + "mode": REWRITE_MODE_EDIT_EXISTING, + "request": REWRITE_REQUEST_TEXT, + "user_instruction": instruction, + "current_rollout": { + "id": rollout_id, + "label": rollout_label, + "segment_prompts": list(prompt_window_prompts), + }, + } + + +def build_rewrite_request_body( + *, + system_prompt: str, + prompt_window_prompts: list[str], + preset_id: str | None, + preset_label: str | None, + rewrite_instruction: str | None, + model: str, + temperature: float, + max_completion_tokens: int, +) -> dict[str, Any]: + user_payload = build_rewrite_user_payload( + prompt_window_prompts=prompt_window_prompts, + preset_id=preset_id, + preset_label=preset_label, + rewrite_instruction=rewrite_instruction, + ) + return { + "model": + model, + "temperature": + temperature, + "max_completion_tokens": + max_completion_tokens, + "response_format": { + "type": "json_object" + }, + "messages": [ + { + "role": "system", + "content": system_prompt, + }, + { + "role": "user", + "content": json.dumps(user_payload, ensure_ascii=False), + }, + ], + } diff --git a/apps/dreamverse/dreamverse/session/__init__.py b/apps/dreamverse/dreamverse/session/__init__.py new file mode 100644 index 000000000..58acc1b43 --- /dev/null +++ b/apps/dreamverse/dreamverse/session/__init__.py @@ -0,0 +1,5 @@ +"""Per-WebSocket-connection code: state, controller, and helpers. + +Nothing in this package outlives a single client session — lifetime matches +one ``websocket.accept()`` to disconnect cycle. +""" diff --git a/apps/dreamverse/dreamverse/session/controller.py b/apps/dreamverse/dreamverse/session/controller.py new file mode 100644 index 000000000..1bea79ff9 --- /dev/null +++ b/apps/dreamverse/dreamverse/session/controller.py @@ -0,0 +1,1810 @@ +"""Per-WebSocket-connection session controller. + +Lift-and-wrap extraction of the original ``main.websocket_endpoint``. The +function body is preserved verbatim inside ``SessionController.run()``; +the only substitutions are references to ``runtime.`` rewritten +to ``self.``, and a local alias ``websocket = self.ws`` at the +top of ``run()`` so the body can keep using the ``websocket`` name it +originally bound as a function parameter. + +Captured locals, nested closures, and ``nonlocal`` declarations are +unchanged — this is a module-boundary move, not a state-semantics +rewrite. Decomposing the 1880-line ``run`` body into smaller methods +and extracting a ``SessionState`` dataclass is explicitly deferred until +a real trigger appears (new WS action, reproducible concurrency bug, +etc.). +""" +from __future__ import annotations +# pyright: reportArgumentType=false, reportMissingImports=false, reportMissingTypeArgument=false, reportOptionalMemberAccess=false +# ruff: noqa: SIM105 +# mypy: ignore-errors + +import asyncio +import time +import uuid +from typing import TYPE_CHECKING + +from fastapi import WebSocket, WebSocketDisconnect +from dreamverse.gpu_pool import GPUSlot +from dreamverse.session_init_image import cleanup_session_init_image, persist_session_init_image +from dreamverse.worker_ipc import MediaChunk, MediaComplete, MediaInit + +from dreamverse.config import ( + DEFAULT_MODEL_ID, + GENERATION_SEGMENT_CAP, + PROMPT_AUTO_SLEEP_MS, + PROMPT_AUTO_TIMEOUT_MS, + PROMPT_TIMEOUT_MS, + SESSION_TIMEOUT_SECONDS, +) +from dreamverse.rewrite_prompt_payload import normalize_prompt_window_prompts + +from dreamverse.session.messages import PromptSubmission, ReadyPrompt +from dreamverse.utils import ( + _main_print, + _resolve_generation_segment_cap, + PROMPT_EXTENSION_FAILURE_USER_MESSAGE, +) + +if TYPE_CHECKING: + from dreamverse.gpu_pool import GPUPool + from dreamverse.session_logger import SessionEventLogger + from dreamverse.prompt_enhancer import PromptEnhancer + from dreamverse.prompt_safety import PromptSafetyFilter + + +class SessionController: + """Runs one WebSocket session from accept() through disconnect.""" + + def __init__( + self, + ws: WebSocket, + gpu_pool: GPUPool, + prompt_enhancer: PromptEnhancer, + prompt_safety_filter: PromptSafetyFilter | None, + session_event_logger: SessionEventLogger | None, + ) -> None: + self.ws = ws + self.gpu_pool = gpu_pool + self.prompt_enhancer = prompt_enhancer + self.prompt_safety_filter = prompt_safety_filter + self.session_event_logger = session_event_logger + + async def run(self) -> None: + websocket = self.ws + # --- lift-and-wrap body begins --- + await websocket.accept() + + if self.gpu_pool is None: + await websocket.send_json({ + "type": "error", + "message": "GPU pool not initialized", + }) + await websocket.close(code=1011) + return + if self.prompt_enhancer is None: + await websocket.send_json({ + "type": "error", + "message": "Prompt enhancer not initialized", + }) + await websocket.close(code=1011) + return + + client_id = str(uuid.uuid4()) + print(f"Client {client_id[:8]} connected") + + async def log_event( + event: str, + payload: dict[str, object] | None = None, + ) -> None: + if self.session_event_logger is None: + return + try: + await self.session_event_logger.write_event( + event=event, + client_id=client_id, + payload=payload, + ) + except Exception as exc: + _main_print( + "WARN", + f"Failed to write session log ({event}): {exc}", + ) + + await log_event("ws_session_start") + + send_lock = asyncio.Lock() + stop_event = asyncio.Event() + + async def ws_send_json(payload: dict): + async with send_lock: + await websocket.send_json(payload) + + async def ws_send_bytes(payload: bytes): + async with send_lock: + await websocket.send_bytes(payload) + + def preview_text(text: str, limit: int = 180) -> str: + normalized = (text or "").replace("\n", "\\n") + if len(normalized) <= limit: + return normalized + return normalized[:limit] + "..." + + def get_prompt_safety_error(prompt: str) -> str | None: + if self.prompt_safety_filter is None: + return None + return self.prompt_safety_filter.get_prompt_safety_error(prompt) + + def get_first_blocked_prompt(prompts: list[str]): + if self.prompt_safety_filter is None: + return None + return self.prompt_safety_filter.get_first_blocked_prompt(prompts) + + # Send initial queue status + status = self.gpu_pool.get_status() + await ws_send_json({ + "type": "queue_status", + "position": status["queue_size"] + 1 if status["available_gpus"] == 0 else 0, + "total_gpus": status["total_gpus"], + "available_gpus": status["available_gpus"], + }) + + gpu_id: int | None = None + slot: GPUSlot | None = None + timeout_task: asyncio.Task | None = None + websocket_reader_task: asyncio.Task | None = None + prompt_worker_task: asyncio.Task | None = None + rewrite_seed_prompts_task: asyncio.Task | None = None + session_init_image = None + + async def session_timeout(): + """Close the session after timeout.""" + await asyncio.sleep(SESSION_TIMEOUT_SECONDS) + if stop_event.is_set(): + return + stop_event.set() + print(f"[GPU {gpu_id}] Session timeout for client {client_id[:8]}") + try: + await ws_send_json({ + "type": "session_timeout", + "message": f"Session expired after {SESSION_TIMEOUT_SECONDS} seconds", + }) + await websocket.close(code=1000, reason="Session timeout") + except Exception: + pass + + async def cancel_task(task: asyncio.Task | None): + if task is None or task.done(): + return + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + try: + # Wait for client init message. + init_data = {} + try: + init_data = await asyncio.wait_for(websocket.receive_json(), timeout=10.0) + except asyncio.TimeoutError: + init_data = {} + + init_type = init_data.get("type") + preset_id = init_data.get("preset_id") + preset_label = str(init_data.get("preset_label") or "").strip() + initial_rollout_prompt = str(init_data.get("initial_rollout_prompt") or "").strip() + enhancement_enabled = bool(init_data.get("enhancement_enabled", True)) + auto_extension_enabled = bool(init_data.get("auto_extension_enabled", False)) + loop_generation_enabled = bool(init_data.get("loop_generation_enabled", False)) + single_clip_mode = bool(init_data.get("single_clip_mode", False)) + rewrite_model = self.prompt_enhancer.resolve_rewrite_model(init_data.get("rewrite_model")) + rewrite_system_prompt_override = str(init_data.get("rewrite_window_system_prompt") or "").strip() + rewrite_user_system_prompt_override = str(init_data.get("rewrite_user_system_prompt") or "").strip() + rewrite_system_prompt = self.prompt_enhancer.resolve_rewrite_system_prompt(rewrite_system_prompt_override) + rewrite_timeout_ms = PROMPT_TIMEOUT_MS + rewrite_temperature = self.prompt_enhancer.resolve_rewrite_temperature(init_data.get("rewrite_temperature")) + + if init_type == "session_init_v2": + incoming_prompts = init_data.get("curated_prompts", []) + elif init_type == "select_model": + incoming_prompts = init_data.get("segment_prompts", []) + else: + incoming_prompts = init_data.get("segment_prompts", []) + + curated_prompts = [] + if isinstance(incoming_prompts, list): + for prompt in incoming_prompts: + if isinstance(prompt, str) and prompt.strip(): + curated_prompts.append(prompt.strip()) + + blocked_init_prompt = get_first_blocked_prompt(curated_prompts) + if blocked_init_prompt is not None: + message = ("Prompt safety filter blocked initial prompt " + f"{blocked_init_prompt.index + 1}. {blocked_init_prompt.error}") + _main_print( + "WARN", + f"Initial prompt blocked for client {client_id[:8]} " + f"index={blocked_init_prompt.index + 1} " + f"preview={preview_text(blocked_init_prompt.prompt)}", + ) + await ws_send_json({ + "type": "error", + "message": message, + }) + await websocket.close(code=1008, reason="Prompt blocked") + return + + try: + session_init_image = persist_session_init_image(init_data.get("initial_image")) + except ValueError as exc: + await ws_send_json({ + "type": "error", + "message": str(exc), + }) + await websocket.close(code=1003, reason="Invalid initial image") + return + + if preset_id: + print(f"Client {client_id[:8]} selected preset: {preset_id} " + f"label={preset_label or '(unset)'} " + f"({len(curated_prompts)} curated prompts, " + f"enhancement_enabled={enhancement_enabled}, " + f"auto_extension_enabled={auto_extension_enabled}, " + f"loop_generation_enabled={loop_generation_enabled}, " + f"single_clip_mode={single_clip_mode})") + if session_init_image is not None: + print(f"Client {client_id[:8]} uploaded initial image: " + f"{session_init_image.display_name}") + + # Acquire a GPU slot. + gpu_id, slot = await self.gpu_pool.acquire(client_id, websocket) + + # Start session timeout. + timeout_task = asyncio.create_task(session_timeout()) + + # Join the engine on this GPU. + await slot.join_user(client_id, model_id=DEFAULT_MODEL_ID) + + # Notify client they're connected to a GPU. + await ws_send_json({ + "type": "gpu_assigned", + "gpu_id": gpu_id, + "session_timeout": SESSION_TIMEOUT_SECONDS, + }) + await log_event( + "gpu_assigned", + { + "gpu_id": gpu_id, + }, + ) + + # Session queues and mutable state. + raw_prompt_queue: asyncio.Queue[PromptSubmission] = asyncio.Queue() + ready_prompt_queue: asyncio.Queue[ReadyPrompt] = asyncio.Queue() + + curated_idx = 0 + segment_idx = 0 + locked_segment_prompts: list[str] = [] + seed_prompt_memory: list[str] = list(curated_prompts) + generated_segment_count = 0 + generation_cap_blocked = False + auto_extension_blocked_segment_idx: int | None = None + prompt_sources_drained_logged = False + generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0) + pending_seed_reset = False + pending_seed_reset_reason = "" + pending_reset_conditioning = False + loop_iteration = 0 if generation_paused else 1 + force_curated_restart_segment = False + pending_simple_prompt_submission: PromptSubmission | None = None + single_clip_waiting_for_request = False + rollout_waiting_for_rewrite = False + initial_rollout_waiting_for_rewrite = generation_paused + segment_generation_active = False + rewrite_restart_pending = False + project_active = True + project_stream_started = False + pending_project_end = False + + def replace_session_init_image(initial_image_payload: object) -> None: + nonlocal session_init_image + next_session_init_image = persist_session_init_image(initial_image_payload) + previous_session_init_image = session_init_image + session_init_image = next_session_init_image + if previous_session_init_image is not None: + cleanup_session_init_image(previous_session_init_image) + + async def schedule_simple_generate_request(payload: dict[str, object]) -> None: + nonlocal preset_id + nonlocal preset_label + nonlocal enhancement_enabled + nonlocal auto_extension_enabled + nonlocal loop_generation_enabled + nonlocal single_clip_mode + nonlocal generation_paused + nonlocal pending_seed_reset + nonlocal pending_seed_reset_reason + nonlocal seed_prompt_memory + nonlocal curated_prompts + nonlocal pending_simple_prompt_submission + + raw_prompt = str(payload.get("prompt", "")).strip() + prompt_id = str(payload.get("prompt_id") or uuid.uuid4()) + if not raw_prompt: + await ws_send_json({ + "type": "error", + "message": "simple_generate requires a non-empty prompt", + }) + return + + blocked_error = get_prompt_safety_error(raw_prompt) + if blocked_error is not None: + _main_print( + "WARN", + f"Simple prompt blocked for client {client_id[:8]} " + f"preview={preview_text(raw_prompt)}", + ) + await log_event( + "prompt_blocked", + { + "kind": "user_raw", + "prompt_id": prompt_id, + "raw_prompt": raw_prompt, + "error": blocked_error, + }, + ) + await ws_send_json({ + "type": "error", + "message": blocked_error, + }) + return + + try: + replace_session_init_image(payload.get("initial_image")) + except ValueError as exc: + await ws_send_json({ + "type": "error", + "message": str(exc), + }) + return + + next_preset_id = str(payload.get("preset_id") or "").strip() + if next_preset_id: + preset_id = next_preset_id + next_preset_label = str(payload.get("preset_label") or "").strip() + if next_preset_label: + preset_label = next_preset_label + + next_enhancement_enabled = bool(payload.get("enhancement_enabled", True)) + + enhancement_enabled = next_enhancement_enabled + auto_extension_enabled = False + loop_generation_enabled = False + single_clip_mode = True + generation_paused = False + pending_seed_reset = True + pending_seed_reset_reason = "simple_generate" + seed_prompt_memory = ([] if next_enhancement_enabled else [raw_prompt]) + curated_prompts = list(seed_prompt_memory) + pending_simple_prompt_submission = (PromptSubmission( + prompt_id=prompt_id, + raw_prompt=raw_prompt, + created_at_s=time.time(), + ) if next_enhancement_enabled else None) + + await log_event( + "simple_generate", + { + "prompt_id": prompt_id, + "prompt": raw_prompt, + "preset_id": preset_id, + "enhancement_enabled": next_enhancement_enabled, + "has_initial_image": session_init_image is not None, + }, + ) + + if session_init_image is not None: + _main_print( + "INFO", f"Client {client_id[:8]} set simple initial image: " + f"{session_init_image.display_name}") + + if next_enhancement_enabled: + await ws_send_json({ + "type": "prompt_received", + "prompt_id": prompt_id, + "queue_depth": 1, + }) + + async def apply_project_init_payload(payload: dict[str, object]) -> bool: + nonlocal preset_id + nonlocal preset_label + nonlocal initial_rollout_prompt + nonlocal enhancement_enabled + nonlocal auto_extension_enabled + nonlocal loop_generation_enabled + nonlocal single_clip_mode + nonlocal generation_paused + nonlocal curated_prompts + nonlocal seed_prompt_memory + nonlocal curated_idx + nonlocal segment_idx + nonlocal locked_segment_prompts + nonlocal pending_seed_reset + nonlocal pending_seed_reset_reason + nonlocal pending_reset_conditioning + nonlocal force_curated_restart_segment + nonlocal generated_segment_count + nonlocal generation_cap_blocked + nonlocal auto_extension_blocked_segment_idx + nonlocal prompt_sources_drained_logged + nonlocal pending_simple_prompt_submission + nonlocal single_clip_waiting_for_request + nonlocal rollout_waiting_for_rewrite + nonlocal initial_rollout_waiting_for_rewrite + nonlocal rewrite_restart_pending + nonlocal loop_iteration + nonlocal rewrite_model + nonlocal rewrite_system_prompt + nonlocal rewrite_system_prompt_override + nonlocal rewrite_user_system_prompt_override + nonlocal rewrite_temperature + nonlocal project_active + nonlocal project_stream_started + nonlocal pending_project_end + + next_initial_rollout_prompt = str(payload.get("initial_rollout_prompt") or "").strip() + next_enhancement_enabled = bool(payload.get("enhancement_enabled", True)) + next_auto_extension_enabled = bool(payload.get("auto_extension_enabled", False)) + next_loop_generation_enabled = bool(payload.get("loop_generation_enabled", False)) + next_single_clip_mode = bool(payload.get("single_clip_mode", False)) + + next_preset_id = str(payload.get("preset_id") or "").strip() + if next_preset_id: + preset_id = next_preset_id + next_preset_label = str(payload.get("preset_label") or "").strip() + if next_preset_label: + preset_label = next_preset_label + + next_rewrite_model = self.prompt_enhancer.resolve_rewrite_model(payload.get("rewrite_model")) + next_rewrite_system_prompt_override = str(payload.get("rewrite_window_system_prompt") or "").strip() + next_rewrite_user_system_prompt_override = str(payload.get("rewrite_user_system_prompt") or "").strip() + next_rewrite_temperature = (self.prompt_enhancer.resolve_rewrite_temperature( + payload.get("rewrite_temperature"))) + + next_curated_prompts = [] + incoming_prompts = payload.get("curated_prompts", []) + if isinstance(incoming_prompts, list): + for prompt in incoming_prompts: + if isinstance(prompt, str) and prompt.strip(): + next_curated_prompts.append(prompt.strip()) + + blocked_init_prompt = get_first_blocked_prompt(next_curated_prompts) + if blocked_init_prompt is not None: + message = ("Prompt safety filter blocked initial prompt " + f"{blocked_init_prompt.index + 1}. " + f"{blocked_init_prompt.error}") + _main_print( + "WARN", + f"Project init prompt blocked for client {client_id[:8]} " + f"index={blocked_init_prompt.index + 1} " + f"preview={preview_text(blocked_init_prompt.prompt)}", + ) + await ws_send_json({ + "type": "error", + "message": message, + }) + return False + + try: + replace_session_init_image(payload.get("initial_image")) + except ValueError as exc: + await ws_send_json({ + "type": "error", + "message": str(exc), + }) + return False + + initial_rollout_prompt = next_initial_rollout_prompt + enhancement_enabled = next_enhancement_enabled + auto_extension_enabled = next_auto_extension_enabled + loop_generation_enabled = next_loop_generation_enabled + single_clip_mode = next_single_clip_mode + rewrite_model = next_rewrite_model + rewrite_system_prompt_override = (next_rewrite_system_prompt_override) + rewrite_user_system_prompt_override = (next_rewrite_user_system_prompt_override) + rewrite_system_prompt = ( + self.prompt_enhancer.resolve_rewrite_system_prompt(rewrite_system_prompt_override)) + rewrite_temperature = next_rewrite_temperature + + curated_prompts = list(next_curated_prompts) + seed_prompt_memory = list(next_curated_prompts) + curated_idx = 0 + segment_idx = 0 + locked_segment_prompts = [] + pending_seed_reset = False + pending_seed_reset_reason = "" + pending_reset_conditioning = True + force_curated_restart_segment = False + generated_segment_count = 0 + generation_cap_blocked = False + auto_extension_blocked_segment_idx = None + prompt_sources_drained_logged = False + pending_simple_prompt_submission = None + single_clip_waiting_for_request = False + rollout_waiting_for_rewrite = False + generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0) + initial_rollout_waiting_for_rewrite = generation_paused + rewrite_restart_pending = False + loop_iteration = 0 + project_active = True + project_stream_started = False + pending_project_end = False + + return True + + async def run_rewrite_seed_prompts( + rewrite_instruction: str, + rewrite_model_name: str, + rewrite_timeout_ms_value: int, + rewrite_temperature_value: float, + prompt_window_prompts: list[str] | None = None, + ): + nonlocal seed_prompt_memory + nonlocal curated_prompts + nonlocal preset_id + nonlocal preset_label + nonlocal rewrite_seed_prompts_task + nonlocal pending_seed_reset + nonlocal pending_seed_reset_reason + nonlocal generated_segment_count + nonlocal generation_cap_blocked + nonlocal generation_paused + nonlocal rollout_waiting_for_rewrite + nonlocal initial_rollout_waiting_for_rewrite + nonlocal rewrite_restart_pending + nonlocal rewrite_system_prompt + nonlocal rewrite_system_prompt_override + nonlocal rewrite_user_system_prompt_override + try: + snapshot_prompts = (list(prompt_window_prompts) if isinstance(prompt_window_prompts, list) + and len(prompt_window_prompts) > 0 else list(seed_prompt_memory)) + effective_system_prompt = (self.prompt_enhancer.resolve_rewrite_new_rollout_system_prompt( + rewrite_user_system_prompt_override or rewrite_system_prompt_override, ) + if len(snapshot_prompts) == 0 else rewrite_system_prompt) + rewrite_result = await self.prompt_enhancer.rewrite_prompt_sequence( + snapshot_prompts, + preset_id=preset_id, + preset_label=preset_label, + rewrite_instruction=rewrite_instruction, + rewrite_model=rewrite_model_name, + rewrite_temperature=rewrite_temperature_value, + timeout_ms=rewrite_timeout_ms_value, + system_prompt_override=effective_system_prompt, + ) + + if stop_event.is_set(): + return + + blocked_rewrite_prompt = get_first_blocked_prompt(rewrite_result.prompts) + if blocked_rewrite_prompt is not None: + error_message = ("Prompt safety filter blocked rewritten prompt " + f"{blocked_rewrite_prompt.index + 1}. " + f"{blocked_rewrite_prompt.error}") + await log_event( + "rewrite_blocked", + { + "kind": "seed_rewrite", + "rewrite_instruction": rewrite_instruction, + "model": rewrite_result.model, + "prompt_window_prompts": snapshot_prompts, + "rewritten_prompts": list(rewrite_result.prompts), + "blocked_prompt_index": blocked_rewrite_prompt.index, + "blocked_prompt": blocked_rewrite_prompt.prompt, + "error": error_message, + }, + ) + await ws_send_json({ + "type": "rewrite_seed_prompts_complete", + "fallback_used": True, + "error": error_message, + "model": rewrite_result.model, + "latency_ms": round(rewrite_result.latency_ms, 2), + }) + rewrite_restart_pending = False + return + + seed_prompt_memory = list(rewrite_result.prompts) + curated_prompts = list(seed_prompt_memory) + if not rewrite_result.fallback_used: + preset_id = rewrite_result.rollout_id + preset_label = rewrite_result.rollout_label + if initial_rollout_waiting_for_rewrite and not rewrite_result.fallback_used: + pending_seed_reset = True + pending_seed_reset_reason = "initial_rewrite" + generation_paused = False + initial_rollout_waiting_for_rewrite = False + await ws_send_json({ + "type": "generation_paused_updated", + "paused": generation_paused, + }) + if rollout_waiting_for_rewrite and not rewrite_result.fallback_used: + pending_seed_reset = True + pending_seed_reset_reason = "rewrite_rollout" + generated_segment_count = 0 + generation_cap_blocked = False + generation_paused = False + elif rewrite_restart_pending and not rewrite_result.fallback_used: + pending_seed_reset = True + pending_seed_reset_reason = "rewrite_during_generation" + generated_segment_count = 0 + generation_cap_blocked = False + generation_paused = False + rewrite_restart_pending = False + elif rewrite_restart_pending: + rewrite_restart_pending = False + await log_event( + "rewrite_done", + { + "kind": "seed_rewrite", + "rewrite_instruction": rewrite_instruction, + "latency_ms": round(rewrite_result.latency_ms, 2), + "response": rewrite_result.raw_response_text or "", + }, + ) + await ws_send_json({ + "type": "seed_prompts_updated", + "prompts": seed_prompt_memory, + "preset_id": preset_id, + "preset_label": preset_label, + "reason": "rewrite", + "fallback_used": rewrite_result.fallback_used, + "error": rewrite_result.error, + "model": rewrite_result.model, + "latency_ms": round(rewrite_result.latency_ms, 2), + "raw_llm_output": rewrite_result.raw_response_text, + "rewrite_instruction": rewrite_instruction, + }) + await ws_send_json({ + "type": "rewrite_seed_prompts_complete", + "fallback_used": rewrite_result.fallback_used, + "error": rewrite_result.error, + "preset_id": preset_id, + "preset_label": preset_label, + "model": rewrite_result.model, + "latency_ms": round(rewrite_result.latency_ms, 2), + }) + except asyncio.CancelledError: + rewrite_restart_pending = False + await log_event( + "rewrite_cancelled", + { + "kind": "seed_rewrite", + "rewrite_instruction": rewrite_instruction, + "rewrite_model": rewrite_model_name, + }, + ) + raise + except Exception as exc: + if stop_event.is_set(): + return + await log_event( + "rewrite_exception", + { + "kind": "seed_rewrite", + "rewrite_instruction": rewrite_instruction, + "rewrite_model": rewrite_model_name, + "error": str(exc), + }, + ) + await ws_send_json({ + "type": "rewrite_seed_prompts_complete", + "fallback_used": True, + "error": str(exc), + "model": rewrite_model_name, + "latency_ms": 0.0, + }) + rewrite_restart_pending = False + finally: + rewrite_seed_prompts_task = None + + async def websocket_reader_loop(): + nonlocal enhancement_enabled + nonlocal auto_extension_enabled + nonlocal loop_generation_enabled + nonlocal generation_paused + nonlocal pending_seed_reset + nonlocal pending_seed_reset_reason + nonlocal generated_segment_count + nonlocal generation_cap_blocked + nonlocal auto_extension_blocked_segment_idx + nonlocal seed_prompt_memory + nonlocal curated_prompts + nonlocal rewrite_seed_prompts_task + nonlocal rewrite_model + nonlocal rewrite_system_prompt + nonlocal rewrite_system_prompt_override + nonlocal rewrite_user_system_prompt_override + nonlocal rewrite_timeout_ms + nonlocal rewrite_temperature + nonlocal rollout_waiting_for_rewrite + nonlocal initial_rollout_waiting_for_rewrite + nonlocal rewrite_restart_pending + nonlocal project_active + nonlocal project_stream_started + nonlocal pending_project_end + while not stop_event.is_set(): + try: + data = await websocket.receive_json() + except WebSocketDisconnect: + loop_generation_enabled = False + stop_event.set() + return + except Exception as exc: + _main_print("INFO", f"Client {client_id[:8]} reader stopped: {exc}") + loop_generation_enabled = False + stop_event.set() + return + + msg_type = data.get("type") + if msg_type == "leave": + loop_generation_enabled = False + stop_event.set() + return + + if msg_type == "end_project_keep_session": + if (rewrite_seed_prompts_task is not None and not rewrite_seed_prompts_task.done()): + await cancel_task(rewrite_seed_prompts_task) + loop_generation_enabled = False + generation_paused = False + rollout_waiting_for_rewrite = False + initial_rollout_waiting_for_rewrite = False + rewrite_restart_pending = False + pending_project_end = True + continue + + if msg_type == "project_init_v1": + if project_active or pending_project_end: + await ws_send_json({ + "type": + "error", + "message": ("Current project is still ending. " + "Wait for the reset to finish before starting another project."), + }) + continue + init_applied = await apply_project_init_payload(data) + if not init_applied: + continue + if initial_rollout_waiting_for_rewrite: + await ws_send_json({ + "type": "rewrite_seed_prompts_started", + "model": rewrite_model, + }) + rewrite_seed_prompts_task = asyncio.create_task( + run_rewrite_seed_prompts( + initial_rollout_prompt, + rewrite_model, + rewrite_timeout_ms, + rewrite_temperature, + [], + )) + else: + pending_seed_reset = True + pending_seed_reset_reason = "project_init" + continue + + if not project_active: + if msg_type in { + "append_prompt", + "rewrite_seed_prompts", + "reset_to_seed_prompts", + "restart_generation", + }: + await ws_send_json({ + "type": "error", + "message": ("Start a new project before sending prompts."), + }) + continue + + if msg_type == "append_prompt": + if ((rollout_waiting_for_rewrite or initial_rollout_waiting_for_rewrite) + and not single_clip_mode): + await ws_send_json({ + "type": "error", + "message": ("Use Rewrite to create or restart the rollout."), + }) + continue + raw_prompt = str(data.get("prompt", "")).strip() + prompt_id = str(data.get("prompt_id") or uuid.uuid4()) + if not raw_prompt: + await ws_send_json({ + "type": "error", + "message": "append_prompt requires a non-empty prompt", + }) + continue + + blocked_error = get_prompt_safety_error(raw_prompt) + if blocked_error is not None: + _main_print( + "WARN", + f"Live prompt blocked for client {client_id[:8]} " + f"preview={preview_text(raw_prompt)}", + ) + await log_event( + "prompt_blocked", + { + "kind": "user_raw", + "prompt_id": prompt_id, + "raw_prompt": raw_prompt, + "error": blocked_error, + }, + ) + await ws_send_json({ + "type": "error", + "message": blocked_error, + }) + continue + + submission = PromptSubmission( + prompt_id=prompt_id, + raw_prompt=raw_prompt, + created_at_s=time.time(), + ) + await raw_prompt_queue.put(submission) + await log_event( + "append_prompt", + { + "prompt": raw_prompt, + }, + ) + await ws_send_json({ + "type": "prompt_received", + "prompt_id": prompt_id, + "queue_depth": raw_prompt_queue.qsize(), + }) + continue + + if msg_type == "simple_generate": + await schedule_simple_generate_request(data) + continue + + if msg_type == "set_enhancement": + enhancement_enabled = bool(data.get("enabled", enhancement_enabled)) + await ws_send_json({ + "type": "enhancement_updated", + "enabled": enhancement_enabled, + }) + continue + + if msg_type == "set_auto_extension": + was_enabled = auto_extension_enabled + auto_extension_enabled = bool(data.get("enabled", auto_extension_enabled)) + if not auto_extension_enabled: + auto_extension_blocked_segment_idx = None + _main_print( + "INFO", + f"Auto extension disabled for client {client_id[:8]} " + "(cleared_blocked_state=true)", + ) + else: + cleared = auto_extension_blocked_segment_idx is not None + auto_extension_blocked_segment_idx = None + _main_print( + "INFO", + f"Auto extension enabled for client {client_id[:8]} " + f"(was_enabled={was_enabled} " + f"cleared_blocked_state={cleared})", + ) + await ws_send_json({ + "type": "auto_extension_updated", + "enabled": auto_extension_enabled, + }) + continue + + if msg_type == "set_loop_generation": + loop_generation_enabled = bool(data.get("enabled", loop_generation_enabled)) + await ws_send_json({ + "type": "loop_generation_updated", + "enabled": loop_generation_enabled, + }) + continue + + if msg_type == "set_generation_paused": + generation_paused = bool(data.get("paused", generation_paused)) + await ws_send_json({ + "type": "generation_paused_updated", + "paused": generation_paused, + }) + continue + + if msg_type == "reset_to_seed_prompts": + pending_seed_reset = True + pending_seed_reset_reason = "manual_reset" + await ws_send_json({ + "type": "seed_prompts_reset_pending", + "seed_prompt_count": len(seed_prompt_memory), + }) + continue + + if msg_type == "restart_generation": + generated_segment_count = 0 + generation_cap_blocked = False + rollout_waiting_for_rewrite = False + pending_seed_reset = True + pending_seed_reset_reason = "cap_restart" + generation_paused = False + await ws_send_json({ + "type": + "generation_restarted", + "reason": + "cap_restart", + "segment_cap": + _resolve_generation_segment_cap( + single_clip_mode=single_clip_mode, + cap=GENERATION_SEGMENT_CAP, + ), + }) + continue + + if msg_type == "rewrite_seed_prompts": + rewrite_instruction = str(data.get("rewrite_instruction", "")).strip() + rewrite_model = self.prompt_enhancer.resolve_rewrite_model(data.get("rewrite_model")) + rewrite_system_prompt_override = str(data.get("rewrite_window_system_prompt") or "").strip() + rewrite_user_system_prompt_override = str(data.get("rewrite_user_system_prompt") or "").strip() + rewrite_system_prompt = self.prompt_enhancer.resolve_rewrite_system_prompt( + rewrite_system_prompt_override) + rewrite_temperature = self.prompt_enhancer.resolve_rewrite_temperature( + data.get("rewrite_temperature")) + if len(seed_prompt_memory) == 0 and not (initial_rollout_waiting_for_rewrite + and rewrite_instruction): + await ws_send_json({ + "type": "error", + "message": "No seed prompts available to rewrite.", + }) + continue + if rewrite_seed_prompts_task is not None and not rewrite_seed_prompts_task.done(): + await ws_send_json({ + "type": "error", + "message": "rewrite_seed_prompts already in progress.", + }) + continue + if (segment_generation_active and not single_clip_mode and not rollout_waiting_for_rewrite + and not initial_rollout_waiting_for_rewrite): + rewrite_restart_pending = True + + await ws_send_json({ + "type": "rewrite_seed_prompts_started", + "model": rewrite_model, + }) + requested_prompt_window = normalize_prompt_window_prompts(data.get("prompt_window_prompts")) + rewrite_seed_prompts_task = asyncio.create_task( + run_rewrite_seed_prompts( + rewrite_instruction, + rewrite_model, + rewrite_timeout_ms, + rewrite_temperature, + requested_prompt_window, + )) + continue + + if msg_type == "set_rewrite_model": + rewrite_model = self.prompt_enhancer.resolve_rewrite_model(data.get("rewrite_model")) + await ws_send_json({ + "type": "rewrite_model_updated", + "rewrite_model": rewrite_model, + }) + continue + + if msg_type == "set_rewrite_temperature": + rewrite_temperature = self.prompt_enhancer.resolve_rewrite_temperature( + data.get("rewrite_temperature")) + await ws_send_json({ + "type": "rewrite_temperature_updated", + "rewrite_temperature": rewrite_temperature, + }) + continue + + async def prompt_worker_loop(): + while not stop_event.is_set(): + try: + submission = await asyncio.wait_for(raw_prompt_queue.get(), timeout=0.1) + except asyncio.TimeoutError: + continue + _main_print("INFO", f"Received user prompt for enhancement: {submission.raw_prompt}") + prompt_id = submission.prompt_id + raw_prompt = submission.raw_prompt + blocked_raw_prompt_error = get_prompt_safety_error(raw_prompt) + if blocked_raw_prompt_error is not None: + _main_print( + "WARN", + f"User prompt blocked for client {client_id[:8]} " + f"prompt_id={prompt_id} preview={preview_text(raw_prompt)}", + ) + await log_event( + "prompt_blocked", + payload={ + "kind": "user_raw", + "prompt_id": prompt_id, + "raw_prompt": raw_prompt, + "error": blocked_raw_prompt_error, + }, + ) + await ws_send_json({ + "type": "error", + "message": blocked_raw_prompt_error, + }) + continue + await log_event( + "enhance_request", + { + "prompt_id": prompt_id, + "raw_prompt": raw_prompt, + "enhancement_enabled": enhancement_enabled, + "rewrite_model": rewrite_model, + }, + ) + if enhancement_enabled: + locked_snapshot = list(locked_segment_prompts) + next_segment_idx = len(locked_snapshot) + 1 + _main_print( + "INFO", + f"Enhancement context: next_segment={next_segment_idx} " + f"locked_count={len(locked_snapshot)}", + ) + await ws_send_json({ + "type": "prompt_enhancing", + "prompt_id": prompt_id, + }) + result = await self.prompt_enhancer.enhance_prompt( + raw_prompt, + locked_segments=locked_snapshot, + next_segment_idx=next_segment_idx, + preset_id=preset_id, + mode="single_clip" if single_clip_mode else "user_live", + model=rewrite_model, + timeout_ms=PROMPT_TIMEOUT_MS, + ) + final_prompt = result.prompt.strip() if result.prompt else "" + blocked_final_prompt_error = get_prompt_safety_error(final_prompt) + _main_print("INFO", f"Enhanced prompt: {final_prompt}") + await log_event( + "rewrite_done", + { + "kind": "enhance_prompt", + "latency_ms": round(result.latency_ms, 2), + "response": final_prompt, + }, + ) + if blocked_final_prompt_error is not None: + await log_event( + "prompt_blocked", + { + "kind": "user_enhanced", + "prompt_id": prompt_id, + "raw_prompt": raw_prompt, + "output_prompt": final_prompt, + "fallback_used": result.fallback_used, + "latency_ms": round(result.latency_ms, 2), + "provider": result.provider, + "model": result.model, + "error": blocked_final_prompt_error, + }, + ) + await ws_send_json({ + "type": "error", + "message": blocked_final_prompt_error, + }) + continue + if result.fallback_used or not final_prompt: + source = "user_enhancement_failed" + _main_print( + "WARN", + f"User prompt enhancement fallback: client={client_id[:8]} " + f"prompt_id={prompt_id} latency_ms={result.latency_ms:.2f} " + f"error={result.error}", + ) + await ws_send_json({ + "type": "prompt_fallback_used", + "prompt_id": prompt_id, + "prompt": final_prompt, + "source": source, + "latency_ms": round(result.latency_ms, 2), + "error": PROMPT_EXTENSION_FAILURE_USER_MESSAGE, + }) + # Enhancement is strict JSON-only; do not enqueue raw + # prompt when enhancement fails. + continue + else: + source = "user_enhanced" + await ws_send_json({ + "type": "prompt_ready", + "prompt_id": prompt_id, + "prompt": final_prompt, + "source": source, + "latency_ms": round(result.latency_ms, 2), + }) + await ready_prompt_queue.put( + ReadyPrompt( + prompt_id=prompt_id, + prompt=final_prompt, + source=source, + fallback_used=result.fallback_used, + loop_iteration=loop_iteration, + )) + else: + await ready_prompt_queue.put( + ReadyPrompt( + prompt_id=prompt_id, + prompt=raw_prompt, + source="user_raw", + fallback_used=False, + loop_iteration=loop_iteration, + )) + await ws_send_json({ + "type": "prompt_ready", + "prompt_id": prompt_id, + "prompt": raw_prompt, + "source": "user_raw", + "latency_ms": 0.0, + }) + + def queue_snapshot() -> dict[str, object]: + return { + "user_ready": ready_prompt_queue.qsize(), + "curated_remaining": max(len(curated_prompts) - curated_idx, 0), + "auto_enabled": auto_extension_enabled, + "auto_blocked_segment": auto_extension_blocked_segment_idx, + "loop_enabled": loop_generation_enabled, + "paused": generation_paused, + "seed_prompt_count": len(seed_prompt_memory), + "loop_iteration": loop_iteration, + } + + def drain_queue_nowait(queue: asyncio.Queue) -> int: + drained = 0 + while True: + try: + queue.get_nowait() + drained += 1 + except asyncio.QueueEmpty: + break + return drained + + async def enter_project_idle() -> None: + nonlocal curated_prompts + nonlocal seed_prompt_memory + nonlocal curated_idx + nonlocal segment_idx + nonlocal locked_segment_prompts + nonlocal generated_segment_count + nonlocal generation_cap_blocked + nonlocal auto_extension_blocked_segment_idx + nonlocal prompt_sources_drained_logged + nonlocal generation_paused + nonlocal pending_seed_reset + nonlocal pending_seed_reset_reason + nonlocal pending_reset_conditioning + nonlocal force_curated_restart_segment + nonlocal pending_simple_prompt_submission + nonlocal single_clip_waiting_for_request + nonlocal rollout_waiting_for_rewrite + nonlocal initial_rollout_waiting_for_rewrite + nonlocal rewrite_restart_pending + nonlocal initial_rollout_prompt + nonlocal project_active + nonlocal project_stream_started + nonlocal pending_project_end + + dropped_raw = drain_queue_nowait(raw_prompt_queue) + dropped_ready = drain_queue_nowait(ready_prompt_queue) + + curated_prompts = [] + seed_prompt_memory = [] + curated_idx = 0 + segment_idx = 0 + locked_segment_prompts = [] + generated_segment_count = 0 + generation_cap_blocked = False + auto_extension_blocked_segment_idx = None + prompt_sources_drained_logged = False + generation_paused = False + pending_seed_reset = False + pending_seed_reset_reason = "" + pending_reset_conditioning = True + force_curated_restart_segment = False + pending_simple_prompt_submission = None + single_clip_waiting_for_request = False + rollout_waiting_for_rewrite = False + initial_rollout_waiting_for_rewrite = False + rewrite_restart_pending = False + initial_rollout_prompt = "" + project_active = False + pending_project_end = False + + if project_stream_started: + project_stream_started = False + await ws_send_json({"type": "ltx2_stream_complete"}) + await log_event("ws_stream_complete") + + await ws_send_json({ + "type": "project_idle", + "dropped_user_raw_queue": dropped_raw, + "dropped_user_ready_queue": dropped_ready, + }) + + def pick_next_prompt_nowait() -> ReadyPrompt | None: + nonlocal curated_idx + try: + return ready_prompt_queue.get_nowait() + except asyncio.QueueEmpty: + pass + + if curated_idx < len(curated_prompts): + prompt_idx = curated_idx + prompt = curated_prompts[prompt_idx] + curated_idx += 1 + return ReadyPrompt( + prompt=prompt, + source="curated", + seed_prompt_index=prompt_idx, + loop_iteration=loop_iteration, + ) + + return None + + websocket_reader_task = asyncio.create_task(websocket_reader_loop()) + prompt_worker_task = asyncio.create_task(prompt_worker_loop()) + await asyncio.sleep(0) + + await ws_send_json({ + "type": "loop_generation_updated", + "enabled": loop_generation_enabled, + }) + await ws_send_json({ + "type": "generation_paused_updated", + "paused": generation_paused, + }) + if pending_project_end: + await enter_project_idle() + elif initial_rollout_waiting_for_rewrite: + await ws_send_json({ + "type": "rewrite_seed_prompts_started", + "model": rewrite_model, + }) + rewrite_seed_prompts_task = asyncio.create_task( + run_rewrite_seed_prompts( + initial_rollout_prompt, + rewrite_model, + rewrite_timeout_ms, + rewrite_temperature, + [], + )) + else: + project_stream_started = True + await ws_send_json({ + "type": + "ltx2_stream_start", + "total_segments": + len(curated_prompts), + "preset_id": + preset_id, + "stream_mode": + "av_fmp4", + "live_mode": + True, + "loop_generation_enabled": + loop_generation_enabled, + "loop_iteration": + loop_iteration, + "generation_segment_cap": + _resolve_generation_segment_cap( + single_clip_mode=single_clip_mode, + cap=GENERATION_SEGMENT_CAP, + ), + }) + await ws_send_json({ + "type": "seed_prompts_updated", + "prompts": seed_prompt_memory, + "preset_id": preset_id, + "preset_label": preset_label, + "reason": "init", + "seed_prompt_count": len(seed_prompt_memory), + }) + + av_event_queue = slot.register_stream_queue(client_id) + + while not stop_event.is_set(): + if pending_project_end: + await enter_project_idle() + await asyncio.sleep(0) + continue + + if not project_active: + await asyncio.sleep(0.05) + continue + + if pending_seed_reset: + nonlocal_reason = pending_seed_reset_reason + curated_prompts = list(seed_prompt_memory) + curated_idx = 0 + segment_idx = 0 + locked_segment_prompts = [] + auto_extension_blocked_segment_idx = None + prompt_sources_drained_logged = False + pending_seed_reset = False + pending_reset_conditioning = True + force_curated_restart_segment = len(curated_prompts) > 0 + single_clip_waiting_for_request = False + dropped_raw = drain_queue_nowait(raw_prompt_queue) + dropped_ready = drain_queue_nowait(ready_prompt_queue) + + loop_iteration += 1 + project_stream_started = True + await ws_send_json({ + "type": + "ltx2_stream_start", + "total_segments": + len(curated_prompts), + "preset_id": + preset_id, + "stream_mode": + "av_fmp4", + "live_mode": + True, + "loop_generation_enabled": + loop_generation_enabled, + "loop_iteration": + loop_iteration, + "generation_segment_cap": + _resolve_generation_segment_cap( + single_clip_mode=single_clip_mode, + cap=GENERATION_SEGMENT_CAP, + ), + }) + if nonlocal_reason == "loop_restart": + await ws_send_json({ + "type": "loop_restarted", + "loop_iteration": loop_iteration, + "seed_prompt_count": len(seed_prompt_memory), + "dropped_user_raw_queue": dropped_raw, + "dropped_user_ready_queue": dropped_ready, + }) + else: + await ws_send_json({ + "type": "seed_prompts_reset_applied", + "reason": nonlocal_reason or "manual_reset", + "loop_iteration": loop_iteration, + "seed_prompt_count": len(seed_prompt_memory), + "dropped_user_raw_queue": dropped_raw, + "dropped_user_ready_queue": dropped_ready, + }) + if nonlocal_reason == "cap_restart": + generation_cap_blocked = False + generated_segment_count = 0 + rollout_waiting_for_rewrite = False + pending_seed_reset_reason = "" + + if pending_simple_prompt_submission is not None: + await raw_prompt_queue.put(pending_simple_prompt_submission) + pending_simple_prompt_submission = None + + if (not single_clip_mode and not generation_cap_blocked and not rollout_waiting_for_rewrite + and GENERATION_SEGMENT_CAP > 0 and generated_segment_count >= GENERATION_SEGMENT_CAP): + loop_generation_enabled = False + rollout_waiting_for_rewrite = True + _main_print( + "INFO", + f"Segment cap reached for client {client_id[:8]} " + f"(cap_segments={GENERATION_SEGMENT_CAP}, " + f"generated_segments={generated_segment_count}); " + "waiting for rollout rewrite", + ) + project_stream_started = False + await ws_send_json({"type": "ltx2_stream_complete"}) + await log_event("ws_stream_complete") + + if generation_cap_blocked: + await asyncio.sleep(0.05) + continue + + if rollout_waiting_for_rewrite: + await asyncio.sleep(0.05) + continue + + if initial_rollout_waiting_for_rewrite: + await asyncio.sleep(0.05) + continue + + if generation_paused: + await asyncio.sleep(0.05) + continue + + if rewrite_restart_pending: + if (rewrite_seed_prompts_task is not None and not rewrite_seed_prompts_task.done()): + await asyncio.sleep(0.05) + continue + rewrite_restart_pending = False + + if force_curated_restart_segment and curated_idx < len(curated_prompts): + prompt_idx = curated_idx + prompt = curated_prompts[prompt_idx] + curated_idx += 1 + selected = ReadyPrompt( + prompt=prompt, + source="curated", + seed_prompt_index=prompt_idx, + loop_iteration=loop_iteration, + ) + force_curated_restart_segment = False + else: + force_curated_restart_segment = False + selected = pick_next_prompt_nowait() + if selected is None and loop_generation_enabled and len(seed_prompt_memory) > 0: + pending_seed_reset = True + pending_seed_reset_reason = "loop_restart" + await asyncio.sleep(0) + continue + + if selected is None and auto_extension_enabled: + next_segment = len(locked_segment_prompts) + 1 + if auto_extension_blocked_segment_idx == next_segment: + if not prompt_sources_drained_logged: + snapshot = queue_snapshot() + _main_print( + "WARN", + f"Prompt sources drained for client {client_id[:8]} " + f"at segment boundary {segment_idx + 1} " + f"(user_ready={snapshot['user_ready']} " + f"curated_remaining={snapshot['curated_remaining']} " + f"auto_enabled={snapshot['auto_enabled']} " + "auto_state=blocked_after_failed_attempt)", + ) + prompt_sources_drained_logged = True + await asyncio.sleep(PROMPT_AUTO_SLEEP_MS / 1000.0) + continue + + locked_snapshot = list(locked_segment_prompts) + last_locked_preview = (preview_text(locked_snapshot[-1]) if locked_snapshot else "(none)") + _main_print( + "INFO", + f"Auto prompt generation request: client={client_id[:8]} " + f"next_segment={next_segment} " + f"locked_count={len(locked_snapshot)} " + f"last_locked_preview={last_locked_preview}", + ) + result = await self.prompt_enhancer.generate_auto_prompt( + locked_segments=locked_snapshot, + next_segment_idx=next_segment, + model=rewrite_model, + timeout_ms=PROMPT_AUTO_TIMEOUT_MS, + ) + _main_print( + "INFO", + f"Auto prompt generation response: client={client_id[:8]} " + f"next_segment={next_segment} " + f"fallback_used={result.fallback_used} " + f"returned_len={len(result.prompt)} " + f"latency_ms={result.latency_ms:.2f} error={result.error} " + f"returned_preview={preview_text(result.prompt)}", + ) + auto_prompt = result.prompt.strip() if result.prompt else "" + blocked_auto_prompt_error = get_prompt_safety_error(auto_prompt) + if result.fallback_used or not auto_prompt: + auto_extension_blocked_segment_idx = next_segment + _main_print( + "WARN", + f"Auto prompt generation failed: client={client_id[:8]} " + f"next_segment={next_segment} " + f"latency_ms={result.latency_ms:.2f} " + f"error={result.error or 'Auto prompt generation returned no prompt.'} " + "policy=blocked_no_retry_until_external_action", + ) + await ws_send_json({ + "type": "auto_prompt_failed", + "segment_idx": next_segment, + "latency_ms": round(result.latency_ms, 2), + "error": result.error or "Auto prompt generation returned no prompt.", + }) + await asyncio.sleep(PROMPT_AUTO_SLEEP_MS / 1000.0) + continue + if blocked_auto_prompt_error is not None: + auto_extension_blocked_segment_idx = next_segment + await log_event( + "auto_prompt_blocked", + { + "next_segment_idx": next_segment, + "output_prompt": auto_prompt, + "error": blocked_auto_prompt_error, + "latency_ms": round(result.latency_ms, 2), + "provider": result.provider, + "model": result.model, + }, + ) + await ws_send_json({ + "type": "auto_prompt_failed", + "segment_idx": next_segment, + "latency_ms": round(result.latency_ms, 2), + "error": blocked_auto_prompt_error, + }) + await asyncio.sleep(PROMPT_AUTO_SLEEP_MS / 1000.0) + continue + + auto_extension_blocked_segment_idx = None + selected = ReadyPrompt( + prompt=auto_prompt, + source="auto_llm", + prompt_id=None, + fallback_used=False, + loop_iteration=loop_iteration, + ) + + if selected is None: + if single_clip_mode: + await asyncio.sleep(PROMPT_AUTO_SLEEP_MS / 1000.0) + continue + if not prompt_sources_drained_logged: + snapshot = queue_snapshot() + _main_print( + "WARN", + f"Prompt sources drained for client {client_id[:8]} " + f"at segment boundary {segment_idx + 1} " + f"(user_ready={snapshot['user_ready']} " + f"curated_remaining={snapshot['curated_remaining']} " + f"auto_enabled={snapshot['auto_enabled']} " + f"auto_blocked_segment={snapshot['auto_blocked_segment']})", + ) + prompt_sources_drained_logged = True + await ws_send_json({ + "type": "prompt_sources_blocked", + "segment_idx": segment_idx + 1, + }) + await asyncio.sleep(PROMPT_AUTO_SLEEP_MS / 1000.0) + continue + + if prompt_sources_drained_logged: + prompt_sources_drained_logged = False + _main_print( + "INFO", + f"Prompt source available again for client {client_id[:8]} " + f"at segment boundary {segment_idx + 1} " + f"(source={selected.source} prompt_id={selected.prompt_id})", + ) + await ws_send_json({ + "type": "prompt_sources_resumed", + "segment_idx": segment_idx + 1, + "source": selected.source, + "prompt_id": selected.prompt_id, + }) + + segment_idx += 1 + single_clip_waiting_for_request = False + total_segments_hint = max(segment_idx, len(curated_prompts)) + prompt = selected.prompt + locked_segment_prompts.append(prompt) + if (auto_extension_blocked_segment_idx is not None + and auto_extension_blocked_segment_idx <= segment_idx): + auto_extension_blocked_segment_idx = None + + await ws_send_json({ + "type": "segment_prompt_source", + "segment_idx": segment_idx, + "source": selected.source, + "prompt_id": selected.prompt_id, + "fallback_used": selected.fallback_used, + "seed_prompt_index": selected.seed_prompt_index, + "loop_iteration": selected.loop_iteration or loop_iteration, + }) + await ws_send_json({ + "type": "ltx2_segment_start", + "segment_idx": segment_idx, + "total_segments": total_segments_hint, + "prompt": prompt, + "source": selected.source, + "seed_prompt_index": selected.seed_prompt_index, + "loop_iteration": selected.loop_iteration or loop_iteration, + }) + await log_event( + "segment_start", + { + "segment_idx": segment_idx, + }, + ) + + t_start = time.perf_counter() + timings: dict = {} + av_chunks_relayed = 0 + av_bytes_relayed = 0 + av_streamed = False + + step_reset_conditioning = pending_reset_conditioning + pending_reset_conditioning = False + step_image_path = (str(session_init_image.file_path) + if segment_idx == 1 and session_init_image is not None else None) + step_task = asyncio.create_task( + slot.user_step( + client_id, + prompt=prompt, + segment_idx=segment_idx, + image_path=step_image_path, + reset_conditioning=step_reset_conditioning, + )) + segment_generation_active = True + try: + while not stop_event.is_set() and (not step_task.done() or not av_event_queue.empty()): + try: + event = await asyncio.wait_for(av_event_queue.get(), timeout=0.05) + except asyncio.TimeoutError: + continue + + if event.segment_idx != segment_idx: + print(f"[GPU {gpu_id}] Ignoring out-of-order AV event: " + f"event_seg={event.segment_idx}, current_seg={segment_idx}, " + f"kind={type(event).__name__}") + continue + + match event: + case MediaInit(mime=mime, stream_id=stream_id): + av_streamed = True + await ws_send_json({ + "type": "media_init", + "segment_idx": segment_idx, + "mime": mime, + "stream_id": stream_id, + "mode": "av_fmp4", + }) + case MediaChunk( + chunk=chunk_bytes, + chunk_offset=chunk_offset, + chunk_length=chunk_length, + uses_shared_buffer=uses_shared, + ): + if (uses_shared and chunk_offset is not None and chunk_length is not None + and slot.shared_stream_buffer is not None): + start = chunk_offset + end = start + chunk_length + # Copy out of shared buffer for websocket send. + chunk = bytes(slot.shared_stream_buffer[start:end]) + else: + chunk = chunk_bytes or b"" + if chunk: + av_chunks_relayed += 1 + av_bytes_relayed += len(chunk) + await ws_send_bytes(chunk) + case MediaComplete(stream_id=stream_id, chunks=n): + await ws_send_json({ + "type": "media_segment_complete", + "segment_idx": segment_idx, + "stream_id": stream_id, + "chunks": n if n is not None else av_chunks_relayed, + }) + case _: + print(f"[GPU {gpu_id}] Unknown AV event: " + f"{type(event).__name__}") + + if not step_task.done(): + step_task.cancel() + else: + timings = await step_task + finally: + segment_generation_active = False + if not step_task.done(): + try: + await step_task + except asyncio.CancelledError: + pass + + if stop_event.is_set(): + break + + t_generation = time.perf_counter() - t_start + worker_e2e_ms = float(timings.get("e2e_latency_ms", 0.0) or 0.0) + main_user_step_ms = t_generation * 1000.0 + overhead_vs_worker_ms = main_user_step_ms - worker_e2e_ms + + ipc_put_start_ns = timings.get("ipc_put_start_ns") + ipc_get_done_ns = timings.get("ipc_get_done_ns") + ipc_queue_transfer_ms = None + if isinstance(ipc_put_start_ns, int) and isinstance(ipc_get_done_ns, int): + ipc_queue_transfer_ms = (ipc_get_done_ns - ipc_put_start_ns) / 1_000_000.0 + + ipc_queue_str = (f"{ipc_queue_transfer_ms:.0f}ms" if ipc_queue_transfer_ms is not None else "n/a") + print(f"[GPU {gpu_id}] Segment {segment_idx}: " + f"source={selected.source}, " + f"worker_e2e={worker_e2e_ms:.0f}ms, " + f"main_user_step={main_user_step_ms:.0f}ms, " + f"overhead={overhead_vs_worker_ms:.0f}ms, " + f"ipc_queue={ipc_queue_str}") + + if not av_streamed: + raise RuntimeError(f"Segment {segment_idx} AV stream did not initialize " + "(no media_init event)") + generated_segment_count += 1 + + print(f"[GPU {gpu_id}] Segment {segment_idx}: " + f"relayed av chunks={av_chunks_relayed}, " + f"bytes={av_bytes_relayed / (1024 * 1024):.1f}MB") + _main_print( + "INFO", + f"Segments generated for client {client_id[:8]}: " + f"{generated_segment_count}", + ) + await log_event( + "segment_complete", + { + "segment_idx": segment_idx, + "latency_ms": { + "total": round(main_user_step_ms, 2), + "worker_e2e": round(worker_e2e_ms, 2), + "main_user_step": round(main_user_step_ms, 2), + "overhead": round(overhead_vs_worker_ms, 2), + }, + "data_size_bytes": av_bytes_relayed, + }, + ) + await ws_send_json({ + "type": "step_complete", + "latency_ms": { + "total": round(main_user_step_ms, 2), + "worker_e2e": round(worker_e2e_ms, 2), + "main_user_step": round(main_user_step_ms, 2), + "overhead": round(overhead_vs_worker_ms, 2), + }, + }) + await ws_send_json({ + "type": "ltx2_segment_complete", + "segment_idx": segment_idx, + "total_segments": total_segments_hint, + }) + if single_clip_mode and not single_clip_waiting_for_request: + single_clip_waiting_for_request = True + project_stream_started = False + await ws_send_json({"type": "ltx2_stream_complete"}) + await log_event("ws_stream_complete") + _main_print( + "INFO", + f"Single-clip session idle for client {client_id[:8]} " + f"after segment {segment_idx}; waiting for next create", + ) + + if not stop_event.is_set() and project_stream_started: + project_stream_started = False + await ws_send_json({"type": "ltx2_stream_complete"}) + await log_event("ws_stream_complete") + print(f"[GPU {gpu_id}] LTX2 streaming complete for {client_id[:8]}") + + except WebSocketDisconnect: + stop_event.set() + _main_print("INFO", f"Client {client_id[:8]} disconnected") + except Exception as exc: + stop_event.set() + _main_print("ERROR", f"Client {client_id[:8]} error: {exc}") + try: + await ws_send_json({ + "type": "error", + "message": f"AV streaming failed: {exc}", + }) + except Exception: + pass + import traceback + traceback.print_exc() + finally: + stop_event.set() + await cancel_task(websocket_reader_task) + await cancel_task(prompt_worker_task) + await cancel_task(rewrite_seed_prompts_task) + + if slot is not None and client_id: + slot.unregister_stream_queue(client_id) + + if timeout_task and not timeout_task.done(): + timeout_task.cancel() + try: + await timeout_task + except asyncio.CancelledError: + pass + + try: + if client_id and self.gpu_pool is not None: + await self.gpu_pool.release(client_id) + finally: + cleanup_session_init_image(session_init_image) diff --git a/apps/dreamverse/dreamverse/session/messages.py b/apps/dreamverse/dreamverse/session/messages.py new file mode 100644 index 000000000..987f3f4ca --- /dev/null +++ b/apps/dreamverse/dreamverse/session/messages.py @@ -0,0 +1,21 @@ +"""Dataclasses for the prompt-submission pipeline queues.""" +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class PromptSubmission: + prompt_id: str + raw_prompt: str + created_at_s: float + + +@dataclass +class ReadyPrompt: + prompt: str + source: str + prompt_id: str | None = None + fallback_used: bool = False + seed_prompt_index: int | None = None + loop_iteration: int | None = None diff --git a/apps/dreamverse/dreamverse/session_init_image.py b/apps/dreamverse/dreamverse/session_init_image.py new file mode 100644 index 000000000..5ad4ae8f1 --- /dev/null +++ b/apps/dreamverse/dreamverse/session_init_image.py @@ -0,0 +1,101 @@ +# pyright: reportMissingTypeArgument=false, reportArgumentType=false +from __future__ import annotations + +import base64 +import io +import re +import shutil +import tempfile +from dataclasses import dataclass +from pathlib import Path + +from PIL import Image + +MAX_SESSION_INIT_IMAGE_BYTES = 15 * 1024 * 1024 +SUPPORTED_SESSION_INIT_IMAGE_MIME_TYPES = { + "image/jpeg", + "image/png", + "image/webp", +} + +_DATA_URL_RE = re.compile(r"^data:(?P[-\w.+/]+);base64,(?P[A-Za-z0-9+/=\s]+)$") + + +@dataclass(frozen=True) +class SessionInitImage: + file_path: Path + temp_dir: Path + display_name: str + mime_type: str + + +def _sanitize_display_name(raw_name: object) -> str: + text = str(raw_name or "").strip() + if not text: + return "uploaded-image" + return Path(text).name or "uploaded-image" + + +def persist_session_init_image( + payload: object, + *, + temp_root: Path | None = None, +) -> SessionInitImage | None: + if payload is None: + return None + if not isinstance(payload, dict): + raise ValueError("initial_image must be an object.") + + data_url = str(payload.get("data_url") or "").strip() + if not data_url: + return None + + data_url_match = _DATA_URL_RE.match(data_url) + if data_url_match is None: + raise ValueError("initial_image.data_url must be a base64 data URL.") + + data_url_mime_type = data_url_match.group("mime").strip().lower() + mime_type = str(payload.get("mime_type") or data_url_mime_type).strip().lower() + if mime_type != data_url_mime_type: + raise ValueError("initial_image.mime_type must match the data URL mime type.") + if mime_type not in SUPPORTED_SESSION_INIT_IMAGE_MIME_TYPES: + raise ValueError("initial_image must be a PNG, JPEG, or WebP image.") + + encoded_bytes = re.sub(r"\s+", "", data_url_match.group("data")) + try: + raw_bytes = base64.b64decode(encoded_bytes, validate=True) + except Exception as exc: + raise ValueError("initial_image.data_url is not valid base64 data.") from exc + + if len(raw_bytes) == 0: + raise ValueError("initial_image.data_url did not contain image bytes.") + if len(raw_bytes) > MAX_SESSION_INIT_IMAGE_BYTES: + raise ValueError("initial_image must be 15 MB or smaller.") + + try: + with Image.open(io.BytesIO(raw_bytes)) as image: + image.load() + normalized = image.convert("RGBA" if "A" in image.getbands() else "RGB") + except Exception as exc: + raise ValueError("initial_image must decode as a valid image.") from exc + + temp_dir = Path( + tempfile.mkdtemp( + prefix="ltx2_session_init_", + dir=str(temp_root) if temp_root is not None else None, + )) + file_path = temp_dir / "initial_frame.png" + normalized.save(file_path, format="PNG") + + return SessionInitImage( + file_path=file_path, + temp_dir=temp_dir, + display_name=_sanitize_display_name(payload.get("name")), + mime_type=mime_type, + ) + + +def cleanup_session_init_image(session_image: SessionInitImage | None) -> None: + if session_image is None: + return + shutil.rmtree(session_image.temp_dir, ignore_errors=True) diff --git a/apps/dreamverse/dreamverse/session_logger.py b/apps/dreamverse/dreamverse/session_logger.py new file mode 100644 index 000000000..f537fdacc --- /dev/null +++ b/apps/dreamverse/dreamverse/session_logger.py @@ -0,0 +1,46 @@ +# pyright: reportMissingTypeArgument=false +from __future__ import annotations + +import asyncio +from datetime import datetime, timezone +import json +import socket +from pathlib import Path +from typing import Any + + +def _utc_now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +class SessionEventLogger: + + def __init__(self, root_dir: Path): + self.hostname = socket.gethostname() + timestamp = datetime.now(timezone.utc).strftime("%y%m%d_%H%M%S") + self.directory = root_dir / self.hostname + self.path = self.directory / f"{timestamp}.jsonl" + self._lock = asyncio.Lock() + + self.directory.mkdir(parents=True, exist_ok=True) + self.path.touch(exist_ok=False) + + async def write_event( + self, + *, + event: str, + client_id: str, + payload: dict[str, Any] | None = None, + ) -> None: + entry = { + "ts": _utc_now_iso(), + "event": event, + "hostname": self.hostname, + "client_id": client_id, + } + if payload: + entry.update(payload) + + async with self._lock: + with self.path.open("a", encoding="utf-8") as fp: + fp.write(json.dumps(entry, ensure_ascii=False) + "\n") diff --git a/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py b/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py new file mode 100644 index 000000000..52552b16a --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py @@ -0,0 +1,1171 @@ +# pyright: reportMissingTypeArgument=false, reportArgumentType=false, reportOptionalSubscript=false, reportOperatorIssue=false +from __future__ import annotations + +import asyncio +import os +import re +import time + + +os.environ.setdefault("CEREBRAS_API_KEY", "dummy") +os.environ.setdefault("GROQ_API_KEY", "dummy") + +import dreamverse.prompt_enhancer as prompt_enhancer_module +from dreamverse.prompt_enhancer import ( + ProviderRuntime, + PromptEnhancer, + _build_prompt_backup_path, + _load_prompt_required, + _parse_json_response, + _resolve_prompt_save_path, +) + + +class _FakeResponse: + def __init__(self, payload: dict): + self._payload = payload + + def model_dump(self, mode: str = "json"): # noqa: ARG002 + return self._payload + + +class _FakeSyncCompletions: + def __init__(self, payload: dict): + self._payload = payload + + def create(self, **kwargs): # noqa: ARG002 + return _FakeResponse(self._payload) + + +class _FakeSyncClient: + def __init__(self, payload: dict): + self.chat = type( + "_FakeChat", + (), + {"completions": _FakeSyncCompletions(payload)}, + )() + + +class _DelayedSyncCompletions: + def __init__(self, payload: dict, delay_s: float = 0.0, exc: Exception | None = None): + self._payload = payload + self._delay_s = delay_s + self._exc = exc + + def create(self, **kwargs): # noqa: ARG002 + if self._delay_s > 0: + time.sleep(self._delay_s) + if self._exc is not None: + raise self._exc + return _FakeResponse(self._payload) + + +class _DelayedSyncClient: + def __init__(self, payload: dict, delay_s: float = 0.0, exc: Exception | None = None): + self.chat = type( + "_FakeChat", + (), + { + "completions": _DelayedSyncCompletions( + payload, + delay_s=delay_s, + exc=exc, + ) + }, + )() + + +def _chat_payload_with_content(content: str) -> dict: + return { + "choices": [ + { + "message": { + "content": content, + } + } + ] + } + + +def _build_test_enhancer(payload: dict) -> PromptEnhancer: + enhancer = PromptEnhancer.__new__(PromptEnhancer) + enhancer.provider = "cerebras" + enhancer.provider_label = "cerebras" + enhancer.model = "gpt-test" + enhancer.rewrite_default_model = "gpt-test" + enhancer.rewrite_model_options = ["gpt-test"] + enhancer.default_timeout_ms = 1200 + enhancer.http_timeout_ms = 1200 + enhancer.temperature = 0.4 + enhancer.rewrite_default_temperature = 0.4 + enhancer.max_completion_tokens = 512 + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + enhancer.rewrite_user_system_prompt_path = "/tmp/rewrite_user_system_prompt.md" + enhancer.rewrite_user_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt = "system prompt" + enhancer.rewrite_user_system_prompt = "system prompt" + enhancer.provider_request_models = { + "cerebras": "gpt-test", + "groq": "openai/gpt-test", + } + enhancer.provider_success_counts = { + "cerebras": 0, + "groq": 0, + } + enhancer.client = _FakeSyncClient(payload) + return enhancer + + +def _build_staged_enhancer( + *, + cerebras_payload: dict, + groq_payload: dict, + cerebras_request_model: str = "gpt-test", + groq_request_model: str = "openai/gpt-test", + cerebras_delay_s: float = 0.0, + groq_delay_s: float = 0.0, + cerebras_exc: Exception | None = None, + groq_exc: Exception | None = None, +) -> PromptEnhancer: + enhancer = _build_test_enhancer(cerebras_payload) + enhancer.provider_request_models = { + "cerebras": cerebras_request_model, + "groq": groq_request_model, + } + enhancer.enhance_system_prompt = "enhance prompt" + enhancer.auto_system_prompt = "auto prompt" + enhancer.provider_runtimes = [ + ProviderRuntime( + name="cerebras", + api_key="cerebras-key", + api_base_url=None, + request_model=cerebras_request_model, + client=_DelayedSyncClient( + cerebras_payload, + delay_s=cerebras_delay_s, + exc=cerebras_exc, + ), + ), + ProviderRuntime( + name="groq", + api_key="groq-key", + api_base_url="https://api.groq.com/openai/v1", + request_model=groq_request_model, + client=_DelayedSyncClient( + groq_payload, + delay_s=groq_delay_s, + exc=groq_exc, + ), + ), + ] + enhancer.provider_runtime_stages = [ + list(enhancer.provider_runtimes), + ] + enhancer.provider_success_counts = { + "cerebras": 0, + "groq": 0, + } + enhancer.client = enhancer.provider_runtimes[0].client + return enhancer + + +class _FakeOpenAIClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.chat = type( + "_FakeChat", + (), + {"completions": _FakeSyncCompletions(_chat_payload_with_content("{}"))}, + )() + + +class _FakeCerebrasClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.chat = type( + "_FakeChat", + (), + {"completions": _FakeSyncCompletions(_chat_payload_with_content("{}"))}, + )() + + +def test_parse_json_response_accepts_fenced_json_with_prose(): + parsed = _parse_json_response( + "Here is the rewrite:\n```json\n{\"segment_prompts\":[\"A\",\"B\"]}\n```\nThanks." + ) + assert parsed == {"segment_prompts": ["A", "B"]} + + +def test_parse_json_response_extracts_first_embedded_object(): + parsed = _parse_json_response( + "Model output:\n{\"segment_prompts\":[\"A\",\"B\"]}\n(complete)" + ) + assert parsed == {"segment_prompts": ["A", "B"]} + + +def test_load_prompt_required_falls_back_to_default_path(tmp_path): + fallback_path = tmp_path / "next_segment_system_prompt.md" + fallback_path.write_text("fallback prompt\n", encoding="utf-8") + + loaded_prompt = _load_prompt_required( + str(tmp_path / "prompts.local" / "next_segment_system_prompt.md"), + "next-segment", + str(fallback_path), + ) + + assert loaded_prompt == "fallback prompt" + + +def test_load_prompt_required_unwraps_python_style_system_prompt_assignment(tmp_path): + prompt_path = tmp_path / "rewrite_window_system_prompt.md" + prompt_path.write_text( + 'SYSTEM_PROMPT = """\nfirst line\nsecond line\n"""\n', + encoding="utf-8", + ) + + loaded_prompt = _load_prompt_required( + str(prompt_path), + "rewrite-window", + ) + + assert loaded_prompt == "first line\nsecond line" + + +def test_build_client_supports_cerebras_provider(monkeypatch): + monkeypatch.setattr(prompt_enhancer_module, "Cerebras", _FakeCerebrasClient) + + enhancer = PromptEnhancer.__new__(PromptEnhancer) + enhancer.provider = "cerebras" + enhancer.api_key = "cerebras-key" + enhancer.api_base_url = None + + client = PromptEnhancer._build_client(enhancer) + + assert isinstance(client, _FakeCerebrasClient) + assert client.kwargs == { + "api_key": "cerebras-key", + } + + +def test_build_client_supports_groq_provider(monkeypatch): + monkeypatch.setattr(prompt_enhancer_module, "OpenAI", _FakeOpenAIClient) + + enhancer = PromptEnhancer.__new__(PromptEnhancer) + enhancer.provider = "groq" + enhancer.api_key = "groq-key" + enhancer.api_base_url = "https://api.groq.com/openai/v1" + + client = PromptEnhancer._build_client(enhancer) + + assert isinstance(client, _FakeOpenAIClient) + assert client.kwargs == { + "api_key": "groq-key", + "base_url": "https://api.groq.com/openai/v1", + } + + +def test_rewrite_prompt_sequence_accepts_segment_prompts_output(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.rollout_id == "preset_a" + assert result.rollout_label == "Preset A" + assert result.prompts == ["A", "B"] + + +def test_rewrite_prompt_sequence_accepts_legacy_rewritten_prompts_output(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"rewritten_prompts":["A","B"]}') + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.rollout_id == "current_rollout" + assert result.rollout_label == "Current rollout" + assert result.prompts == ["A", "B"] + + +def test_rewrite_prompt_sequence_accepts_segment_dicts_without_top_level_rollout_metadata(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"segments":[{"prompt":"A"},{"text":"B"}]}' + ) + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + preset_id="preset_a", + preset_label="Preset A", + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.rollout_id == "preset_a" + assert result.rollout_label == "Preset A" + assert result.prompts == ["A", "B"] + + +def test_rewrite_prompt_sequence_accepts_numbered_prose_output(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + "The user is asking for a cinematic rewrite.\n\n" + "1. A dog bounds across the moon's dusty surface, kicking up silver regolith as it chases a rabbit beneath the black sky.\n" + "2. The rabbit darts around a crater rim while the dog lunges after it, Earth glowing blue in the distance.\n" + ) + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.rollout_id == "current_rollout" + assert result.rollout_label == "Current rollout" + assert result.prompts == [ + "A dog bounds across the moon's dusty surface, kicking up silver regolith as it chases a rabbit beneath the black sky.", + "The rabbit darts around a crater rim while the dog lunges after it, Earth glowing blue in the distance.", + ] + + +def test_enhance_prompt_prefers_cerebras_before_groq_fallback(): + enhancer = _build_staged_enhancer( + cerebras_payload=_chat_payload_with_content('{"prompt":"Cerebras prompt"}'), + groq_payload=_chat_payload_with_content('{"prompt":"Groq prompt"}'), + cerebras_delay_s=0.01, + groq_delay_s=0.01, + ) + + result = asyncio.run( + enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + ) + ) + + assert result.fallback_used is False + assert result.error is None + assert result.provider == "cerebras" + assert result.model == "gpt-test" + assert result.prompt == "Cerebras prompt" + assert enhancer.get_provider_success_counts() == { + "cerebras": 1, + "groq": 0, + } + + +def test_enhance_prompt_uses_groq_when_cerebras_fails(): + enhancer = _build_staged_enhancer( + cerebras_payload=_chat_payload_with_content("{}"), + groq_payload=_chat_payload_with_content('{"prompt":"Groq prompt"}'), + cerebras_exc=RuntimeError("cerebras overloaded"), + groq_delay_s=0.01, + ) + + result = asyncio.run( + enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + ) + ) + + assert result.fallback_used is False + assert result.error is None + assert result.provider == "groq" + assert result.prompt == "Groq prompt" + assert enhancer.get_provider_success_counts() == { + "cerebras": 0, + "groq": 1, + } + + +def test_enhance_prompt_can_use_groq_when_cerebras_times_out(): + enhancer = _build_staged_enhancer( + cerebras_payload=_chat_payload_with_content('{"prompt":"Slow Cerebras prompt"}'), + groq_payload=_chat_payload_with_content('{"prompt":"Groq prompt"}'), + cerebras_delay_s=0.4, + groq_delay_s=0.01, + ) + enhancer.http_timeout_ms = 50 + enhancer.default_timeout_ms = 50 + + result = asyncio.run( + enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + timeout_ms=50, + ) + ) + + assert result.fallback_used is False + assert result.error is None + assert result.provider == "groq" + assert result.prompt == "Groq prompt" + assert enhancer.get_provider_success_counts() == { + "cerebras": 0, + "groq": 1, + } + + +def test_enhance_prompt_can_use_cerebras_when_it_returns_first(): + enhancer = _build_staged_enhancer( + cerebras_payload=_chat_payload_with_content('{"prompt":"Primary Cerebras prompt"}'), + groq_payload=_chat_payload_with_content('{"prompt":"Groq prompt"}'), + cerebras_delay_s=0.01, + groq_delay_s=0.08, + ) + + result = asyncio.run( + enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + ) + ) + + assert result.fallback_used is False + assert result.error is None + assert result.provider == "cerebras" + assert result.model == "gpt-test" + assert result.prompt == "Primary Cerebras prompt" + assert enhancer.get_provider_success_counts() == { + "cerebras": 1, + "groq": 0, + } + + +def test_rewrite_prompt_sequence_keeps_raw_output_on_parse_error(): + enhancer = _build_test_enhancer( + _chat_payload_with_content("I cannot comply with JSON right now.") + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is True + assert "No JSON object found in assistant response." in (result.error or "") + assert result.raw_response_text == "I cannot comply with JSON right now." + assert result.rollout_id == "current_rollout" + assert result.rollout_label == "Current rollout" + assert result.prompts == ["prompt one", "prompt two"] + + +def test_rewrite_prompt_sequence_uses_current_rollout_payload_shape(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}' + ) + ) + captured = { + "body": None, + "timeout_seconds": None, + } + + async def _fake_request_content_with_body(*, body: dict, timeout_seconds: float): + captured["body"] = body + captured["timeout_seconds"] = timeout_seconds + return ( + _chat_payload_with_content( + '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}' + ), + '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}', + ) + + enhancer._request_content_with_body = _fake_request_content_with_body # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + preset_id="preset_a", + preset_label="Preset A", + rewrite_instruction="make it cinematic", + rewrite_model="gpt-test", + rewrite_temperature=0.2, + timeout_ms=800, + ) + ) + + assert result.fallback_used is False + assert captured["body"]["messages"][0] == { + "role": "system", + "content": "system prompt", + } + assert captured["body"]["messages"][1]["role"] == "user" + assert prompt_enhancer_module.json.loads(captured["body"]["messages"][1]["content"]) == { + "mode": "edit_existing_rollout", + "request": ( + "Rewrite all segment prompts with improved continuity and cinematic detail. " + "Keep count and ordering identical." + ), + "user_instruction": "make it cinematic", + "current_rollout": { + "id": "preset_a", + "label": "Preset A", + "segment_prompts": ["prompt one", "prompt two"], + }, + } + + +def test_rewrite_prompt_sequence_supports_new_rollout_mode(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' + '"A","B","C","D","E","F"]}' + ) + ) + captured = { + "body": None, + } + + async def _fake_request_content_with_body(*, body: dict, timeout_seconds: float): + del timeout_seconds + captured["body"] = body + return ( + _chat_payload_with_content( + '{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' + '"A","B","C","D","E","F"]}' + ), + '{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' + '"A","B","C","D","E","F"]}', + ) + + enhancer._request_content_with_body = _fake_request_content_with_body # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + [], + preset_id="custom_editable", + preset_label="Custom rollout", + rewrite_instruction="A moonbase corridor thriller with flooding and red alarms", + rewrite_model="gpt-test", + rewrite_temperature=0.2, + timeout_ms=800, + ) + ) + + assert result.fallback_used is False + assert result.prompts == ["A", "B", "C", "D", "E", "F"] + assert prompt_enhancer_module.json.loads(captured["body"]["messages"][1]["content"]) == { + "mode": "new_rollout", + "request": ( + "Rewrite all segment prompts with improved continuity and cinematic detail. " + "Keep count and ordering identical." + ), + "user_instruction": "A moonbase corridor thriller with flooding and red alarms", + "desired_segment_count": 6, + "rollout_id_hint": "custom_editable", + "rollout_label_hint": "Custom rollout", + } + + +def test_rewrite_prompt_sequence_uses_session_override_system_prompt(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.rewrite_all_system_prompt = "shared system prompt" + captured = { + "body": None, + } + + async def _fake_request_content_with_body(*, body: dict, timeout_seconds: float): + del timeout_seconds + captured["body"] = body + return ( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ), + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}', + ) + + enhancer._request_content_with_body = _fake_request_content_with_body # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + preset_id="preset_a", + preset_label="Preset A", + rewrite_instruction="make it cinematic", + rewrite_model="gpt-test", + system_prompt_override="session specific system prompt", + ) + ) + + assert result.fallback_used is False + assert captured["body"]["messages"][0] == { + "role": "system", + "content": "session specific system prompt", + } + + +def test_resolve_rewrite_new_rollout_system_prompt_uses_dedicated_prompt(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.rewrite_all_system_prompt = "shared rewrite system prompt" + enhancer.rewrite_user_system_prompt = "new rollout rewrite system prompt" + + resolved = enhancer.resolve_rewrite_new_rollout_system_prompt(None) + + assert resolved == "new rollout rewrite system prompt" + + +def test_resolve_rewrite_new_rollout_system_prompt_prefers_override(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.rewrite_all_system_prompt = "shared rewrite system prompt" + enhancer.rewrite_user_system_prompt = "new rollout rewrite system prompt" + + resolved = enhancer.resolve_rewrite_new_rollout_system_prompt( + "session specific system prompt" + ) + + assert resolved == "session specific system prompt" + + +def test_generate_auto_prompt_uses_selected_model(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"next_prompt":"Auto next"}') + ) + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] + enhancer.rewrite_default_model = "gpt-test" + captured_model = {"value": None} + + async def _fake_request_content( + *, + system_prompt: str, + user_payload: dict, + timeout_ms: int, + model: str | None = None, + temperature: float | None = None, + ): + captured_model["value"] = model + return ( + _chat_payload_with_content('{"next_prompt":"Auto next"}'), + '{"next_prompt":"Auto next"}', + ) + + enhancer._request_content = _fake_request_content # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.generate_auto_prompt( + locked_segments=["One"], + next_segment_idx=2, + model="gpt-alt", + timeout_ms=800, + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.prompt == "Auto next" + assert captured_model["value"] == "gpt-alt" + assert result.model == "gpt-alt" + + +def test_enhance_prompt_uses_selected_model(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"next_prompt":"Enhanced next"}') + ) + enhancer.enhance_system_prompt = "enhance system prompt" + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] + enhancer.rewrite_default_model = "gpt-test" + captured_model = {"value": None} + + async def _fake_request_content( + *, + system_prompt: str, + user_payload: dict, + timeout_ms: int, + model: str | None = None, + temperature: float | None = None, + ): + captured_model["value"] = model + return ( + _chat_payload_with_content('{"next_prompt":"Enhanced next"}'), + '{"next_prompt":"Enhanced next"}', + ) + + enhancer._request_content = _fake_request_content # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.enhance_prompt( + "A user live prompt", + locked_segments=["Segment 1"], + next_segment_idx=2, + model="gpt-alt", + timeout_ms=800, + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.prompt == "Enhanced next" + assert captured_model["value"] == "gpt-alt" + assert result.model == "gpt-alt" + + +def test_enhance_prompt_single_clip_uses_auto_extension_prompt_and_prompt_field(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"prompt":"Extended single clip"}') + ) + enhancer.enhance_system_prompt = "enhance system prompt" + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] + enhancer.rewrite_default_model = "gpt-test" + captured = { + "model": None, + "system_prompt": None, + "user_payload": None, + } + + async def _fake_request_content( + *, + system_prompt: str, + user_payload: dict, + timeout_ms: int, + model: str | None = None, + temperature: float | None = None, + ): + del timeout_ms, temperature + captured["model"] = model + captured["system_prompt"] = system_prompt + captured["user_payload"] = user_payload + return ( + _chat_payload_with_content('{"prompt":"Extended single clip"}'), + '{"prompt":"Extended single clip"}', + ) + + enhancer._request_content = _fake_request_content # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.enhance_prompt( + "short 5s idea", + mode="single_clip", + model="gpt-alt", + timeout_ms=800, + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.prompt == "Extended single clip" + assert result.model == "gpt-alt" + assert captured["model"] == "gpt-alt" + assert captured["system_prompt"] == "auto system prompt" + assert captured["user_payload"] == { + "request": ( + "Expand the user prompt into one detailed prompt for a " + "single 5-second LTX-2.3 video clip. Respond with " + 'valid JSON only as {"prompt": "..."}.' # noqa: E501 + ), + "user_prompt": "short 5s idea", + } + + +def test_enhance_prompt_single_clip_rejects_plain_text_response(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + "Medium shot of a woman by a rainy cafe window as she lifts her " + "phone, exhales softly, and the camera makes a slow push in." + ) + ) + enhancer.auto_system_prompt = "auto system prompt" + + result = asyncio.run( + enhancer.enhance_prompt( + "woman in a cafe looking at her phone", + mode="single_clip", + model="gpt-test", + timeout_ms=800, + ) + ) + assert result.fallback_used is True + assert "No JSON object found in assistant response." in result.error + assert result.prompt == "" + + +def test_enhance_prompt_single_clip_rejects_segment_prompts_json(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"segment_prompts":["A","B"]}') + ) + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test"] + enhancer.rewrite_default_model = "gpt-test" + + result = asyncio.run( + enhancer.enhance_prompt( + "short 5s idea", + mode="single_clip", + model="gpt-test", + timeout_ms=800, + ) + ) + assert result.fallback_used is True + assert result.prompt == "" + assert "Missing prompt string." in (result.error or "") + + +def test_enhance_prompt_requires_json_and_does_not_fallback_to_raw_text(): + enhancer = _build_test_enhancer( + _chat_payload_with_content("A cinematic continuation with slow dolly movement.") + ) + enhancer.enhance_system_prompt = "enhance system prompt" + enhancer.rewrite_model_options = ["gpt-test"] + enhancer.rewrite_default_model = "gpt-test" + + result = asyncio.run( + enhancer.enhance_prompt( + "raw input prompt", + locked_segments=["Segment 1"], + next_segment_idx=2, + model="gpt-test", + timeout_ms=800, + ) + ) + assert result.fallback_used is True + assert result.prompt == "" + assert "No JSON object found in assistant response." in (result.error or "") + + +def test_generate_auto_prompt_requires_json_and_does_not_fallback_to_raw_text(): + enhancer = _build_test_enhancer( + _chat_payload_with_content("A calm, grounded continuation with subtle motion.") + ) + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test"] + enhancer.rewrite_default_model = "gpt-test" + + result = asyncio.run( + enhancer.generate_auto_prompt( + locked_segments=["Segment 1"], + next_segment_idx=2, + model="gpt-test", + timeout_ms=800, + ) + ) + assert result.fallback_used is True + assert result.prompt == "" + assert "No JSON object found in assistant response." in (result.error or "") + + +def test_rewrite_prompt_sequence_includes_raw_json_when_content_empty(): + enhancer = _build_test_enhancer( + { + "choices": [ + { + "finish_reason": "length", + "message": { + "content": [], + "refusal": None, + }, + } + ], + "usage": {"completion_tokens": 0}, + } + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is True + assert "No rewrite segment prompts found in assistant response." in (result.error or "") + assert isinstance(result.raw_response_text, str) + assert '"finish_reason": "length"' in result.raw_response_text + + +def test_get_rewrite_model_config_returns_fixed_defaults(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.rewrite_default_model = "gpt-oss-120b" + enhancer.rewrite_model_options = ["gpt-oss-120b"] + + config = enhancer.get_rewrite_model_config() + + assert config["default_model"] == "gpt-oss-120b" + assert config["options"] == ["gpt-oss-120b"] + + +def test_get_prompt_config_includes_auto_extension_prompt(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.enhance_system_prompt_path = "/tmp/next.md" + enhancer.auto_system_prompt_path = "/tmp/auto.md" + enhancer.rewrite_all_system_prompt_path = "/tmp/rewrite.md" + enhancer.rewrite_user_system_prompt_path = "/tmp/rewrite_user.md" + enhancer.enhance_system_prompt = "next prompt" + enhancer.auto_system_prompt = "auto prompt" + enhancer.rewrite_all_system_prompt = "rewrite prompt" + enhancer.rewrite_user_system_prompt = "rewrite user prompt" + + config = enhancer.get_prompt_config() + + assert config["next_segment_system_prompt_path"] == "/tmp/next.md" + assert config["auto_extension_system_prompt_path"] == "/tmp/auto.md" + assert config["rewrite_window_system_prompt_path"] == "/tmp/rewrite.md" + assert config["rewrite_user_system_prompt_path"] == "/tmp/rewrite_user.md" + assert config["next_segment_system_prompt"] == "next prompt" + assert config["auto_extension_system_prompt"] == "auto prompt" + assert config["rewrite_window_system_prompt"] == "rewrite prompt" + assert config["rewrite_user_system_prompt"] == "rewrite user prompt" + assert config["rewrite_model"] == "gpt-test" + assert config["rewrite_model_options"] == ["gpt-test"] + assert config["rewrite_temperature"] == 0.4 + + +def test_get_prompt_config_reports_loaded_fallback_prompt_path(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + rewrite_fallback_path = tmp_path / "rewrite_window_system_prompt.md" + rewrite_fallback_path.write_text("rewrite prompt\n", encoding="utf-8") + next_path = tmp_path / "next.md" + next_path.write_text("next prompt\n", encoding="utf-8") + auto_path = tmp_path / "auto.md" + auto_path.write_text("auto prompt\n", encoding="utf-8") + enhancer.rewrite_all_system_prompt_path = str( + tmp_path / "prompts.local" / "rewrite_window_system_prompt.md" + ) + enhancer.rewrite_all_system_prompt_fallback_path = str(rewrite_fallback_path) + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.reload_system_prompts() + + config = enhancer.get_prompt_config() + + assert config["rewrite_window_system_prompt_path"] == str(rewrite_fallback_path) + + +def test_reload_system_prompts_falls_back_to_rewrite_window_when_user_prompt_empty( + tmp_path, +): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite_window_system_prompt.md" + rewrite_user_path = tmp_path / "rewrite_user_system_prompt.md" + next_path.write_text("next prompt\n", encoding="utf-8") + auto_path.write_text("auto prompt\n", encoding="utf-8") + rewrite_path.write_text("rewrite window prompt\n", encoding="utf-8") + rewrite_user_path.write_text("\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.rewrite_user_system_prompt_path = str(rewrite_user_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + enhancer.rewrite_user_system_prompt_fallback_path = None + + enhancer.reload_system_prompts() + + assert enhancer.rewrite_all_system_prompt == "rewrite window prompt" + assert enhancer.rewrite_user_system_prompt == "rewrite window prompt" + assert enhancer.rewrite_user_system_prompt_source_path == str(rewrite_path) + + +def test_save_prompt_config_updates_auto_extension_prompt(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite.md" + next_path.write_text("next original\n", encoding="utf-8") + auto_path.write_text("auto original\n", encoding="utf-8") + rewrite_path.write_text("rewrite original\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + + config = enhancer.save_prompt_config( + auto_extension_system_prompt="auto updated", + ) + + assert auto_path.read_text(encoding="utf-8").strip() == "auto updated" + assert config["auto_extension_system_prompt"] == "auto updated" + + +def test_save_prompt_config_updates_rewrite_user_prompt(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite.md" + rewrite_user_path = tmp_path / "rewrite_user.md" + next_path.write_text("next original\n", encoding="utf-8") + auto_path.write_text("auto original\n", encoding="utf-8") + rewrite_path.write_text("rewrite original\n", encoding="utf-8") + rewrite_user_path.write_text("rewrite user original\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.rewrite_user_system_prompt_path = str(rewrite_user_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + enhancer.rewrite_user_system_prompt_fallback_path = None + + config = enhancer.save_prompt_config( + rewrite_user_system_prompt="rewrite user updated", + ) + + assert rewrite_user_path.read_text(encoding="utf-8").strip() == "rewrite user updated" + assert config["rewrite_user_system_prompt"] == "rewrite user updated" + + +def test_save_prompt_config_updates_rewrite_model(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite.md" + next_path.write_text("next prompt\n", encoding="utf-8") + auto_path.write_text("auto prompt\n", encoding="utf-8") + rewrite_path.write_text("rewrite prompt\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + enhancer.rewrite_default_model = "gpt-test" + enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] + + config = enhancer.save_prompt_config( + rewrite_model="gpt-alt", + ) + + assert enhancer.rewrite_default_model == "gpt-alt" + assert config["rewrite_model"] == "gpt-alt" + assert config["rewrite_model_options"] == ["gpt-test", "gpt-alt"] + + +def test_save_prompt_config_updates_rewrite_temperature(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite.md" + next_path.write_text("next prompt\n", encoding="utf-8") + auto_path.write_text("auto prompt\n", encoding="utf-8") + rewrite_path.write_text("rewrite prompt\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + + config = enhancer.save_prompt_config( + rewrite_temperature=1.3, + ) + + assert enhancer.rewrite_default_temperature == 1.3 + assert config["rewrite_temperature"] == 1.3 + + +def test_save_prompt_config_creates_versioned_backup_for_existing_prompt(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite_window_system_prompt.md" + next_path.write_text("next original\n", encoding="utf-8") + auto_path.write_text("auto original\n", encoding="utf-8") + rewrite_path.write_text("rewrite original\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + + enhancer.save_prompt_config( + rewrite_window_system_prompt="rewrite updated", + ) + + backup_paths = sorted( + tmp_path.glob("rewrite_window_system_prompt.*.bak.md") + ) + + assert rewrite_path.read_text(encoding="utf-8").strip() == "rewrite updated" + assert len(backup_paths) == 1 + assert backup_paths[0].read_text(encoding="utf-8").strip() == "rewrite original" + + +def test_resolve_prompt_save_path_prefers_existing_fallback_file(tmp_path): + primary_path = tmp_path / "prompts.local" / "rewrite.md" + fallback_path = tmp_path / "rewrite.md" + fallback_path.write_text("existing\n", encoding="utf-8") + + resolved = _resolve_prompt_save_path(str(primary_path), str(fallback_path)) + + assert resolved == str(fallback_path) + + +def test_build_prompt_backup_path_preserves_extension(tmp_path): + prompt_path = tmp_path / "rewrite_window_system_prompt.md" + + backup_path = _build_prompt_backup_path(prompt_path) + + assert backup_path.parent == tmp_path + assert re.match( + r"rewrite_window_system_prompt\.\d{8}_\d{6}\.bak\.md$", + backup_path.name, + ) diff --git a/apps/dreamverse/dreamverse/tests/test_realtime_stress.py b/apps/dreamverse/dreamverse/tests/test_realtime_stress.py new file mode 100644 index 000000000..2c492b351 --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_realtime_stress.py @@ -0,0 +1,951 @@ +"""Multiprocess realtime stress test for LTX2 streaming websocket service.""" +# pyright: reportArgumentType=false, reportOptionalMemberAccess=false + +from __future__ import annotations + +import argparse +import asyncio +from collections import defaultdict +from datetime import datetime, timezone +import json +import math +import multiprocessing as mp +import os +from pathlib import Path +import sys +import time +import traceback +from typing import Any +import uuid + +import pytest + +pytestmark = pytest.mark.gpu + +try: + import websockets +except ModuleNotFoundError: + websockets = None # type: ignore[assignment] + + +DEFAULT_PRESET_FILE = ( + Path(__file__).resolve().parents[2] + / "web" + / "prompts" + / "selected_ltx2_continuation_story_presets.json" +) + + +def utc_now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +def iso_from_epoch(epoch_s: float) -> str: + return datetime.fromtimestamp(epoch_s, tz=timezone.utc).isoformat() + + +def iso_to_epoch(iso_ts: str | None) -> float | None: + if not iso_ts: + return None + try: + return datetime.fromisoformat(iso_ts).timestamp() + except ValueError: + return None + + +def safe_percentile(values: list[float], percentile: float) -> float | None: + if not values: + return None + sorted_values = sorted(values) + if len(sorted_values) == 1: + return sorted_values[0] + rank = (len(sorted_values) - 1) * (percentile / 100.0) + lower = math.floor(rank) + upper = math.ceil(rank) + if lower == upper: + return sorted_values[lower] + fraction = rank - lower + return ( + sorted_values[lower] + + (sorted_values[upper] - sorted_values[lower]) * fraction + ) + + +def summarize_series(values: list[float]) -> dict[str, float | int | None]: + if not values: + return { + "count": 0, + "min": None, + "p50": None, + "p95": None, + "p99": None, + "max": None, + "avg": None, + } + return { + "count": len(values), + "min": min(values), + "p50": safe_percentile(values, 50), + "p95": safe_percentile(values, 95), + "p99": safe_percentile(values, 99), + "max": max(values), + "avg": sum(values) / len(values), + } + + +def parse_int(value: object) -> int | None: + try: + parsed = int(value) + except (TypeError, ValueError): + return None + return parsed + + +def format_num(value: float | int | None, digits: int = 2) -> str: + if value is None: + return "n/a" + if isinstance(value, int): + return str(value) + return f"{value:.{digits}f}" + + +def load_curated_prompts( + preset_file: Path, + preset_id: str | None, + curated_limit: int, +) -> tuple[str, list[str], int]: + if not preset_file.is_file(): + raise ValueError(f"Preset file not found: {preset_file}") + try: + payload = json.loads(preset_file.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid JSON in preset file: {preset_file}") from exc + + if not isinstance(payload, list) or len(payload) == 0: + raise ValueError("Preset file must contain a non-empty JSON array.") + + selected: dict[str, Any] | None = None + if preset_id: + for item in payload: + if not isinstance(item, dict): + continue + if str(item.get("id", "")).strip() == preset_id: + selected = item + break + if selected is None: + raise ValueError(f"Preset id not found: {preset_id}") + else: + for item in payload: + if isinstance(item, dict): + selected = item + break + if selected is None: + raise ValueError("No valid preset object found in preset file.") + + selected_id = str(selected.get("id", "")).strip() or "unknown_preset" + raw_prompts = selected.get("segment_prompts", []) + if not isinstance(raw_prompts, list): + raise ValueError( + f"Preset {selected_id} has invalid segment_prompts (must be list)." + ) + + prompts = [ + str(prompt).strip() + for prompt in raw_prompts + if isinstance(prompt, str) and str(prompt).strip() + ] + if not prompts: + raise ValueError(f"Preset {selected_id} has no non-empty prompts.") + + limited = prompts[:curated_limit] + if not limited: + raise ValueError( + f"curated_limit={curated_limit} produced no prompts for preset " + f"{selected_id}." + ) + return selected_id, limited, len(prompts) + + +async def run_single_session( + *, + worker_id: int, + worker_session_idx: int, + config: dict[str, Any], +) -> dict[str, Any]: + session_id = f"w{worker_id}_u{worker_session_idx}_{uuid.uuid4().hex[:8]}" + process_id = os.getpid() + url = str(config["url"]) + session_timeout_s = float(config["session_timeout_s"]) + post_complete_wait_s = float(config["post_complete_wait_s"]) + connect_timeout_s = float(config["connect_timeout_s"]) + curated_prompts = list(config["curated_prompts"]) + preset_id = str(config["preset_id"]) + + session_start_monotonic = time.monotonic() + session_start_epoch = time.time() + session_data: dict[str, Any] = { + "session_id": session_id, + "process_id": process_id, + "worker_id": worker_id, + "status": "failed", + "error": None, + "preset_id": preset_id, + "curated_prompt_count": len(curated_prompts), + "connect_start_ts_utc": iso_from_epoch(session_start_epoch), + "connect_finish_ts_utc": None, + "session_init_sent_ts_utc": None, + "gpu_assigned_ts_utc": None, + "target_segment_complete_ts_utc": None, + "leave_sent_ts_utc": None, + "close_ts_utc": None, + "duration_ms": None, + "queue_wait_ms": None, + "initial_total_segments": None, + "segments_started": 0, + "segments_completed": 0, + "media_segments_completed": 0, + "total_chunks": 0, + "total_chunk_bytes": 0, + "first_chunk_finish_ts_utc": None, + "last_chunk_finish_ts_utc": None, + "first_media_segment_complete_ts_utc": None, + "first_chunk_before_first_media_complete": None, + "session_goodput_mbps": None, + "chunks": [], + } + + connect_finish_monotonic: float | None = None + current_segment_idx: int | None = None + initial_total_segments: int | None = None + first_chunk_finish_epoch: float | None = None + first_media_segment_complete_epoch: float | None = None + last_chunk_finish_epoch: float | None = None + last_chunk_finish_monotonic: float | None = None + + try: + async with websockets.connect( + url, + max_size=None, + ping_interval=None, + open_timeout=connect_timeout_s, + close_timeout=2.0, + ) as ws: + connect_finish_monotonic = time.monotonic() + session_data["connect_finish_ts_utc"] = utc_now_iso() + + init_payload = { + "type": "session_init_v2", + "preset_id": preset_id, + "curated_prompts": curated_prompts, + "enhancement_enabled": False, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + } + await ws.send(json.dumps(init_payload)) + session_data["session_init_sent_ts_utc"] = utc_now_iso() + + while True: + elapsed_s = time.monotonic() - session_start_monotonic + timeout_remaining = session_timeout_s - elapsed_s + if timeout_remaining <= 0: + session_data["status"] = "timeout" + session_data["error"] = ( + f"Session timed out after {session_timeout_s:.1f}s." + ) + break + + recv_start_epoch = time.time() + recv_start_monotonic = time.monotonic() + recv_start_iso = iso_from_epoch(recv_start_epoch) + + try: + message = await asyncio.wait_for( + ws.recv(), + timeout=timeout_remaining, + ) + except asyncio.TimeoutError: + session_data["status"] = "timeout" + session_data["error"] = ( + "Timed out waiting for websocket message." + ) + break + except Exception as exc: + session_data["status"] = "failed" + session_data["error"] = f"WebSocket receive failed: {exc}" + break + + recv_finish_epoch = time.time() + recv_finish_monotonic = time.monotonic() + recv_finish_iso = iso_from_epoch(recv_finish_epoch) + + if isinstance(message, bytes): + session_data["total_chunks"] += 1 + session_data["total_chunk_bytes"] += len(message) + + if first_chunk_finish_epoch is None: + first_chunk_finish_epoch = recv_finish_epoch + session_data["first_chunk_finish_ts_utc"] = recv_finish_iso + + chunk_gap_ms: float | None = None + if last_chunk_finish_monotonic is not None: + chunk_gap_ms = ( + recv_finish_monotonic - last_chunk_finish_monotonic + ) * 1000.0 + + session_data["chunks"].append( + { + "segment_idx": current_segment_idx, + "chunk_idx": session_data["total_chunks"], + "size_bytes": len(message), + "chunk_start_ts_utc": recv_start_iso, + "chunk_finish_ts_utc": recv_finish_iso, + "chunk_gap_ms": chunk_gap_ms, + } + ) + last_chunk_finish_monotonic = recv_finish_monotonic + last_chunk_finish_epoch = recv_finish_epoch + session_data["last_chunk_finish_ts_utc"] = recv_finish_iso + continue + + if not isinstance(message, str): + continue + + try: + data = json.loads(message) + except json.JSONDecodeError as exc: + session_data["status"] = "protocol_error" + session_data["error"] = f"Invalid JSON message: {exc}" + break + + msg_type = data.get("type") + if msg_type == "gpu_assigned": + session_data["gpu_assigned_ts_utc"] = recv_finish_iso + if connect_finish_monotonic is not None: + session_data["queue_wait_ms"] = ( + recv_finish_monotonic - connect_finish_monotonic + ) * 1000.0 + elif msg_type == "ltx2_stream_start": + if initial_total_segments is None: + parsed_total = parse_int(data.get("total_segments")) + if parsed_total is not None and parsed_total > 0: + initial_total_segments = parsed_total + session_data["initial_total_segments"] = parsed_total + elif msg_type == "ltx2_segment_start": + parsed_idx = parse_int(data.get("segment_idx")) + current_segment_idx = parsed_idx + session_data["segments_started"] += 1 + elif msg_type == "media_segment_complete": + session_data["media_segments_completed"] += 1 + if first_media_segment_complete_epoch is None: + first_media_segment_complete_epoch = recv_finish_epoch + session_data[ + "first_media_segment_complete_ts_utc" + ] = recv_finish_iso + elif msg_type == "ltx2_segment_complete": + session_data["segments_completed"] += 1 + seg_idx = parse_int(data.get("segment_idx")) + if ( + initial_total_segments is not None + and seg_idx is not None + and seg_idx >= initial_total_segments + ): + session_data[ + "target_segment_complete_ts_utc" + ] = recv_finish_iso + await asyncio.sleep(post_complete_wait_s) + session_data["leave_sent_ts_utc"] = utc_now_iso() + try: + await ws.send(json.dumps({"type": "leave"})) + except Exception: + pass + session_data["status"] = "success" + break + elif msg_type == "session_timeout": + session_data["status"] = "timeout" + session_data["error"] = str( + data.get("message") or "Backend session timeout" + ) + break + elif msg_type == "error": + session_data["status"] = "failed" + session_data["error"] = str( + data.get("message") or "Backend error message" + ) + break + + if session_data["status"] == "failed" and session_data["error"] is None: + session_data["error"] = "Session ended without success." + except Exception as exc: + session_data["status"] = "failed" + session_data["error"] = f"WebSocket connect/run failed: {exc}" + + if ( + first_chunk_finish_epoch is not None + and last_chunk_finish_epoch is not None + and session_data["total_chunk_bytes"] > 0 + ): + duration_s = last_chunk_finish_epoch - first_chunk_finish_epoch + if duration_s > 0: + session_data["session_goodput_mbps"] = ( + session_data["total_chunk_bytes"] * 8.0 / duration_s / 1_000_000.0 + ) + + if ( + first_chunk_finish_epoch is not None + and first_media_segment_complete_epoch is not None + ): + session_data["first_chunk_before_first_media_complete"] = ( + first_chunk_finish_epoch < first_media_segment_complete_epoch + ) + + session_data["close_ts_utc"] = utc_now_iso() + session_data["duration_ms"] = ( + time.monotonic() - session_start_monotonic + ) * 1000.0 + return session_data + + +async def run_worker_sessions( + *, + worker_id: int, + session_count: int, + config: dict[str, Any], +) -> list[dict[str, Any]]: + tasks = [ + asyncio.create_task( + run_single_session( + worker_id=worker_id, + worker_session_idx=idx, + config=config, + ) + ) + for idx in range(session_count) + ] + if not tasks: + return [] + return await asyncio.gather(*tasks) + + +def worker_entry( + worker_id: int, + session_count: int, + config: dict[str, Any], + start_event: Any, + ready_queue: Any, + result_queue: Any, +) -> None: + try: + ready_queue.put({"worker_id": worker_id, "status": "ready"}) + start_event.wait() + sessions = asyncio.run( + run_worker_sessions( + worker_id=worker_id, + session_count=session_count, + config=config, + ) + ) + result_queue.put( + { + "worker_id": worker_id, + "status": "ok", + "sessions": sessions, + } + ) + except Exception as exc: + result_queue.put( + { + "worker_id": worker_id, + "status": "error", + "error": str(exc), + "traceback": traceback.format_exc(), + } + ) + + +def build_summary( + *, + sessions: list[dict[str, Any]], + chunk_gap_threshold_ms: float, +) -> dict[str, Any]: + status_counts: dict[str, int] = defaultdict(int) + chunk_gaps: list[float] = [] + queue_waits: list[float] = [] + session_goodputs: list[float] = [] + progressive_eligible = 0 + progressive_success = 0 + + all_chunk_finish_epochs: list[float] = [] + bucket_bytes: dict[int, int] = defaultdict(int) + total_chunk_bytes = 0 + + for session in sessions: + status = str(session.get("status") or "unknown") + status_counts[status] += 1 + + queue_wait_ms = session.get("queue_wait_ms") + if isinstance(queue_wait_ms, (int, float)): + queue_waits.append(float(queue_wait_ms)) + + session_goodput = session.get("session_goodput_mbps") + if isinstance(session_goodput, (int, float)): + session_goodputs.append(float(session_goodput)) + + progressive_value = session.get("first_chunk_before_first_media_complete") + if isinstance(progressive_value, bool): + progressive_eligible += 1 + if progressive_value: + progressive_success += 1 + + for chunk in session.get("chunks", []): + gap = chunk.get("chunk_gap_ms") + if isinstance(gap, (int, float)): + chunk_gaps.append(float(gap)) + + size_bytes = int(chunk.get("size_bytes") or 0) + finish_epoch = iso_to_epoch(chunk.get("chunk_finish_ts_utc")) + if finish_epoch is None or size_bytes <= 0: + continue + total_chunk_bytes += size_bytes + all_chunk_finish_epochs.append(finish_epoch) + bucket_bytes[int(finish_epoch)] += size_bytes + + chunk_gap_stats = summarize_series(chunk_gaps) + queue_wait_stats = summarize_series(queue_waits) + session_goodput_stats = summarize_series(session_goodputs) + + global_goodput_mbps: float | None = None + if len(all_chunk_finish_epochs) >= 2 and total_chunk_bytes > 0: + duration_s = max(all_chunk_finish_epochs) - min(all_chunk_finish_epochs) + if duration_s > 0: + global_goodput_mbps = ( + total_chunk_bytes * 8.0 / duration_s / 1_000_000.0 + ) + + bucket_throughputs_mbps = [ + (bytes_count * 8.0) / 1_000_000.0 + for _, bytes_count in sorted(bucket_bytes.items()) + ] + bucket_stats = summarize_series(bucket_throughputs_mbps) + + chunk_gap_threshold_breaches = [ + value for value in chunk_gaps if value >= chunk_gap_threshold_ms + ] + non_success = len(sessions) - status_counts.get("success", 0) + + fail_reasons: list[str] = [] + if non_success > 0: + fail_reasons.append( + f"{non_success} session(s) did not complete successfully." + ) + if not chunk_gaps: + fail_reasons.append("No chunk gap data collected.") + if chunk_gap_threshold_breaches: + fail_reasons.append( + f"{len(chunk_gap_threshold_breaches)} chunk gap(s) were >= " + f"{chunk_gap_threshold_ms:.0f}ms." + ) + + passed = len(fail_reasons) == 0 + progressive_ratio = None + if progressive_eligible > 0: + progressive_ratio = progressive_success / progressive_eligible + + return { + "passed": passed, + "fail_reasons": fail_reasons, + "sessions": { + "total": len(sessions), + "success": status_counts.get("success", 0), + "failed": status_counts.get("failed", 0), + "timeout": status_counts.get("timeout", 0), + "protocol_error": status_counts.get("protocol_error", 0), + "other": ( + len(sessions) + - ( + status_counts.get("success", 0) + + status_counts.get("failed", 0) + + status_counts.get("timeout", 0) + + status_counts.get("protocol_error", 0) + ) + ), + }, + "chunk_gap_ms": { + **chunk_gap_stats, + "threshold_ms": chunk_gap_threshold_ms, + "breach_count": len(chunk_gap_threshold_breaches), + }, + "queue_wait_ms": queue_wait_stats, + "progressive_streaming": { + "eligible_sessions": progressive_eligible, + "success_sessions": progressive_success, + "ratio": progressive_ratio, + }, + "bandwidth_mbps": { + "per_session": session_goodput_stats, + "global_goodput_mbps": global_goodput_mbps, + "bucketed_1s": { + "count": bucket_stats.get("count"), + "avg_mbps": bucket_stats.get("avg"), + "peak_mbps": bucket_stats.get("max"), + }, + }, + } + + +def print_summary( + *, + run_info: dict[str, Any], + summary: dict[str, Any], +) -> None: + sessions = summary["sessions"] + chunk_gap = summary["chunk_gap_ms"] + queue_wait = summary["queue_wait_ms"] + progressive = summary["progressive_streaming"] + bandwidth = summary["bandwidth_mbps"] + per_session_bw = bandwidth["per_session"] + bucket_bw = bandwidth["bucketed_1s"] + + print("=== LTX2 Realtime Stress Test Summary ===") + print( + "Run: " + f"url={run_info['url']} clients={run_info['clients']} " + f"processes={run_info['processes']} " + f"preset={run_info['preset_id']} " + f"curated_limit={run_info['curated_limit']}" + ) + print( + "Sessions: " + f"total={sessions['total']} success={sessions['success']} " + f"failed={sessions['failed']} timeout={sessions['timeout']} " + f"protocol_error={sessions['protocol_error']}" + ) + print( + "Chunk gap ms: " + f"min={format_num(chunk_gap['min'])} " + f"p50={format_num(chunk_gap['p50'])} " + f"p95={format_num(chunk_gap['p95'])} " + f"p99={format_num(chunk_gap['p99'])} " + f"max={format_num(chunk_gap['max'])} " + f"threshold={format_num(chunk_gap['threshold_ms'])} " + f"breaches={chunk_gap['breach_count']}" + ) + print( + "Queue wait ms: " + f"min={format_num(queue_wait['min'])} " + f"p50={format_num(queue_wait['p50'])} " + f"p95={format_num(queue_wait['p95'])} " + f"max={format_num(queue_wait['max'])}" + ) + ratio = progressive["ratio"] + ratio_text = "n/a" if ratio is None else f"{ratio * 100:.2f}%" + print( + "Progressive streaming: " + f"{progressive['success_sessions']}/" + f"{progressive['eligible_sessions']} ({ratio_text})" + ) + print( + "Bandwidth Mbps: " + f"per_session_avg={format_num(per_session_bw['avg'])} " + f"per_session_p95={format_num(per_session_bw['p95'])} " + f"global={format_num(bandwidth['global_goodput_mbps'])} " + f"bucket_avg={format_num(bucket_bw['avg_mbps'])} " + f"bucket_peak={format_num(bucket_bw['peak_mbps'])}" + ) + print(f"VERDICT: {'PASS' if summary['passed'] else 'FAIL'}") + if summary["fail_reasons"]: + print("Fail reasons:") + for reason in summary["fail_reasons"]: + print(f"- {reason}") + + +def distribute_sessions(total_clients: int, process_count: int) -> list[int]: + base = total_clients // process_count + remainder = total_clients % process_count + counts = [] + for idx in range(process_count): + count = base + (1 if idx < remainder else 0) + counts.append(count) + return counts + + +def run_stress(args: argparse.Namespace) -> tuple[dict[str, Any], int]: + if websockets is None: + raise RuntimeError( + "Missing dependency: websockets. Install it before running this " + "stress test." + ) + + preset_file = Path(args.preset_file).expanduser().resolve() + selected_preset_id, curated_prompts, total_prompt_count = load_curated_prompts( + preset_file=preset_file, + preset_id=args.preset_id, + curated_limit=args.curated_limit, + ) + + process_count = args.processes + if process_count is None: + process_count = min(args.clients, os.cpu_count() or 1) + process_count = max(1, min(process_count, args.clients)) + + mp_ctx = mp.get_context("spawn") + start_event = mp_ctx.Event() + ready_queue = mp_ctx.Queue() + result_queue = mp_ctx.Queue() + + worker_config = { + "url": args.url, + "preset_id": selected_preset_id, + "curated_prompts": curated_prompts, + "connect_timeout_s": args.connect_timeout_s, + "session_timeout_s": args.session_timeout_s, + "post_complete_wait_s": args.post_complete_wait_s, + } + + session_counts = distribute_sessions(args.clients, process_count) + run_start_epoch = time.time() + run_start_monotonic = time.monotonic() + run_start_iso = iso_from_epoch(run_start_epoch) + + processes: list[mp.Process] = [] + for worker_id, session_count in enumerate(session_counts): + proc = mp_ctx.Process( + target=worker_entry, + args=( + worker_id, + session_count, + worker_config, + start_event, + ready_queue, + result_queue, + ), + ) + proc.start() + processes.append(proc) + + try: + ready_workers = 0 + ready_deadline = time.monotonic() + 60.0 + while ready_workers < len(processes): + timeout_s = max(0.1, ready_deadline - time.monotonic()) + if timeout_s <= 0: + raise RuntimeError("Timed out waiting for workers to become ready.") + msg = ready_queue.get(timeout=timeout_s) + if msg.get("status") == "ready": + ready_workers += 1 + + start_event.set() + + result_deadline = ( + time.monotonic() + + args.connect_timeout_s + + args.session_timeout_s + + args.post_complete_wait_s + + 180.0 + ) + worker_results: list[dict[str, Any]] = [] + while len(worker_results) < len(processes): + timeout_s = max(0.1, result_deadline - time.monotonic()) + if timeout_s <= 0: + break + try: + result = result_queue.get(timeout=timeout_s) + except Exception: + break + worker_results.append(result) + + for proc in processes: + proc.join(timeout=5.0) + if proc.is_alive(): + proc.terminate() + proc.join(timeout=2.0) + + sessions: list[dict[str, Any]] = [] + worker_errors: list[dict[str, Any]] = [] + for result in worker_results: + if result.get("status") == "ok": + sessions.extend(result.get("sessions", [])) + else: + worker_errors.append( + { + "worker_id": result.get("worker_id"), + "error": result.get("error"), + "traceback": result.get("traceback"), + } + ) + + received_workers = {result.get("worker_id") for result in worker_results} + expected_workers = set(range(len(processes))) + missing_workers = sorted(expected_workers - received_workers) + for worker_id in missing_workers: + worker_errors.append( + { + "worker_id": worker_id, + "error": "No worker result received.", + } + ) + + run_end_epoch = time.time() + run_end_iso = iso_from_epoch(run_end_epoch) + run_duration_ms = (time.monotonic() - run_start_monotonic) * 1000.0 + + summary = build_summary( + sessions=sessions, + chunk_gap_threshold_ms=args.chunk_gap_threshold_ms, + ) + + if worker_errors: + summary["passed"] = False + summary["fail_reasons"] = list(summary["fail_reasons"]) + [ + f"{len(worker_errors)} worker error(s) occurred." + ] + + output_payload = { + "run_info": { + "url": args.url, + "clients": args.clients, + "processes": process_count, + "preset_file": str(preset_file), + "preset_id": selected_preset_id, + "curated_limit": args.curated_limit, + "selected_prompt_count": len(curated_prompts), + "preset_total_prompt_count": total_prompt_count, + "chunk_gap_threshold_ms": args.chunk_gap_threshold_ms, + "connect_timeout_s": args.connect_timeout_s, + "session_timeout_s": args.session_timeout_s, + "post_complete_wait_s": args.post_complete_wait_s, + "run_start_ts_utc": run_start_iso, + "run_end_ts_utc": run_end_iso, + "run_duration_ms": run_duration_ms, + }, + "summary": summary, + "full_data": { + "sessions": sessions, + "worker_errors": worker_errors, + }, + } + exit_code = 0 if summary["passed"] else 1 + return output_payload, exit_code + finally: + for proc in processes: + if proc.is_alive(): + proc.terminate() + proc.join(timeout=1.0) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Multiprocess realtime stress test for LTX2 streaming.", + ) + parser.add_argument( + "-u", + "--url", + required=True, + help="WebSocket URL, e.g. wss://your-domain/ws", + ) + parser.add_argument( + "-c", + "--clients", + type=int, + required=True, + help="Total concurrent virtual users.", + ) + parser.add_argument( + "--processes", + type=int, + default=None, + help="Worker process count (default: min(clients, cpu_count)).", + ) + parser.add_argument( + "--preset-file", + default=str(DEFAULT_PRESET_FILE), + help="Path to curated presets JSON file.", + ) + parser.add_argument( + "--preset-id", + default=None, + help="Preset id to use (default: first preset in file).", + ) + parser.add_argument( + "--curated-limit", + type=int, + default=6, + help="Number of curated prompts to send from selected preset.", + ) + parser.add_argument( + "--chunk-gap-threshold-ms", + type=float, + default=5000.0, + help="Fail if any chunk gap is >= this value.", + ) + parser.add_argument( + "--post-complete-wait-s", + type=float, + default=5.0, + help="Seconds to wait after target segment completion before leave.", + ) + parser.add_argument( + "--connect-timeout-s", + type=float, + default=20.0, + help="WebSocket connect timeout in seconds.", + ) + parser.add_argument( + "--session-timeout-s", + type=float, + default=180.0, + help="Max session runtime per virtual user in seconds.", + ) + parser.add_argument( + "-o", + "--output-json", + required=True, + help="Required output JSON path (summary + full data).", + ) + args = parser.parse_args() + + if args.clients <= 0: + parser.error("--clients must be > 0") + if args.processes is not None and args.processes <= 0: + parser.error("--processes must be > 0") + if args.curated_limit <= 0: + parser.error("--curated-limit must be > 0") + if args.chunk_gap_threshold_ms <= 0: + parser.error("--chunk-gap-threshold-ms must be > 0") + if args.post_complete_wait_s < 0: + parser.error("--post-complete-wait-s must be >= 0") + if args.connect_timeout_s <= 0: + parser.error("--connect-timeout-s must be > 0") + if args.session_timeout_s <= 0: + parser.error("--session-timeout-s must be > 0") + return args + + +def main() -> int: + try: + args = parse_args() + print("Starting LTX2 realtime stress test with the following parameters:") + for arg, value in vars(args).items(): + print(f" {arg}: {value}") + output_payload, exit_code = run_stress(args) + + output_path = Path(args.output_json).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(output_payload, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + + print_summary( + run_info=output_payload["run_info"], + summary=output_payload["summary"], + ) + return exit_code + except SystemExit: + raise + except Exception as exc: + print(f"Error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/dreamverse/dreamverse/tests/test_session_init_image.py b/apps/dreamverse/dreamverse/tests/test_session_init_image.py new file mode 100644 index 000000000..52e351c0c --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_session_init_image.py @@ -0,0 +1,75 @@ +# pyright: reportMissingTypeArgument=false +import base64 +import io +from pathlib import Path + +from PIL import Image +import pytest + +from dreamverse.session_init_image import ( + MAX_SESSION_INIT_IMAGE_BYTES, + cleanup_session_init_image, + persist_session_init_image, +) + + +def make_data_url(format_name: str = "PNG") -> str: + image = Image.new("RGB", (2, 2), color=(16, 32, 64)) + buffer = io.BytesIO() + image.save(buffer, format=format_name) + encoded = base64.b64encode(buffer.getvalue()).decode("ascii") + mime_type = "image/png" if format_name == "PNG" else "image/jpeg" + return f"data:{mime_type};base64,{encoded}" + + +def test_persist_session_init_image_saves_normalized_file(tmp_path: Path): + session_image = persist_session_init_image( + { + "name": "frame.png", + "mime_type": "image/png", + "data_url": make_data_url("PNG"), + }, + temp_root=tmp_path, + ) + + assert session_image is not None + assert session_image.display_name == "frame.png" + assert session_image.file_path.is_file() + + cleanup_session_init_image(session_image) + assert not session_image.temp_dir.exists() + + +def test_persist_session_init_image_returns_none_when_missing_data(): + assert persist_session_init_image(None) is None + assert persist_session_init_image({"name": "frame.png", "data_url": ""}) is None + + +def test_persist_session_init_image_rejects_unsupported_mime(): + with pytest.raises(ValueError, match="PNG, JPEG, or WebP"): + persist_session_init_image( + { + "name": "frame.gif", + "mime_type": "image/gif", + "data_url": "data:image/gif;base64,R0lGODlhAQABAAAAACw=", + } + ) + + +def test_persist_session_init_image_rejects_large_payload(monkeypatch): + data_url = make_data_url("PNG") + oversized = "a" * (MAX_SESSION_INIT_IMAGE_BYTES + 1) + + def fake_b64decode(value: str, validate: bool = True): + return oversized.encode("ascii") + + monkeypatch.setattr(base64, "b64decode", fake_b64decode) + + with pytest.raises(ValueError, match="15 MB or smaller"): + persist_session_init_image( + { + "name": "frame.png", + "mime_type": "image/png", + "data_url": data_url, + } + ) diff --git a/apps/dreamverse/dreamverse/tests/test_session_logging.py b/apps/dreamverse/dreamverse/tests/test_session_logging.py new file mode 100644 index 000000000..33c643328 --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_session_logging.py @@ -0,0 +1,1276 @@ +from __future__ import annotations +# pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportMissingImports=false, reportMissingTypeArgument=false, reportOperatorIssue=false + +import asyncio +import base64 +import io +import json +import os +import re +import socket +from types import SimpleNamespace + +from fastapi import WebSocketDisconnect +from PIL import Image +import pytest + + +os.environ.setdefault("CEREBRAS_API_KEY", "dummy") +os.environ.setdefault("GROQ_API_KEY", "dummy") + +import dreamverse.main as server_main # noqa: E402 +import dreamverse.runtime as runtime # noqa: E402 +from dreamverse.session_logger import SessionEventLogger # noqa: E402 +from dreamverse.worker_ipc import MediaChunk, MediaComplete, MediaInit # noqa: E402 +from dreamverse.session import controller as session_controller # noqa: E402 + +pytestmark = pytest.mark.gpu + + +class _FakeSessionLogger: + def __init__(self): + self.events: list[dict[str, object]] = [] + + async def write_event( + self, + *, + event: str, + client_id: str, + payload: dict[str, object] | None = None, + ) -> None: + entry: dict[str, object] = { + "event": event, + "client_id": client_id, + } + if payload: + entry.update(payload) + self.events.append(entry) + + +class _FakeWebSocket: + def __init__(self, messages: list[tuple[float, dict[str, object]]]): + self._messages = messages + self._index = 0 + self.sent_json: list[dict[str, object]] = [] + self.sent_bytes: list[bytes] = [] + + async def accept(self) -> None: + return None + + async def send_json(self, payload: dict[str, object]) -> None: + self.sent_json.append(payload) + + async def send_bytes(self, payload: bytes) -> None: + self.sent_bytes.append(payload) + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + del code, reason + + async def receive_json(self) -> dict[str, object]: + if self._index >= len(self._messages): + raise WebSocketDisconnect() + delay_s, payload = self._messages[self._index] + self._index += 1 + if delay_s > 0: + await asyncio.sleep(delay_s) + return payload + + +class _FakeSlot: + def __init__(self, step_delay_s: float = 0.0): + self.shared_stream_buffer = None + self._stream_queues: dict[str, asyncio.Queue] = {} + self.calls: list[dict[str, object]] = [] + self._step_delay_s = max(0.0, float(step_delay_s)) + + async def join_user(self, client_id: str, model_id: str) -> None: + del client_id, model_id + + def register_stream_queue(self, client_id: str) -> asyncio.Queue: + queue: asyncio.Queue = asyncio.Queue() + self._stream_queues[client_id] = queue + return queue + + def unregister_stream_queue(self, client_id: str) -> None: + self._stream_queues.pop(client_id, None) + + async def user_step( + self, + client_id: str, + prompt: str, + segment_idx: int, + reset_conditioning: bool, + image_path: str | None = None, + ): + self.calls.append({ + "client_id": client_id, + "prompt": prompt, + "segment_idx": segment_idx, + "image_path": image_path, + "reset_conditioning": reset_conditioning, + }) + queue = self._stream_queues[client_id] + await queue.put(MediaInit( + user_id=client_id, + segment_idx=segment_idx, + stream_id=f"stream-{segment_idx}", + mime='video/mp4; codecs="avc1.640028,mp4a.40.2"', + uses_shared_buffer=False, + )) + await queue.put(MediaChunk( + user_id=client_id, + segment_idx=segment_idx, + stream_id=f"stream-{segment_idx}", + chunk=b"abc123", + uses_shared_buffer=False, + )) + await queue.put(MediaComplete( + user_id=client_id, + segment_idx=segment_idx, + stream_id=f"stream-{segment_idx}", + chunks=1, + )) + if self._step_delay_s > 0: + await asyncio.sleep(self._step_delay_s) + return {"e2e_latency_ms": 10.0} + + +class _FakeGPUPool: + def __init__(self, step_delay_s: float = 0.0): + self._slot = _FakeSlot(step_delay_s=step_delay_s) + + def get_status(self) -> dict[str, int]: + return { + "queue_size": 0, + "available_gpus": 1, + "total_gpus": 1, + } + + async def acquire(self, client_id: str, websocket) -> tuple[int, _FakeSlot]: + del client_id, websocket + return 0, self._slot + + async def release(self, client_id: str) -> None: + del client_id + + +def _make_png_data_url() -> str: + image = Image.new("RGB", (2, 2), color=(16, 32, 64)) + buffer = io.BytesIO() + image.save(buffer, format="PNG") + encoded = base64.b64encode(buffer.getvalue()).decode("ascii") + return f"data:image/png;base64,{encoded}" + + +class _FakePromptEnhancer: + def set_provider(self, provider: object, preferred_model: object = None): + del provider, preferred_model + return { + "provider": "openai", + "provider_options": ["openai", "cerebras"], + "options": ["gpt-5-nano-2025-08-07"], + } + + def get_rewrite_model_config(self): + return self.set_provider(None, None) + + def resolve_rewrite_model(self, value: object) -> str: + if isinstance(value, str) and value.strip(): + return value.strip() + return "gpt-5-nano-2025-08-07" + + def resolve_rewrite_system_prompt(self, value: object) -> str: + if isinstance(value, str) and value.strip(): + return value.strip() + return "shared rewrite system prompt" + + def resolve_rewrite_new_rollout_system_prompt(self, value: object) -> str: + if isinstance(value, str) and value.strip(): + return value.strip() + return "new rollout rewrite system prompt" + + def resolve_rewrite_temperature(self, value: float | None) -> float: + return 1.0 if value is None else float(value) + + async def enhance_prompt( + self, + raw_prompt: str, + *, + locked_segments: list[str], + next_segment_idx: int, + preset_id: str | None, + mode: str, + model: str, + timeout_ms: int, + ): + del locked_segments, next_segment_idx, preset_id, mode, model, timeout_ms + return SimpleNamespace( + prompt=f"enhanced: {raw_prompt}", + fallback_used=False, + error=None, + provider="openai", + model="gpt-5-nano-2025-08-07", + latency_ms=12.5, + ) + + async def rewrite_prompt_sequence( + self, + snapshot_prompts: list[str], + *, + preset_id: str | None, + preset_label: str | None, + rewrite_instruction: str, + rewrite_model: str, + rewrite_temperature: float, + timeout_ms: int, + system_prompt_override: str | None = None, + ): + del ( + snapshot_prompts, + preset_id, + preset_label, + rewrite_instruction, + rewrite_model, + rewrite_temperature, + timeout_ms, + system_prompt_override, + ) + return SimpleNamespace( + prompts=["rewrite a", "rewrite b"], + fallback_used=False, + error=None, + provider="openai", + model="gpt-5-nano-2025-08-07", + latency_ms=23.0, + rollout_id="rewrite_rollout", + rollout_label="Rewrite Rollout", + raw_response_text=( + '{"id":"rewrite_rollout","label":"Rewrite Rollout",' + '"segment_prompts":["rewrite a","rewrite b"]}' + ), + ) + + +class _FallbackPromptEnhancer(_FakePromptEnhancer): + async def enhance_prompt( + self, + raw_prompt: str, + *, + locked_segments: list[str], + next_segment_idx: int, + preset_id: str | None, + mode: str, + model: str, + timeout_ms: int, + ): + del ( + raw_prompt, + locked_segments, + next_segment_idx, + preset_id, + mode, + model, + timeout_ms, + ) + return SimpleNamespace( + prompt="", + fallback_used=True, + error="mock enhancement timeout", + provider="openai", + model="gpt-5-nano-2025-08-07", + latency_ms=18.0, + ) + + +class _UnsafeEnhancedPromptEnhancer(_FakePromptEnhancer): + async def enhance_prompt( + self, + raw_prompt: str, + *, + locked_segments: list[str], + next_segment_idx: int, + preset_id: str | None, + mode: str, + model: str, + timeout_ms: int, + ): + del ( + raw_prompt, + locked_segments, + next_segment_idx, + preset_id, + mode, + model, + timeout_ms, + ) + return SimpleNamespace( + prompt="unsafe enhanced prompt", + fallback_used=False, + error=None, + provider="openai", + model="gpt-5-nano-2025-08-07", + latency_ms=14.0, + ) + + +class _FakePromptSafetyFilter: + def __init__(self, blocked_prompts: dict[str, str]): + self._blocked_prompts = { + str(prompt).strip(): str(error) + for prompt, error in blocked_prompts.items() + if str(prompt).strip() + } + + def get_prompt_safety_error(self, prompt: str) -> str | None: + return self._blocked_prompts.get(str(prompt).strip()) + + def get_first_blocked_prompt(self, prompts: list[str]): + for index, prompt in enumerate(prompts): + error = self.get_prompt_safety_error(prompt) + if error is not None: + return SimpleNamespace(index=index, prompt=prompt, error=error) + return None + + +def test_session_event_logger_initializes_hostname_folder_and_utc_filename( + tmp_path, +): + logger = SessionEventLogger(tmp_path) + assert logger.directory == tmp_path / socket.gethostname() + assert logger.directory.is_dir() + assert re.match(r"^\d{6}_\d{6}\.jsonl$", logger.path.name) + assert logger.path.is_file() + + +def test_session_event_logger_appends_valid_jsonl(tmp_path): + logger = SessionEventLogger(tmp_path) + + async def _write() -> None: + await logger.write_event( + event="append_prompt", + client_id="client-123", + payload={"prompt": "hello world"}, + ) + + asyncio.run(_write()) + + lines = logger.path.read_text(encoding="utf-8").strip().splitlines() + assert len(lines) == 1 + record = json.loads(lines[0]) + assert record["event"] == "append_prompt" + assert record["client_id"] == "client-123" + assert record["hostname"] == socket.gethostname() + assert isinstance(record["ts"], str) + assert record["prompt"] == "hello world" + + +def test_websocket_flow_emits_required_session_events(): + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + try: + fake_logger = _FakeSessionLogger() + runtime.gpu_pool = _FakeGPUPool() + runtime.prompt_enhancer = _FakePromptEnhancer() + runtime.session_event_logger = fake_logger + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "curated_prompts": ["segment one prompt"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.01, {"type": "append_prompt", "prompt": "make it rainy"}), + ( + 0.01, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "more dramatic", + }, + ), + (0.30, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + event_names = [event["event"] for event in fake_logger.events] + assert "ws_session_start" in event_names + assert "gpu_assigned" in event_names + assert "append_prompt" in event_names + assert "segment_start" in event_names + assert "segment_complete" in event_names + + rewrite_events = [ + event + for event in fake_logger.events + if event["event"] == "rewrite_done" + ] + rewrite_kinds = {event.get("kind") for event in rewrite_events} + assert "seed_rewrite" in rewrite_kinds + assert "enhance_prompt" in rewrite_kinds + + segment_complete = next( + event + for event in fake_logger.events + if event["event"] == "segment_complete" + ) + latency = segment_complete.get("latency_ms") + assert isinstance(latency, dict) + assert set(latency.keys()) == { + "total", + "worker_e2e", + "main_user_step", + "overhead", + } + assert isinstance(segment_complete.get("data_size_bytes"), int) + assert segment_complete["data_size_bytes"] > 0 + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + + +def test_simple_generate_reuses_session_and_resets_conditioning(): + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool() + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = _FakePromptEnhancer() + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = None + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_prompt_1", + "curated_prompts": ["a selected prompt"], + "single_clip_mode": True, + "enhancement_enabled": False, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "simple_generate", + "preset_id": "simple_custom_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "a fresh custom prompt", + "enhancement_enabled": True, + "initial_image": { + "name": "frame.png", + "mime_type": "image/png", + "data_url": _make_png_data_url(), + }, + }, + ), + (0.50, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert [call["segment_idx"] for call in fake_pool._slot.calls] == [1, 1] + assert fake_pool._slot.calls[0]["prompt"] == "a selected prompt" + assert fake_pool._slot.calls[0]["reset_conditioning"] is False + assert fake_pool._slot.calls[0]["image_path"] is None + assert fake_pool._slot.calls[1]["prompt"] == "enhanced: a fresh custom prompt" + assert fake_pool._slot.calls[1]["reset_conditioning"] is True + assert isinstance(fake_pool._slot.calls[1]["image_path"], str) + + event_names = [event["event"] for event in fake_logger.events] + assert "simple_generate" in event_names + + stream_start_events = [ + payload + for payload in ws.sent_json + if payload.get("type") == "ltx2_stream_start" + ] + assert len(stream_start_events) == 2 + assert all( + payload.get("generation_segment_cap") == 0 + for payload in stream_start_events + ) + stream_complete_events = [ + payload + for payload in ws.sent_json + if payload.get("type") == "ltx2_stream_complete" + ] + assert len(stream_complete_events) == 2 + assert "prompt_sources_blocked" not in [ + payload.get("type") for payload in ws.sent_json + ] + + segment_start_events = [ + payload + for payload in ws.sent_json + if payload.get("type") == "ltx2_segment_start" + ] + assert [payload["segment_idx"] for payload in segment_start_events] == [1, 1] + + step_complete_events = [ + payload + for payload in ws.sent_json + if payload.get("type") == "step_complete" + ] + assert len(step_complete_events) == 2 + assert all( + set(payload.get("latency_ms", {}).keys()) == { + "total", + "worker_e2e", + "main_user_step", + "overhead", + } + for payload in step_complete_events + ) + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter + + +def test_single5s_ignores_global_segment_cap(): + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + old_generation_segment_cap = session_controller.GENERATION_SEGMENT_CAP + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool() + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = _FakePromptEnhancer() + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = None + session_controller.GENERATION_SEGMENT_CAP = 1 + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_prompt_1", + "curated_prompts": ["a selected prompt"], + "single_clip_mode": True, + "enhancement_enabled": False, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "simple_generate", + "preset_id": "simple_custom_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "a fresh custom prompt", + "enhancement_enabled": True, + "initial_image": None, + }, + ), + (0.50, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert [call["segment_idx"] for call in fake_pool._slot.calls] == [1, 1] + assert all( + payload.get("generation_segment_cap") == 0 + for payload in ws.sent_json + if payload.get("type") == "ltx2_stream_start" + ) + assert "generation_cap_reached" not in [ + payload.get("type") for payload in ws.sent_json + ] + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter + session_controller.GENERATION_SEGMENT_CAP = old_generation_segment_cap + + +def test_regular_segment_cap_completes_stream_and_allows_rewrite_rollout_restart(): + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + old_generation_segment_cap = session_controller.GENERATION_SEGMENT_CAP + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool() + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = _FakePromptEnhancer() + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = None + session_controller.GENERATION_SEGMENT_CAP = 1 + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": ["segment one prompt"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.20, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "start a new rollout", + }, + ), + (0.50, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert [call["segment_idx"] for call in fake_pool._slot.calls] == [1, 1] + assert fake_pool._slot.calls[0]["prompt"] == "segment one prompt" + assert fake_pool._slot.calls[1]["prompt"] == "rewrite a" + assert fake_pool._slot.calls[1]["reset_conditioning"] is True + + message_types = [payload["type"] for payload in ws.sent_json] + assert message_types.count("ltx2_stream_start") == 2 + assert message_types.count("ltx2_stream_complete") == 2 + assert "generation_cap_reached" not in message_types + assert "prompt_sources_blocked" not in message_types + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter + session_controller.GENERATION_SEGMENT_CAP = old_generation_segment_cap + + +def test_rewrite_during_active_segment_restarts_from_rewritten_seed_window(): + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool(step_delay_s=0.15) + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = _FakePromptEnhancer() + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = None + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": [ + "segment one prompt", + "segment two prompt", + ], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "restart from rewrite", + }, + ), + (0.45, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert [call["prompt"] for call in fake_pool._slot.calls[:2]] == [ + "segment one prompt", + "rewrite a", + ] + assert all( + call["prompt"] != "segment two prompt" + for call in fake_pool._slot.calls[1:] + ) + assert fake_pool._slot.calls[1]["segment_idx"] == 1 + assert fake_pool._slot.calls[1]["reset_conditioning"] is True + + reset_events = [ + payload + for payload in ws.sent_json + if payload.get("type") == "seed_prompts_reset_applied" + ] + assert reset_events + assert any( + payload.get("reason") == "rewrite_during_generation" + for payload in reset_events + ) + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter + + +def test_initial_custom_rollout_prompt_generates_seed_window_before_streaming(): + class _InitialRolloutPromptEnhancer(_FakePromptEnhancer): + def __init__(self): + self.rewrite_calls: list[dict[str, object]] = [] + + async def rewrite_prompt_sequence( + self, + snapshot_prompts: list[str], + *, + preset_id: str | None, + preset_label: str | None, + rewrite_instruction: str, + rewrite_model: str, + rewrite_temperature: float, + timeout_ms: int, + system_prompt_override: str | None = None, + ): + self.rewrite_calls.append({ + "snapshot_prompts": list(snapshot_prompts), + "preset_id": preset_id, + "preset_label": preset_label, + "rewrite_instruction": rewrite_instruction, + "rewrite_model": rewrite_model, + "rewrite_temperature": rewrite_temperature, + "timeout_ms": timeout_ms, + "system_prompt_override": system_prompt_override, + }) + del rewrite_model, rewrite_temperature, timeout_ms, system_prompt_override + return SimpleNamespace( + prompts=[ + "rewrite 1", + "rewrite 2", + "rewrite 3", + "rewrite 4", + "rewrite 5", + "rewrite 6", + ], + fallback_used=False, + error=None, + provider="openai", + model="gpt-5-nano-2025-08-07", + latency_ms=23.0, + rollout_id=preset_id or "custom_editable", + rollout_label=preset_label or "Custom rollout", + raw_response_text=( + '{"id":"custom_editable","label":"Custom rollout",' + '"segment_prompts":["rewrite 1","rewrite 2","rewrite 3",' + '"rewrite 4","rewrite 5","rewrite 6"]}' + ), + ) + + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool() + fake_enhancer = _InitialRolloutPromptEnhancer() + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = fake_enhancer + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = None + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "custom_editable", + "preset_label": "Custom rollout", + "curated_prompts": [], + "initial_rollout_prompt": "A moonbase corridor thriller with flooding", + "rewrite_model": "gpt-4.1-mini", + "rewrite_temperature": 0.4, + "rewrite_window_system_prompt": "Session-specific rewrite prompt", + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.20, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert len(fake_enhancer.rewrite_calls) == 1 + assert fake_enhancer.rewrite_calls[0] == { + "snapshot_prompts": [], + "preset_id": "custom_editable", + "preset_label": "Custom rollout", + "rewrite_instruction": "A moonbase corridor thriller with flooding", + "rewrite_model": "gpt-4.1-mini", + "rewrite_temperature": 0.4, + "timeout_ms": server_main.PROMPT_TIMEOUT_MS, + "system_prompt_override": "Session-specific rewrite prompt", + } + assert fake_pool._slot.calls + assert fake_pool._slot.calls[0]["prompt"] == "rewrite 1" + + message_types = [payload["type"] for payload in ws.sent_json] + assert "rewrite_seed_prompts_started" in message_types + assert "seed_prompts_updated" in message_types + assert "rewrite_seed_prompts_complete" in message_types + assert "seed_prompts_reset_applied" in message_types + assert "ltx2_stream_start" in message_types + assert "prompt_sources_blocked" not in message_types + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter + + +def test_initial_custom_rollout_prompt_uses_dedicated_user_prompt_by_default(): + class _InitialRolloutPromptEnhancer(_FakePromptEnhancer): + def __init__(self): + self.rewrite_calls: list[dict[str, object]] = [] + + async def rewrite_prompt_sequence( + self, + snapshot_prompts: list[str], + *, + preset_id: str | None, + preset_label: str | None, + rewrite_instruction: str, + rewrite_model: str, + rewrite_temperature: float, + timeout_ms: int, + system_prompt_override: str | None = None, + ): + self.rewrite_calls.append({ + "snapshot_prompts": list(snapshot_prompts), + "preset_id": preset_id, + "preset_label": preset_label, + "rewrite_instruction": rewrite_instruction, + "rewrite_model": rewrite_model, + "rewrite_temperature": rewrite_temperature, + "timeout_ms": timeout_ms, + "system_prompt_override": system_prompt_override, + }) + return SimpleNamespace( + prompts=[ + "rewrite 1", + "rewrite 2", + "rewrite 3", + "rewrite 4", + "rewrite 5", + "rewrite 6", + ], + fallback_used=False, + error=None, + provider="openai", + model="gpt-5-nano-2025-08-07", + latency_ms=23.0, + rollout_id=preset_id or "custom_editable", + rollout_label=preset_label or "Custom rollout", + raw_response_text=( + '{"id":"custom_editable","label":"Custom rollout",' + '"segment_prompts":["rewrite 1","rewrite 2","rewrite 3",' + '"rewrite 4","rewrite 5","rewrite 6"]}' + ), + ) + + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool() + fake_enhancer = _InitialRolloutPromptEnhancer() + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = fake_enhancer + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = None + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "custom_editable", + "preset_label": "Custom rollout", + "curated_prompts": [], + "initial_rollout_prompt": "A moonbase corridor thriller with flooding", + "rewrite_model": "gpt-4.1-mini", + "rewrite_temperature": 0.4, + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.20, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert len(fake_enhancer.rewrite_calls) == 1 + assert fake_enhancer.rewrite_calls[0]["snapshot_prompts"] == [] + assert ( + fake_enhancer.rewrite_calls[0]["system_prompt_override"] + == "new rollout rewrite system prompt" + ) + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter + + +def test_initial_custom_rollout_prompt_prefers_session_specific_user_prompt(): + class _InitialRolloutPromptEnhancer(_FakePromptEnhancer): + def __init__(self): + self.rewrite_calls: list[dict[str, object]] = [] + + async def rewrite_prompt_sequence( + self, + snapshot_prompts: list[str], + *, + preset_id: str | None, + preset_label: str | None, + rewrite_instruction: str, + rewrite_model: str, + rewrite_temperature: float, + timeout_ms: int, + system_prompt_override: str | None = None, + ): + self.rewrite_calls.append({ + "snapshot_prompts": list(snapshot_prompts), + "preset_id": preset_id, + "preset_label": preset_label, + "rewrite_instruction": rewrite_instruction, + "rewrite_model": rewrite_model, + "rewrite_temperature": rewrite_temperature, + "timeout_ms": timeout_ms, + "system_prompt_override": system_prompt_override, + }) + return SimpleNamespace( + prompts=[ + "rewrite 1", + "rewrite 2", + "rewrite 3", + "rewrite 4", + "rewrite 5", + "rewrite 6", + ], + fallback_used=False, + error=None, + provider="openai", + model="gpt-5-nano-2025-08-07", + latency_ms=23.0, + rollout_id=preset_id or "custom_editable", + rollout_label=preset_label or "Custom rollout", + raw_response_text=( + '{"id":"custom_editable","label":"Custom rollout",' + '"segment_prompts":["rewrite 1","rewrite 2","rewrite 3",' + '"rewrite 4","rewrite 5","rewrite 6"]}' + ), + ) + + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool() + fake_enhancer = _InitialRolloutPromptEnhancer() + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = fake_enhancer + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = None + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "custom_editable", + "preset_label": "Custom rollout", + "curated_prompts": [], + "initial_rollout_prompt": "A moonbase corridor thriller with flooding", + "rewrite_model": "gpt-4.1-mini", + "rewrite_temperature": 0.4, + "rewrite_window_system_prompt": "Window rewrite prompt", + "rewrite_user_system_prompt": "User rewrite prompt", + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.20, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert len(fake_enhancer.rewrite_calls) == 1 + assert fake_enhancer.rewrite_calls[0]["snapshot_prompts"] == [] + assert fake_enhancer.rewrite_calls[0]["system_prompt_override"] == "User rewrite prompt" + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter + + +def test_single5s_raw_prompt_block_prevents_generation(): + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool() + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = _FakePromptEnhancer() + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = _FakePromptSafetyFilter({ + "blocked raw prompt": "Test raw prompt blocked.", + }) + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_custom_prompt", + "curated_prompts": [], + "single_clip_mode": True, + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "append_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "blocked raw prompt", + }, + ), + (0.20, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert fake_pool._slot.calls == [] + assert { + payload["type"] for payload in ws.sent_json + }.isdisjoint({"ltx2_segment_start", "ltx2_segment_complete"}) + error_payloads = [ + payload for payload in ws.sent_json if payload.get("type") == "error" + ] + assert any( + payload.get("message") == "Test raw prompt blocked." + for payload in error_payloads + ) + + prompt_blocked_events = [ + event + for event in fake_logger.events + if event["event"] == "prompt_blocked" + ] + assert len(prompt_blocked_events) == 1 + assert prompt_blocked_events[0]["kind"] == "user_raw" + assert prompt_blocked_events[0]["raw_prompt"] == "blocked raw prompt" + assert all(event["event"] != "segment_start" for event in fake_logger.events) + assert all(event["event"] != "rewrite_done" for event in fake_logger.events) + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter + + +def test_single5s_enhanced_prompt_block_prevents_generation(): + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool() + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = _UnsafeEnhancedPromptEnhancer() + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = _FakePromptSafetyFilter({ + "unsafe enhanced prompt": "Test enhanced prompt blocked.", + }) + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_custom_prompt", + "curated_prompts": [], + "single_clip_mode": True, + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "append_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "safe raw prompt", + }, + ), + (0.20, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert fake_pool._slot.calls == [] + assert { + payload["type"] for payload in ws.sent_json + }.isdisjoint({"ltx2_segment_start", "ltx2_segment_complete"}) + error_payloads = [ + payload for payload in ws.sent_json if payload.get("type") == "error" + ] + assert any( + payload.get("message") == "Test enhanced prompt blocked." + for payload in error_payloads + ) + + rewrite_events = [ + event + for event in fake_logger.events + if event["event"] == "rewrite_done" + ] + assert len(rewrite_events) == 1 + assert rewrite_events[0]["kind"] == "enhance_prompt" + + prompt_blocked_events = [ + event + for event in fake_logger.events + if event["event"] == "prompt_blocked" + ] + assert len(prompt_blocked_events) == 1 + assert prompt_blocked_events[0]["kind"] == "user_enhanced" + assert prompt_blocked_events[0]["output_prompt"] == "unsafe enhanced prompt" + assert all(event["event"] != "segment_start" for event in fake_logger.events) + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter + + +def test_single5s_enhancement_fallback_does_not_start_generation(): + old_gpu_pool = runtime.gpu_pool + old_prompt_enhancer = runtime.prompt_enhancer + old_session_logger = runtime.session_event_logger + old_prompt_safety_filter = runtime.prompt_safety_filter + try: + fake_logger = _FakeSessionLogger() + fake_pool = _FakeGPUPool() + runtime.gpu_pool = fake_pool + runtime.prompt_enhancer = _FallbackPromptEnhancer() + runtime.session_event_logger = fake_logger + runtime.prompt_safety_filter = None + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_custom_prompt", + "curated_prompts": [], + "single_clip_mode": True, + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "append_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "prompt that times out", + }, + ), + (0.20, {"type": "leave"}), + ] + ) + + asyncio.run(server_main.websocket_endpoint(ws)) + + assert fake_pool._slot.calls == [] + assert { + payload["type"] for payload in ws.sent_json + }.isdisjoint({"ltx2_segment_start", "ltx2_segment_complete"}) + + fallback_payloads = [ + payload + for payload in ws.sent_json + if payload.get("type") == "prompt_fallback_used" + ] + assert len(fallback_payloads) == 1 + assert fallback_payloads[0]["prompt_id"] == "simple_custom_prompt" + assert ( + fallback_payloads[0]["error"] + == server_main.PROMPT_EXTENSION_FAILURE_USER_MESSAGE + ) + assert fallback_payloads[0]["source"] == "user_enhancement_failed" + + rewrite_events = [ + event + for event in fake_logger.events + if event["event"] == "rewrite_done" + ] + assert len(rewrite_events) == 1 + assert rewrite_events[0]["kind"] == "enhance_prompt" + assert all(event["event"] != "segment_start" for event in fake_logger.events) + assert all(event["event"] != "prompt_blocked" for event in fake_logger.events) + finally: + runtime.gpu_pool = old_gpu_pool + runtime.prompt_enhancer = old_prompt_enhancer + runtime.session_event_logger = old_session_logger + runtime.prompt_safety_filter = old_prompt_safety_filter