| 1 | """First-frame reference image persisted on session metadata and workplace state.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from typing import Any |
| 6 | from urllib.parse import urlparse |
| 7 | |
| 8 | SESSION_REFERENCE_IMAGE_KEY = "reference_image" |
| 9 | SESSION_REFERENCE_IMAGE_LOCKED_KEY = "reference_image_locked" |
| 10 | SESSION_REFERENCE_IMAGE_REWRITE_KEY = "reference_image_needs_story_rewrite" |
| 11 | SESSION_STORY_REWRITE_SUPPRESS_KEY = "story_rewrite_suppressed" |
| 12 | |
| 13 | STORY_CONFIRM_INJECT_PREFIX = ( |
| 14 | "STORY_DIRECTION_CONFIRM lock=true. The user accepted the CURRENT story_md " |
| 15 | "on the confirmation card. Ignore reference_image_needs_story_rewrite even if " |
| 16 | "it is true. Do NOT rewrite plot, characters, wardrobe, or setting. Call " |
| 17 | "write_story with the existing story_md and confirmed=true. Then ask shot " |
| 18 | "count via ask_user. Do not show the story-direction card again." |
| 19 | ) |
| 20 | |
| 21 | _BLOCKED_HOSTS = frozenset({"127.0.0.1", "localhost", "::1", "0.0.0.0"}) |
| 22 | |
| 23 | |
| 24 | def normalize_reference_image(payload: Any) -> dict[str, Any] | None: |
| 25 | """Return a stored reference-image dict, or None when empty/invalid.""" |
| 26 | if not isinstance(payload, dict): |
| 27 | return None |
| 28 | url = ( |
| 29 | payload.get("url") |
| 30 | or payload.get("reference_image_url") |
| 31 | or payload.get("referenceImageUrl") |
| 32 | ) |
| 33 | if not isinstance(url, str) or not url.strip(): |
| 34 | return None |
| 35 | name = ( |
| 36 | payload.get("name") |
| 37 | or payload.get("reference_image_name") |
| 38 | or payload.get("referenceImageName") |
| 39 | or "" |
| 40 | ) |
| 41 | width = ( |
| 42 | payload.get("width") |
| 43 | or payload.get("reference_image_width") |
| 44 | or payload.get("referenceImageWidth") |
| 45 | or 0 |
| 46 | ) |
| 47 | height = ( |
| 48 | payload.get("height") |
| 49 | or payload.get("reference_image_height") |
| 50 | or payload.get("referenceImageHeight") |
| 51 | or 0 |
| 52 | ) |
| 53 | try: |
| 54 | width_i = int(width or 0) |
| 55 | except (TypeError, ValueError): |
| 56 | width_i = 0 |
| 57 | try: |
| 58 | height_i = int(height or 0) |
| 59 | except (TypeError, ValueError): |
| 60 | height_i = 0 |
| 61 | return { |
| 62 | "url": url.strip(), |
| 63 | "name": name.strip() if isinstance(name, str) else "", |
| 64 | "width": width_i, |
| 65 | "height": height_i, |
| 66 | } |
| 67 | |
| 68 | |
| 69 | def reference_image_present(payload: Any) -> bool: |
| 70 | ref = normalize_reference_image(payload) |
| 71 | return bool(ref and ref.get("url")) |
| 72 | |
| 73 | |
| 74 | def reference_image_url(payload: Any) -> str: |
| 75 | ref = normalize_reference_image(payload) |
| 76 | return str(ref.get("url") or "").strip() if ref else "" |
| 77 | |
| 78 | |
| 79 | def story_rewrite_suppressed(metadata: Any) -> bool: |
| 80 | return bool( |
| 81 | isinstance(metadata, dict) and metadata.get(SESSION_STORY_REWRITE_SUPPRESS_KEY) |
| 82 | ) |
| 83 | |
| 84 | |
| 85 | def suppress_story_rewrite(metadata: dict[str, Any]) -> None: |
| 86 | """Hide rewrite from the agent until the user picks 我想修改/增删参考图.""" |
| 87 | metadata[SESSION_STORY_REWRITE_SUPPRESS_KEY] = True |
| 88 | |
| 89 | |
| 90 | def allow_story_rewrite(metadata: dict[str, Any]) -> None: |
| 91 | metadata.pop(SESSION_STORY_REWRITE_SUPPRESS_KEY, None) |
| 92 | |
| 93 | |
| 94 | def apply_story_direction_answer(metadata: dict[str, Any], answer: str) -> str | None: |
| 95 | """Confirm locks the current story; edit-image re-enables rewrite. |
| 96 | |
| 97 | Returns an agent inject note for 可以，按这个来; otherwise None. |
| 98 | """ |
| 99 | from nanobot.agent.tools.ask_user import ( |
| 100 | is_reference_image_edit_option, |
| 101 | is_story_confirm_option, |
| 102 | ) |
| 103 | |
| 104 | if is_story_confirm_option(answer): |
| 105 | suppress_story_rewrite(metadata) |
| 106 | return STORY_CONFIRM_INJECT_PREFIX |
| 107 | if is_reference_image_edit_option(answer): |
| 108 | allow_story_rewrite(metadata) |
| 109 | return None |
| 110 | |
| 111 | |
| 112 | def reference_image_needs_story_rewrite(metadata: Any) -> bool: |
| 113 | if story_rewrite_suppressed(metadata): |
| 114 | return False |
| 115 | return bool( |
| 116 | isinstance(metadata, dict) and metadata.get(SESSION_REFERENCE_IMAGE_REWRITE_KEY) |
| 117 | ) |
| 118 | |
| 119 | |
| 120 | def mark_reference_image_needs_story_rewrite( |
| 121 | metadata: dict[str, Any], |
| 122 | *, |
| 123 | previous_url: str, |
| 124 | next_url: str, |
| 125 | ) -> bool: |
| 126 | """Flag a rewrite when an existing first-frame is replaced or removed. |
| 127 | |
| 128 | The flag is still stored while suppress is on so a later 我想修改/增删参考图 |
| 129 | can see the change; reads stay hidden until then. |
| 130 | """ |
| 131 | prev = (previous_url or "").strip() |
| 132 | nxt = (next_url or "").strip() |
| 133 | if prev == nxt or not prev: |
| 134 | return False |
| 135 | metadata[SESSION_REFERENCE_IMAGE_REWRITE_KEY] = True |
| 136 | return True |
| 137 | |
| 138 | |
| 139 | def clear_reference_image_needs_story_rewrite(metadata: Any) -> None: |
| 140 | if isinstance(metadata, dict): |
| 141 | metadata.pop(SESSION_REFERENCE_IMAGE_REWRITE_KEY, None) |
| 142 | |
| 143 | |
| 144 | def story_reference_image_inject_note(*, replaced: bool) -> str: |
| 145 | """Vision prompt attached with the current first-frame image.""" |
| 146 | note = ( |
| 147 | "The attached image is the user-uploaded CURRENT first-frame reference. " |
| 148 | "Write the screenplay and shot-1 opening so they continue from this image as frame 0." |
| 149 | ) |
| 150 | if replaced: |
| 151 | return ( |
| 152 | note |
| 153 | + " This image REPLACES a previous first-frame. Discard the previous " |
| 154 | "screenplay, characters, wardrobe, and setting unless they are visible " |
| 155 | "in THIS image. Do not paraphrase the old story. The confirmation card " |
| 156 | "must describe this new screenplay." |
| 157 | ) |
| 158 | return note |
| 159 | |
| 160 | |
| 161 | def is_blocked_local_url(url: str) -> bool: |
| 162 | """True when *url* is a local/file locator that must not be stored as a public address.""" |
| 163 | raw = (url or "").strip() |
| 164 | if not raw: |
| 165 | return True |
| 166 | parsed = urlparse(raw) |
| 167 | scheme = (parsed.scheme or "").lower() |
| 168 | if scheme in {"file", ""}: |
| 169 | return True |
| 170 | host = (parsed.hostname or "").lower() |
| 171 | if host in _BLOCKED_HOSTS: |
| 172 | return True |
| 173 | if host.startswith("127.") or host.endswith(".localhost"): |
| 174 | return True |
| 175 | return False |
| 176 | |
| 177 | |
| 178 | def is_reference_image_locked(state: dict[str, Any] | None) -> bool: |
| 179 | """Lock after shot_count is confirmed (PDF: 确认完分镜数后不可改).""" |
| 180 | if not isinstance(state, dict): |
| 181 | return False |
| 182 | if state.get(SESSION_REFERENCE_IMAGE_LOCKED_KEY) is True: |
| 183 | return True |
| 184 | goal = state.get("goal") if isinstance(state.get("goal"), dict) else {} |
| 185 | try: |
| 186 | shot_count = int(goal.get("shot_count") or 0) |
| 187 | except (TypeError, ValueError): |
| 188 | shot_count = 0 |
| 189 | return shot_count > 0 |
| 190 | |
| 191 | |
| 192 | def lock_reference_image(state: dict[str, Any]) -> None: |
| 193 | state[SESSION_REFERENCE_IMAGE_LOCKED_KEY] = True |
| 194 | |
| 195 | |
| 196 | def unlock_reference_image(state: dict[str, Any]) -> None: |
| 197 | state[SESSION_REFERENCE_IMAGE_LOCKED_KEY] = False |
| 198 | |
| 199 | |
| 200 | def download_reference_image_data_uri(url: str) -> str: |
| 201 | """Return a reference image as a data URI for LLM vision.""" |
| 202 | if url.startswith("data:image/"): |
| 203 | return url |
| 204 | from nanobot.security.http_download import HttpDownloadError, download_http_bytes |
| 205 | from nanobot.security.url_validator import UrlValidationError, validate_external_url |
| 206 | |
| 207 | if is_blocked_local_url(url): |
| 208 | raise ValueError("reference image url must be a public HTTP(S) address") |
| 209 | try: |
| 210 | validated = validate_external_url(url) |
| 211 | except UrlValidationError as exc: |
| 212 | raise ValueError(str(exc)) from exc |
| 213 | try: |
| 214 | result = download_http_bytes(validated, max_bytes=20 * 1024 * 1024, timeout_s=15.0) |
| 215 | except HttpDownloadError as exc: |
| 216 | raise ValueError(str(exc)) from exc |
| 217 | mime_type = (result.content_type or "").split(";", 1)[0].strip() |
| 218 | if (not mime_type) or mime_type == "application/octet-stream": |
| 219 | from os.path import splitext |
| 220 | from urllib.parse import urlparse |
| 221 | |
| 222 | ext = splitext(urlparse(validated).path)[1].lower() |
| 223 | mime_type = { |
| 224 | ".jpg": "image/jpeg", |
| 225 | ".jpeg": "image/jpeg", |
| 226 | ".png": "image/png", |
| 227 | ".webp": "image/webp", |
| 228 | ".gif": "image/gif", |
| 229 | ".bmp": "image/bmp", |
| 230 | }.get(ext, "image/jpeg") |
| 231 | import base64 |
| 232 | |
| 233 | b64 = base64.b64encode(result.data).decode("ascii") |
| 234 | return f"data:{mime_type};base64,{b64}" |
| 235 |