| 1 | """Approve selected memories and prepare exactly one next R2V shot.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import hashlib |
| 6 | import json |
| 7 | import re |
| 8 | import shutil |
| 9 | from pathlib import Path |
| 10 | from typing import Any, Callable |
| 11 | |
| 12 | from nanobot.storage.files import configured_file_publisher |
| 13 | from nanobot.utils.helpers import write_json_atomic |
| 14 | |
| 15 | Publisher = Callable[[str, str], str] |
| 16 | |
| 17 | |
| 18 | def _read(path: Path) -> dict[str, Any]: |
| 19 | value = json.loads(path.read_text(encoding="utf-8")) |
| 20 | if not isinstance(value, dict): |
| 21 | raise ValueError(f"invalid JSON object: {path}") |
| 22 | return value |
| 23 | |
| 24 | |
| 25 | def _write(path: Path, value: dict[str, Any]) -> None: |
| 26 | write_json_atomic(path, value) |
| 27 | |
| 28 | |
| 29 | def stage_after_memory_advance(state: dict[str, Any], *, has_next: bool) -> str: |
| 30 | """Keep merge in flight. Memory approval must not clobber a pending merge.""" |
| 31 | pending = state.get("pending_remote_jobs") |
| 32 | kinds: set[str] = set() |
| 33 | if isinstance(pending, dict): |
| 34 | kinds = { |
| 35 | str(item.get("kind")) |
| 36 | for item in pending.values() |
| 37 | if isinstance(item, dict) |
| 38 | } |
| 39 | if "merge_shot" in kinds or str(state.get("stage") or "") == "merging": |
| 40 | return "merging" |
| 41 | if "generate_echo_shot" in kinds: |
| 42 | return "shot_generating" |
| 43 | return "shot_generating" if has_next else "shot_reviewing" |
| 44 | |
| 45 | |
| 46 | def publish_memory_record( |
| 47 | record: dict[str, Any], |
| 48 | *, |
| 49 | prefix: str, |
| 50 | publisher: Publisher, |
| 51 | ) -> dict[str, Any]: |
| 52 | published = dict(record) |
| 53 | image = str(record.get("local_image_path") or "").strip() |
| 54 | audio = str(record.get("local_audio_path") or "").strip() |
| 55 | source_shot_id = int(record.get("source_shot_id") or 0) |
| 56 | frame_index = int(record.get("frame_index") or 0) |
| 57 | object_prefix = ( |
| 58 | f"shot_{source_shot_id:03d}/" |
| 59 | f"{prefix}_frame_{frame_index:06d}" |
| 60 | ) |
| 61 | if image: |
| 62 | published["image_path"] = publisher(image, f"{object_prefix}.jpg") |
| 63 | if audio: |
| 64 | published["audio_path"] = publisher(audio, f"{object_prefix}.wav") |
| 65 | published.pop("local_image_path", None) |
| 66 | published.pop("local_audio_path", None) |
| 67 | if getattr(publisher, "delete_local_after_upload", False): |
| 68 | for path in (image, audio): |
| 69 | if path: |
| 70 | Path(path).unlink(missing_ok=True) |
| 71 | return published |
| 72 | |
| 73 | |
| 74 | def publish_memory_records( |
| 75 | records: dict[str, dict[str, Any]], |
| 76 | *, |
| 77 | publisher: Publisher, |
| 78 | ) -> dict[str, dict[str, Any]]: |
| 79 | """Publish a character bank while preserving its metadata and key order.""" |
| 80 | return { |
| 81 | memory_id: publish_memory_record(record, prefix=memory_id, publisher=publisher) |
| 82 | for memory_id, record in records.items() |
| 83 | if isinstance(record, dict) |
| 84 | } |
| 85 | |
| 86 | |
| 87 | def _character_ids(caption: str) -> list[str]: |
| 88 | return list(dict.fromkeys(re.findall( |
| 89 | r"(?<![A-Za-z0-9_-])ID_[A-Za-z0-9_-]+(?![A-Za-z0-9_-])", |
| 90 | caption, |
| 91 | ))) |
| 92 | |
| 93 | |
| 94 | def _memory_slots( |
| 95 | bank: dict[str, dict[str, Any]], |
| 96 | caption: str, |
| 97 | previous: dict[str, Any] | None, |
| 98 | ) -> list[dict[str, Any]]: |
| 99 | slots: list[dict[str, Any]] = [] |
| 100 | for memory_id in _character_ids(caption): |
| 101 | record = bank.get(memory_id) |
| 102 | if not isinstance(record, dict) or not record.get("image_path"): |
| 103 | continue |
| 104 | slot = { |
| 105 | "image_url": record["image_path"], |
| 106 | "metadata": { |
| 107 | "id": memory_id, |
| 108 | "workspace_asset_id": _automatic_asset_id( |
| 109 | record, memory_id=memory_id, kind="character" |
| 110 | ), |
| 111 | "source": "character_memory_bank", |
| 112 | "visual_status": record.get("visual_status", "confirmed"), |
| 113 | "source_shot_id": record.get("source_shot_id"), |
| 114 | "frame_index": record.get("frame_index"), |
| 115 | "timestamp_sec": record.get("timestamp_sec"), |
| 116 | "confidence": record.get("confidence"), |
| 117 | }, |
| 118 | } |
| 119 | if record.get("audio_path"): |
| 120 | slot["audio_url"] = record["audio_path"] |
| 121 | slot["metadata"]["audio_source_shot_id"] = record.get( |
| 122 | "audio_source_shot_id") |
| 123 | else: |
| 124 | slot["audio_mode"] = "empty" |
| 125 | slots.append(slot) |
| 126 | if previous and previous.get("image_path"): |
| 127 | slot = { |
| 128 | "image_url": previous["image_path"], |
| 129 | "metadata": { |
| 130 | "id": "PREVIOUS_SHOT", |
| 131 | "workspace_asset_id": _automatic_asset_id( |
| 132 | previous, memory_id="PREVIOUS_SHOT", kind="previous_shot" |
| 133 | ), |
| 134 | "source": "previous_shot_representative", |
| 135 | "visual_status": "representative", |
| 136 | "source_shot_id": previous.get("source_shot_id"), |
| 137 | "frame_index": previous.get("frame_index"), |
| 138 | "timestamp_sec": previous.get("timestamp_sec"), |
| 139 | "confidence": previous.get("confidence"), |
| 140 | }, |
| 141 | } |
| 142 | if previous.get("audio_path"): |
| 143 | slot["audio_url"] = previous["audio_path"] |
| 144 | slot["metadata"]["audio_source_shot_id"] = previous.get( |
| 145 | "audio_source_shot_id") |
| 146 | else: |
| 147 | slot["audio_mode"] = "empty" |
| 148 | slots.append(slot) |
| 149 | if len(slots) > 7: |
| 150 | raise ValueError("active character memory count exceeds 7") |
| 151 | return slots |
| 152 | |
| 153 | |
| 154 | def _automatic_asset_id( |
| 155 | raw: dict[str, Any], *, memory_id: str, kind: str |
| 156 | ) -> str: |
| 157 | """Match the stable id used by the WebSocket Memory Workspace.""" |
| 158 | fingerprint = json.dumps( |
| 159 | [ |
| 160 | memory_id, |
| 161 | int(raw.get("source_shot_id") or 0), |
| 162 | int(raw.get("frame_index") or 0), |
| 163 | kind, |
| 164 | ], |
| 165 | ensure_ascii=False, |
| 166 | separators=(",", ":"), |
| 167 | ) |
| 168 | return "auto_" + hashlib.sha256(fingerprint.encode("utf-8")).hexdigest()[:20] |
| 169 | |
| 170 | |
| 171 | def _remote_review_selections( |
| 172 | selections: Any, |
| 173 | bank: dict[str, dict[str, Any]], |
| 174 | previous: dict[str, Any] | None, |
| 175 | ) -> list[dict[str, Any]]: |
| 176 | projected: list[dict[str, Any]] = [] |
| 177 | for raw in selections if isinstance(selections, list) else []: |
| 178 | if not isinstance(raw, dict): |
| 179 | continue |
| 180 | memory_id = str(raw.get("memory_id") or "").strip() |
| 181 | source = previous if memory_id == "PREVIOUS_SHOT" else bank.get(memory_id) |
| 182 | item = dict(raw) |
| 183 | item.pop("local_image_path", None) |
| 184 | item.pop("local_audio_path", None) |
| 185 | if isinstance(source, dict): |
| 186 | for key in ("image_path", "audio_path"): |
| 187 | value = source.get(key) |
| 188 | if isinstance(value, str) and value.strip(): |
| 189 | item[key] = value.strip() |
| 190 | projected.append(item) |
| 191 | return projected |
| 192 | |
| 193 | |
| 194 | def _cleanup_local_review_artifacts( |
| 195 | work: Path, |
| 196 | shot_id: int, |
| 197 | review: dict[str, Any], |
| 198 | ) -> None: |
| 199 | review.pop("candidate_paths", None) |
| 200 | for attempt in review.get("history", []): |
| 201 | if not isinstance(attempt, dict): |
| 202 | continue |
| 203 | for item in attempt.get("selections", []): |
| 204 | if isinstance(item, dict): |
| 205 | item.pop("local_image_path", None) |
| 206 | item.pop("local_audio_path", None) |
| 207 | memory = work / "memory" |
| 208 | (memory / "videos" / f"shot_{shot_id:03d}.mp4").unlink(missing_ok=True) |
| 209 | shutil.rmtree(memory / "candidates" / f"shot_{shot_id:03d}", ignore_errors=True) |
| 210 | proposals = memory / "proposals" |
| 211 | for path in proposals.glob(f"shot_{shot_id:03d}_attempt_*"): |
| 212 | shutil.rmtree(path, ignore_errors=True) |
| 213 | |
| 214 | |
| 215 | def approve_review_and_prepare_next( |
| 216 | *, |
| 217 | workspace: Path, |
| 218 | work_id: str, |
| 219 | shot_id: int, |
| 220 | publisher: Publisher | None = None, |
| 221 | ) -> int | None: |
| 222 | work = workspace / "director" / "works" / work_id |
| 223 | current_path = work / "shots" / f"shot_{shot_id:03d}.json" |
| 224 | current = _read(current_path) |
| 225 | review = current.get("memory_review") |
| 226 | if not isinstance(review, dict) or review.get("status") != "approved": |
| 227 | raise ValueError("memory review must be approved before advancing") |
| 228 | proposed = review.get("proposed_bank") |
| 229 | if not isinstance(proposed, dict): |
| 230 | raise ValueError("approved memory review has no proposed bank") |
| 231 | |
| 232 | selections = [ |
| 233 | item for item in review.get("selections", []) if isinstance(item, dict) |
| 234 | ] |
| 235 | previous_raw = review.get("previous_shot") |
| 236 | if "retained_memory_ids" in review: |
| 237 | retained_ids = { |
| 238 | str(value) |
| 239 | for value in review.get("retained_memory_ids", []) |
| 240 | if value |
| 241 | } |
| 242 | bank_path = work / "memory" / "memory_bank.json" |
| 243 | current_bank = _read(bank_path) if bank_path.is_file() else {} |
| 244 | retained_proposed = dict(current_bank) |
| 245 | for item in selections: |
| 246 | memory_id = str(item.get("memory_id") or "") |
| 247 | record = proposed.get(memory_id) |
| 248 | if ( |
| 249 | memory_id |
| 250 | and memory_id != "PREVIOUS_SHOT" |
| 251 | and memory_id in retained_ids |
| 252 | and isinstance(record, dict) |
| 253 | ): |
| 254 | retained_proposed[memory_id] = record |
| 255 | proposed = retained_proposed |
| 256 | selections = [ |
| 257 | item |
| 258 | for item in selections |
| 259 | if str(item.get("memory_id") or "") in retained_ids |
| 260 | ] |
| 261 | if "PREVIOUS_SHOT" not in retained_ids: |
| 262 | previous_raw = None |
| 263 | |
| 264 | publisher = publisher or configured_file_publisher(work_id) |
| 265 | published_bank = publish_memory_records(proposed, publisher=publisher) |
| 266 | previous = ( |
| 267 | publish_memory_record( |
| 268 | previous_raw, |
| 269 | prefix=f"shot_{shot_id:03d}_PREVIOUS_SHOT", |
| 270 | publisher=publisher, |
| 271 | ) |
| 272 | if isinstance(previous_raw, dict) |
| 273 | else None |
| 274 | ) |
| 275 | review["proposed_bank"] = published_bank |
| 276 | review["previous_shot"] = previous |
| 277 | review["selections"] = _remote_review_selections( |
| 278 | selections, published_bank, previous |
| 279 | ) |
| 280 | if getattr(publisher, "delete_local_after_upload", False): |
| 281 | _cleanup_local_review_artifacts(work, shot_id, review) |
| 282 | current["memory_review"] = review |
| 283 | _write(current_path, current) |
| 284 | _write(work / "memory" / "memory_bank.json", published_bank) |
| 285 | if previous: |
| 286 | _write(work / "memory" / "previous_shot.json", previous) |
| 287 | else: |
| 288 | (work / "memory" / "previous_shot.json").unlink(missing_ok=True) |
| 289 | |
| 290 | next_path: Path | None = None |
| 291 | next_shot: dict[str, Any] | None = None |
| 292 | for path in sorted((work / "shots").glob("shot_*.json")): |
| 293 | candidate = _read(path) |
| 294 | candidate_id = int(candidate.get("shot_id") or 0) |
| 295 | if candidate_id <= shot_id: |
| 296 | continue |
| 297 | if str(candidate.get("status") or "") in { |
| 298 | "queued", "generated", "review_pass", "approved" |
| 299 | }: |
| 300 | continue |
| 301 | next_path = path |
| 302 | next_shot = candidate |
| 303 | break |
| 304 | |
| 305 | state_path = work / "state.json" |
| 306 | state = _read(state_path) |
| 307 | if next_path is None or next_shot is None: |
| 308 | state["stage"] = stage_after_memory_advance(state, has_next=False) |
| 309 | _write(state_path, state) |
| 310 | return None |
| 311 | |
| 312 | next_caption = str(next_shot.get("caption") or "") |
| 313 | # Extraction creates a recommendation draft only. In interactive mode the |
| 314 | # user must explicitly apply that draft in Build Memory before generation. |
| 315 | # An agent recommendation may replace this conservative identity-based |
| 316 | # draft later, but neither path may silently approve it. |
| 317 | if not next_shot.get("memory_slots_user_configured"): |
| 318 | recommended = _memory_slots(published_bank, next_caption, previous) |
| 319 | next_shot["recommended_memory_slots"] = recommended |
| 320 | next_shot["recommended_memory_display_slots"] = recommended |
| 321 | next_shot["memory_recommendation_source"] = ( |
| 322 | "profile_fallback" |
| 323 | if state.get("auto_generate") |
| 324 | else "pending_agent" |
| 325 | ) |
| 326 | next_shot.pop("approved_memory_slots", None) |
| 327 | next_shot.pop("approved_memory_display_slots", None) |
| 328 | _write(next_path, next_shot) |
| 329 | advanced_stage = stage_after_memory_advance(state, has_next=True) |
| 330 | state["stage"] = ( |
| 331 | "awaiting_memory_build" |
| 332 | if advanced_stage == "shot_generating" |
| 333 | else advanced_stage |
| 334 | ) |
| 335 | _write(state_path, state) |
| 336 | return int(next_shot["shot_id"]) |
| 337 | |
| 338 | |
| 339 | def auto_approve_review_and_prepare_next( |
| 340 | *, |
| 341 | workspace: Path, |
| 342 | work_id: str, |
| 343 | shot_id: int, |
| 344 | publisher: Publisher | None = None, |
| 345 | ) -> int | None: |
| 346 | """Approve a fresh Memory proposal and prepare its next shot exactly once.""" |
| 347 | work = workspace / "director" / "works" / work_id |
| 348 | current_path = work / "shots" / f"shot_{shot_id:03d}.json" |
| 349 | current = _read(current_path) |
| 350 | review = current.get("memory_review") |
| 351 | if not isinstance(review, dict): |
| 352 | raise ValueError("generated shot has no memory review") |
| 353 | |
| 354 | if review.get("auto_advance_complete") is True: |
| 355 | next_id = review.get("auto_advance_next_shot_id") |
| 356 | return int(next_id) if next_id is not None else None |
| 357 | if review.get("status") == "awaiting_review": |
| 358 | review["status"] = "approved" |
| 359 | current["memory_review"] = review |
| 360 | _write(current_path, current) |
| 361 | elif review.get("status") != "approved": |
| 362 | raise ValueError("memory review is not ready for automatic approval") |
| 363 | |
| 364 | next_id = approve_review_and_prepare_next( |
| 365 | workspace=workspace, |
| 366 | work_id=work_id, |
| 367 | shot_id=shot_id, |
| 368 | publisher=publisher, |
| 369 | ) |
| 370 | if next_id is not None: |
| 371 | next_path = work / "shots" / f"shot_{next_id:03d}.json" |
| 372 | next_shot = _read(next_path) |
| 373 | recommended = next_shot.get("recommended_memory_slots") |
| 374 | if isinstance(recommended, list): |
| 375 | next_shot["approved_memory_slots"] = recommended |
| 376 | next_shot["approved_memory_display_slots"] = next_shot.get( |
| 377 | "recommended_memory_display_slots", recommended |
| 378 | ) |
| 379 | next_shot["memory_slots_auto_approved"] = True |
| 380 | _write(next_path, next_shot) |
| 381 | state_path = work / "state.json" |
| 382 | state = _read(state_path) |
| 383 | if str(state.get("stage") or "") == "awaiting_memory_build": |
| 384 | state["stage"] = "shot_generating" |
| 385 | _write(state_path, state) |
| 386 | current = _read(current_path) |
| 387 | review = current.get("memory_review") |
| 388 | if isinstance(review, dict): |
| 389 | review["auto_advance_complete"] = True |
| 390 | review["auto_advance_next_shot_id"] = next_id |
| 391 | current["memory_review"] = review |
| 392 | _write(current_path, current) |
| 393 | return next_id |
| 394 |