| 1 | """VLM-based vision selection for director memory reviews.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import base64 |
| 6 | import json |
| 7 | import re |
| 8 | import time |
| 9 | import urllib.error |
| 10 | import urllib.request |
| 11 | from pathlib import Path |
| 12 | from typing import Any, Callable |
| 13 | |
| 14 | from loguru import logger |
| 15 | |
| 16 | Transport = Callable[[str, dict[str, str], dict[str, Any], float], dict[str, Any]] |
| 17 | Sleeper = Callable[[float], None] |
| 18 | |
| 19 | |
| 20 | class MemorySelectorError(ValueError): |
| 21 | """The selector request or model response is invalid.""" |
| 22 | |
| 23 | |
| 24 | class SelectorTransportError(RuntimeError): |
| 25 | """The configured VLM provider request failed.""" |
| 26 | |
| 27 | def __init__( |
| 28 | self, |
| 29 | status_code: int, |
| 30 | message: str, |
| 31 | *, |
| 32 | retry_after_s: float | None = None, |
| 33 | ) -> None: |
| 34 | super().__init__(message) |
| 35 | self.status_code = int(status_code) |
| 36 | self.retryable = self.status_code == 429 |
| 37 | self.retry_after_s = retry_after_s |
| 38 | |
| 39 | |
| 40 | def _default_transport( |
| 41 | endpoint: str, |
| 42 | headers: dict[str, str], |
| 43 | payload: dict[str, Any], |
| 44 | timeout_s: float, |
| 45 | ) -> dict[str, Any]: |
| 46 | request = urllib.request.Request( |
| 47 | endpoint, |
| 48 | data=json.dumps(payload).encode("utf-8"), |
| 49 | headers=headers, |
| 50 | method="POST", |
| 51 | ) |
| 52 | try: |
| 53 | with urllib.request.urlopen(request, timeout=timeout_s) as response: |
| 54 | parsed = json.loads(response.read().decode("utf-8")) |
| 55 | except urllib.error.HTTPError as exc: |
| 56 | detail = exc.read().decode("utf-8", errors="replace") |
| 57 | retry_after = exc.headers.get("Retry-After") if exc.headers else None |
| 58 | try: |
| 59 | retry_after_s = float(retry_after) if retry_after is not None else None |
| 60 | except ValueError: |
| 61 | retry_after_s = None |
| 62 | raise SelectorTransportError( |
| 63 | exc.code, |
| 64 | detail or str(exc), |
| 65 | retry_after_s=retry_after_s, |
| 66 | ) from exc |
| 67 | except urllib.error.URLError as exc: |
| 68 | raise SelectorTransportError(0, str(exc.reason)) from exc |
| 69 | if not isinstance(parsed, dict): |
| 70 | raise MemorySelectorError("VLM provider response must be a JSON object") |
| 71 | return parsed |
| 72 | |
| 73 | |
| 74 | def _extract_json(text: str) -> dict[str, Any]: |
| 75 | value = (text or "").strip() |
| 76 | fenced = re.search(r"```(?:json)?\s*(\{.*\})\s*```", value, re.DOTALL) |
| 77 | if fenced: |
| 78 | value = fenced.group(1) |
| 79 | start, end = value.find("{"), value.rfind("}") |
| 80 | if start >= 0 and end > start: |
| 81 | value = value[start : end + 1] |
| 82 | try: |
| 83 | parsed = json.loads(value) |
| 84 | except json.JSONDecodeError as exc: |
| 85 | logger.warning("VLM returned non-JSON response: {}", (text or "")[:500]) |
| 86 | raise MemorySelectorError("VLM response is not valid JSON") from exc |
| 87 | if not isinstance(parsed, dict): |
| 88 | raise MemorySelectorError("VLM response must be a JSON object") |
| 89 | return parsed |
| 90 | |
| 91 | |
| 92 | class MemoryVlmSelector: |
| 93 | """Select reviewable memories using the configured VLM model.""" |
| 94 | |
| 95 | def __init__( |
| 96 | self, |
| 97 | *, |
| 98 | api_base: str, |
| 99 | api_key: str, |
| 100 | model: str, |
| 101 | transport: Transport = _default_transport, |
| 102 | sleeper: Sleeper = time.sleep, |
| 103 | timeout_s: float = 180, |
| 104 | max_attempts: int = 3, |
| 105 | ) -> None: |
| 106 | if not api_base.strip() or not api_key.strip() or not model.strip(): |
| 107 | raise MemorySelectorError( |
| 108 | "VLM provider apiBase, apiKey, and model are required for memory review" |
| 109 | ) |
| 110 | if max_attempts < 1: |
| 111 | raise MemorySelectorError("max_attempts must be positive") |
| 112 | self.model = model.strip() |
| 113 | self.endpoint = f"{api_base.rstrip('/')}/chat/completions" |
| 114 | self.headers = { |
| 115 | "Authorization": f"Bearer {api_key}", |
| 116 | "Content-Type": "application/json", |
| 117 | } |
| 118 | self.transport = transport |
| 119 | self.sleeper = sleeper |
| 120 | self.timeout_s = timeout_s |
| 121 | self.max_attempts = max_attempts |
| 122 | |
| 123 | def _complete(self, *, content: list[dict[str, Any]], max_tokens: int) -> dict[str, Any]: |
| 124 | body = { |
| 125 | "model": self.model, |
| 126 | "messages": [{"role": "user", "content": content}], |
| 127 | "max_tokens": max_tokens, |
| 128 | } |
| 129 | for attempt in range(1, self.max_attempts + 1): |
| 130 | try: |
| 131 | response = self.transport( |
| 132 | self.endpoint, self.headers, body, self.timeout_s) |
| 133 | logger.debug( |
| 134 | "VLM response model={} finish_reason={} content_len={}", |
| 135 | self.model, |
| 136 | response.get("choices", [{}])[0].get("finish_reason", ""), |
| 137 | len(response.get("choices", [{}])[0].get("message", {}).get("content", "") or ""), |
| 138 | ) |
| 139 | text = response["choices"][0]["message"]["content"] |
| 140 | if not isinstance(text, str): |
| 141 | raise TypeError |
| 142 | if not text.strip(): |
| 143 | logger.warning( |
| 144 | "VLM returned empty content model={} full_response={}", |
| 145 | self.model, |
| 146 | json.dumps(response, ensure_ascii=False)[:1000], |
| 147 | ) |
| 148 | return _extract_json(text) |
| 149 | except SelectorTransportError as exc: |
| 150 | if not exc.retryable or attempt >= self.max_attempts: |
| 151 | raise |
| 152 | self.sleeper( |
| 153 | exc.retry_after_s |
| 154 | if exc.retry_after_s is not None and exc.retry_after_s >= 0 |
| 155 | else 60.0 |
| 156 | ) |
| 157 | except (KeyError, IndexError, TypeError) as exc: |
| 158 | logger.warning( |
| 159 | "VLM response missing content model={} full_response={}", |
| 160 | self.model, |
| 161 | json.dumps(response, ensure_ascii=False)[:1000], |
| 162 | ) |
| 163 | raise MemorySelectorError("VLM response has no message content") from exc |
| 164 | raise AssertionError("unreachable") |
| 165 | |
| 166 | @staticmethod |
| 167 | def _available_candidates( |
| 168 | candidates: list[dict[str, Any]], rejected: set[int] | None |
| 169 | ) -> list[dict[str, Any]]: |
| 170 | rejected = rejected or set() |
| 171 | available = [item for item in candidates if int(item["candidate_index"]) not in rejected] |
| 172 | if not available: |
| 173 | raise MemorySelectorError("no memory candidates remain after exclusions") |
| 174 | return available |
| 175 | |
| 176 | @staticmethod |
| 177 | def _content(prompt: str, candidates: list[dict[str, Any]]) -> list[dict[str, Any]]: |
| 178 | content: list[dict[str, Any]] = [{"type": "text", "text": prompt}] |
| 179 | for item in candidates: |
| 180 | raw = Path(str(item["path"])).read_bytes() |
| 181 | encoded = base64.b64encode(raw).decode("ascii") |
| 182 | content.extend([ |
| 183 | {"type": "text", "text": ( |
| 184 | f"candidate_index={item['candidate_index']} " |
| 185 | f"frame_index={item['frame_index']} " |
| 186 | f"timestamp_sec={item['timestamp_sec']}")}, |
| 187 | {"type": "image_url", "image_url": { |
| 188 | "url": f"data:image/jpeg;base64,{encoded}"}}, |
| 189 | ]) |
| 190 | return content |
| 191 | |
| 192 | def select_characters( |
| 193 | self, |
| 194 | *, |
| 195 | shot_id: int, |
| 196 | caption: str, |
| 197 | character_ids: list[str], |
| 198 | candidates: list[dict[str, Any]], |
| 199 | rejected_candidate_indices: set[int] | None = None, |
| 200 | ) -> dict[str, Any]: |
| 201 | available = self._available_candidates(candidates, rejected_candidate_indices) |
| 202 | if len(available) > 8: |
| 203 | results = [ |
| 204 | self.select_characters( |
| 205 | shot_id=shot_id, |
| 206 | caption=caption, |
| 207 | character_ids=character_ids, |
| 208 | candidates=available[index:index + 8], |
| 209 | ) |
| 210 | for index in range(0, len(available), 8) |
| 211 | ] |
| 212 | best_by_id: dict[str, dict[str, Any]] = {} |
| 213 | for result in results: |
| 214 | for item in result["selections"]: |
| 215 | memory_id = str(item["character_id"]) |
| 216 | previous = best_by_id.get(memory_id) |
| 217 | if previous is None or float(item["confidence"]) > float( |
| 218 | previous["confidence"] |
| 219 | ): |
| 220 | best_by_id[memory_id] = item |
| 221 | selected: list[dict[str, Any]] = [] |
| 222 | used_candidates: set[int] = set() |
| 223 | for memory_id in character_ids: |
| 224 | item = best_by_id.get(memory_id) |
| 225 | if item is None or int(item["candidate_index"]) in used_candidates: |
| 226 | continue |
| 227 | used_candidates.add(int(item["candidate_index"])) |
| 228 | selected.append(item) |
| 229 | return { |
| 230 | "reasoning": " | ".join(str(result["reasoning"]) for result in results), |
| 231 | "selections": selected, |
| 232 | } |
| 233 | manifest = [{key: item[key] for key in ( |
| 234 | "candidate_index", "frame_index", "timestamp_sec")} for item in available] |
| 235 | prompt = ( |
| 236 | f"SHOT_ID: {shot_id}\nTARGET_CHARACTER_IDS: {character_ids}\n" |
| 237 | f"SHOT_CAPTION: {caption}\nCANDIDATES: {json.dumps(manifest)}\n\n" |
| 238 | "For each target ID, select one frame only when the exact person is clearly " |
| 239 | "identifiable. Never map one visible person to two IDs. Return only JSON: " |
| 240 | "{\"reasoning\":\"...\",\"selections\":[{\"character_id\":\"ID_A\"," |
| 241 | "\"candidate_index\":0,\"confidence\":0.0,\"target_only\":true," |
| 242 | "\"visible_character_ids\":[\"ID_A\"],\"reasoning\":\"...\"}]}.") |
| 243 | available_indices = {int(item["candidate_index"]) for item in available} |
| 244 | for attempt in range(2): |
| 245 | current_prompt = prompt |
| 246 | if attempt: |
| 247 | current_prompt += ( |
| 248 | "\nThe previous pass returned no character selections. Re-check every " |
| 249 | "candidate against the target descriptions and choose the clearest " |
| 250 | "identifiable frame for each visible target. A different person may be " |
| 251 | "partially visible at the edge; in that case set target_only=false and " |
| 252 | "list every visible ID. Return an empty selections list only when none " |
| 253 | "of the candidates visibly contains any target ID." |
| 254 | ) |
| 255 | result = self._complete( |
| 256 | content=self._content(current_prompt, available), |
| 257 | max_tokens=4000, |
| 258 | ) |
| 259 | self._validate_characters( |
| 260 | result, |
| 261 | set(character_ids), |
| 262 | available_indices, |
| 263 | ) |
| 264 | if result["selections"] or attempt == 1: |
| 265 | return result |
| 266 | raise AssertionError("unreachable") |
| 267 | |
| 268 | def decide_scene_transition( |
| 269 | self, |
| 270 | *, |
| 271 | previous_shot_id: int, |
| 272 | previous_caption: str, |
| 273 | next_shot_id: int, |
| 274 | next_caption: str, |
| 275 | ) -> dict[str, Any]: |
| 276 | """Decide whether the next outer shot changes the story scene.""" |
| 277 | prompt = ( |
| 278 | "Compare two adjacent OUTER shots from one screenplay. Decide whether " |
| 279 | "the next shot changes scene. A scene transition means a material change " |
| 280 | "in location, time, environment, or story situation. A camera cut, angle " |
| 281 | "change, framing change, or continued action in the same setting is NOT " |
| 282 | "a scene transition. Return only JSON: " |
| 283 | '{"scene_transition":true,"reasoning":"..."}.\n\n' |
| 284 | f"PREVIOUS_SHOT_ID: {previous_shot_id}\n" |
| 285 | f"PREVIOUS_CAPTION:\n{previous_caption.strip()}\n\n" |
| 286 | f"NEXT_SHOT_ID: {next_shot_id}\n" |
| 287 | f"NEXT_CAPTION:\n{next_caption.strip()}" |
| 288 | ) |
| 289 | result = self._complete( |
| 290 | content=[{"type": "text", "text": prompt}], |
| 291 | max_tokens=2000, |
| 292 | ) |
| 293 | if ( |
| 294 | not isinstance(result.get("scene_transition"), bool) |
| 295 | or not isinstance(result.get("reasoning"), str) |
| 296 | or not result["reasoning"].strip() |
| 297 | ): |
| 298 | raise MemorySelectorError("invalid VLM scene transition decision") |
| 299 | return { |
| 300 | "scene_transition": result["scene_transition"], |
| 301 | "reasoning": result["reasoning"].strip(), |
| 302 | } |
| 303 | |
| 304 | @staticmethod |
| 305 | def _validate_characters( |
| 306 | result: dict[str, Any], requested: set[str], available: set[int] |
| 307 | ) -> None: |
| 308 | selections = result.get("selections") |
| 309 | if not isinstance(result.get("reasoning"), str) or not isinstance(selections, list): |
| 310 | raise MemorySelectorError("invalid VLM character selection schema") |
| 311 | seen_ids: set[str] = set() |
| 312 | seen_candidates: set[int] = set() |
| 313 | for item in selections: |
| 314 | if not isinstance(item, dict): |
| 315 | raise MemorySelectorError("character selection must be an object") |
| 316 | character_id = item.get("character_id") |
| 317 | try: |
| 318 | candidate_index = int(item["candidate_index"]) |
| 319 | confidence = float(item["confidence"]) |
| 320 | except (KeyError, TypeError, ValueError) as exc: |
| 321 | raise MemorySelectorError("invalid character candidate fields") from exc |
| 322 | if (character_id not in requested or character_id in seen_ids |
| 323 | or candidate_index not in available or candidate_index in seen_candidates |
| 324 | or not 0 <= confidence <= 1 |
| 325 | or not isinstance(item.get("target_only"), bool) |
| 326 | or not isinstance(item.get("visible_character_ids"), list) |
| 327 | or not all(isinstance(value, str) for value in item["visible_character_ids"]) |
| 328 | or not isinstance(item.get("reasoning"), str) |
| 329 | or not item["reasoning"].strip()): |
| 330 | raise MemorySelectorError("invalid VLM character selection") |
| 331 | seen_ids.add(character_id) |
| 332 | seen_candidates.add(candidate_index) |
| 333 | |
| 334 | def select_representative( |
| 335 | self, |
| 336 | *, |
| 337 | shot_id: int, |
| 338 | caption: str, |
| 339 | candidates: list[dict[str, Any]], |
| 340 | rejected_candidate_indices: set[int] | None = None, |
| 341 | ) -> dict[str, Any]: |
| 342 | available = self._available_candidates(candidates, rejected_candidate_indices) |
| 343 | manifest = [{key: item[key] for key in ( |
| 344 | "candidate_index", "frame_index", "timestamp_sec")} for item in available] |
| 345 | prompt = ( |
| 346 | f"SHOT_ID: {shot_id}\nSHOT_CAPTION: {caption}\n" |
| 347 | f"CANDIDATES: {json.dumps(manifest)}\n\n" |
| 348 | "Select exactly one clear, information-rich continuity frame. Return only JSON: " |
| 349 | "{\"candidate_index\":0,\"confidence\":0.0,\"reasoning\":\"...\"}.") |
| 350 | result = self._complete(content=self._content(prompt, available), max_tokens=4000) |
| 351 | try: |
| 352 | candidate_index = int(result["candidate_index"]) |
| 353 | confidence = float(result["confidence"]) |
| 354 | except (KeyError, TypeError, ValueError) as exc: |
| 355 | raise MemorySelectorError("invalid VLM representative fields") from exc |
| 356 | if (candidate_index not in {int(item["candidate_index"]) for item in available} |
| 357 | or not 0 <= confidence <= 1 |
| 358 | or not isinstance(result.get("reasoning"), str) |
| 359 | or not result["reasoning"].strip()): |
| 360 | raise MemorySelectorError("invalid VLM representative selection") |
| 361 | return result |
| 362 | |
| 363 | def profile_image( |
| 364 | self, |
| 365 | *, |
| 366 | image_path: Path, |
| 367 | display_name: str = "", |
| 368 | ) -> dict[str, Any]: |
| 369 | """Create a compact, editable retrieval profile for one uploaded image.""" |
| 370 | raw = image_path.read_bytes() |
| 371 | encoded = base64.b64encode(raw).decode("ascii") |
| 372 | suffix = image_path.suffix.lower() |
| 373 | mime = "image/png" if suffix == ".png" else "image/webp" if suffix == ".webp" else "image/jpeg" |
| 374 | prompt = ( |
| 375 | "Describe this reusable video-reference asset for a director agent. " |
| 376 | "Be concise and factual. Include visible people, appearance, clothing, " |
| 377 | "objects, location, lighting, composition, and continuity cues when present. " |
| 378 | "Do not invent names or identity IDs. Return only JSON: " |
| 379 | '{"profile_text":"...","identity_ids":[]}.' |
| 380 | f"\nFILE_LABEL: {display_name.strip() or image_path.name}" |
| 381 | ) |
| 382 | result = self._complete( |
| 383 | content=[ |
| 384 | {"type": "text", "text": prompt}, |
| 385 | {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{encoded}"}}, |
| 386 | ], |
| 387 | max_tokens=1200, |
| 388 | ) |
| 389 | profile_text = result.get("profile_text") |
| 390 | identity_ids = result.get("identity_ids", []) |
| 391 | if ( |
| 392 | not isinstance(profile_text, str) |
| 393 | or not profile_text.strip() |
| 394 | or not isinstance(identity_ids, list) |
| 395 | or not all(isinstance(value, str) for value in identity_ids) |
| 396 | ): |
| 397 | raise MemorySelectorError("invalid VLM asset profile schema") |
| 398 | return { |
| 399 | "profile_text": profile_text.strip(), |
| 400 | "identity_ids": [value.strip() for value in identity_ids if value.strip()], |
| 401 | } |
| 402 |