""" Shared infrastructure for the OCR scripts (novel_ocr.py, manga_ocr_llm.py). Everything here is genuinely identical between the two "page layout" scripts — config/state/logging plumbing, the OpenAI-compatible request helper, and the translate-only helpers (glossary, sliding-window context). What's deliberately NOT here is anything that depends on page layout (reading order, output format, the prompt files themselves) — that stays in each script, since novel and manga pages need genuinely different logic there. system/user prompt split (for provider-side prompt caching): Everything that stays IDENTICAL across every page of a run — the format/transcription instructions, and (for --translate) the glossary — is sent as the `system` message. Everything that changes per page — the image, and any sliding-window context — goes in `user`. Most OpenAI-compatible providers cache a repeated system prefix, so with a long book this avoids re-billing/re-processing the same instructions on every single page. Per-page context (previous/upcoming pages) is NOT cached-friendly by nature (it changes every page), so it belongs in `user`, never appended to the cacheable `system` text. """ import base64 import io import json import re import sys import time from datetime import datetime from pathlib import Path from openai import OpenAI from PIL import Image IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".tif", ".tiff", ".bmp"} STATE_FILE = "state.json" # --------------------------------------------------------------------------- # config.json # --------------------------------------------------------------------------- def load_config_section(config_path: Path, section: str) -> dict: """Loads the shared config.json and returns one top-level section (e.g. "ocr"). Falls back to treating the whole file as that section if the key is missing, so a bare {"api_key": ...} style file still works. """ if not config_path.exists(): return {} try: data = json.loads(config_path.read_text(encoding="utf-8")) except json.JSONDecodeError as e: print(f"Failed to parse {config_path}: {e}", file=sys.stderr) sys.exit(1) return data.get(section, data) if isinstance(data, dict) else {} def load_prompt_file(prompt_path: Path) -> str: if not prompt_path.exists(): print(f"Prompt file not found: {prompt_path}", file=sys.stderr) sys.exit(1) return prompt_path.read_text(encoding="utf-8").strip() # --------------------------------------------------------------------------- # Page selection (--pages "3,7-12,20", --start-page) # --------------------------------------------------------------------------- def parse_page_selection(spec: str, images: list) -> set | None: """Parses a --pages selection like "3,7-12,20" into a set of page stems. Numbers refer to the 1-based position in the naturally-sorted page list, so "3" is the third page, not a file literally named "3". Returns None when the spec is empty/absent, meaning "no explicit selection". """ if not spec: return None selected = set() for part in spec.split(","): part = part.strip() if not part: continue m = re.fullmatch(r"(\d+)\s*(?:-\s*(\d+))?", part) if not m: raise ValueError(f"Invalid --pages entry: {part!r} (expected N or N-M)") start = int(m.group(1)) end = int(m.group(2)) if m.group(2) else start if start < 1 or end < start: raise ValueError(f"Invalid --pages range: {part!r}") for n in range(start, end + 1): if n > len(images): raise ValueError(f"--pages refers to page {n}, but only {len(images)} page(s) were found") selected.add(images[n - 1].stem) return selected # --------------------------------------------------------------------------- # state.json manifest (resume/redo tracking, prompt-change detection) # --------------------------------------------------------------------------- def prompt_fingerprint(prompt: str) -> str: """Short stable hash of the prompt text, used to detect prompt edits.""" import hashlib return hashlib.sha256(prompt.encode("utf-8")).hexdigest()[:12] def load_state(output_dir: Path) -> dict: """Loads the pages manifest, mapping page stem -> what produced it. Returns {} if the manifest is missing or unreadable; callers must treat a page with no entry as "unverified" rather than "stale". """ path = output_dir / STATE_FILE if not path.exists(): return {} try: data = json.loads(path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError) as e: # noqa: BLE001 print(f"Warning: could not read {path} ({e}); existing pages are reused unverified.", file=sys.stderr) return {} return data.get("pages", {}) if isinstance(data, dict) else {} def save_state(output_dir: Path, pages: dict, meta: dict) -> None: """Atomically rewrites the pages manifest (temp file + replace).""" path = output_dir / STATE_FILE tmp = path.with_suffix(".json.tmp") tmp.write_text( json.dumps({"meta": meta, "pages": pages}, ensure_ascii=False, indent=2, default=str), encoding="utf-8", ) tmp.replace(path) def warn_if_prompt_changed(state: dict, images: list, prompt_hash: str, prompt_file_name: str) -> None: """Warns (does not fail) if some already-transcribed pages were done with a prompt whose hash no longer matches the current one.""" if not state: return stale = [ p.stem for p in images if p.stem in state and state[p.stem].get("prompt_hash") not in (None, prompt_hash) ] if stale: print( f"Note: {len(stale)} page(s) were transcribed with a different prompt " f"({prompt_file_name} changed). Use --redo to re-OCR them.", file=sys.stderr, ) # --------------------------------------------------------------------------- # Request/response logging # --------------------------------------------------------------------------- def write_log(log_dir: Path, page_name: str, entry: dict) -> None: """Writes one JSON log file per request/response. The base64 image payload itself is never included (megabytes of no debugging value, one per page) — only its path/size — but everything else (full prompt text, model params, every retry attempt, the full raw API response) is recorded as-is. """ log_dir.mkdir(parents=True, exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") safe_name = re.sub(r"[^A-Za-z0-9_.-]", "_", page_name) log_path = log_dir / f"{timestamp}_{safe_name}.json" log_path.write_text(json.dumps(entry, ensure_ascii=False, indent=2, default=str), encoding="utf-8") # --------------------------------------------------------------------------- # Request params (CLI flag > config.json > default) # --------------------------------------------------------------------------- def build_request_kwargs(config: dict, args) -> tuple[dict, dict]: """Resolves OpenAI-API-style request params: CLI flag > config.json > a sensible default. Returns (kwargs, extra_body) — standard params go straight into the request; reasoning_effort goes through extra_body since support for it varies by provider/model and extra_body is the designed passthrough for exactly that. """ temperature = args.temperature if args.temperature is not None else config.get("temperature", 0) max_tokens = args.max_tokens if args.max_tokens is not None else config.get("max_tokens") top_p = args.top_p if args.top_p is not None else config.get("top_p") reasoning_effort = args.reasoning_effort or config.get("reasoning_effort") kwargs = {"temperature": temperature} if max_tokens is not None: kwargs["max_tokens"] = max_tokens if top_p is not None: kwargs["top_p"] = top_p extra_body = {} if reasoning_effort: extra_body["reasoning_effort"] = reasoning_effort return kwargs, extra_body # --------------------------------------------------------------------------- # Image encoding # --------------------------------------------------------------------------- def image_to_data_url(path: Path, max_dim: int = 2200) -> str: """Downscale (if needed) and encode an image as a base64 data URL.""" img = Image.open(path) if img.mode not in ("RGB", "L"): img = img.convert("RGB") if max(img.size) > max_dim: ratio = max_dim / max(img.size) img = img.resize((int(img.width * ratio), int(img.height * ratio))) buf = io.BytesIO() img.save(buf, format="JPEG", quality=90) b64 = base64.b64encode(buf.getvalue()).decode("utf-8") return f"data:image/jpeg;base64,{b64}" # --------------------------------------------------------------------------- # LLM OCR call, with system/user split for prompt caching # --------------------------------------------------------------------------- def usage_dict(response) -> dict | None: """Extracts token usage from a response, tolerating providers that omit it.""" usage = getattr(response, "usage", None) if usage is None: return None try: data = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) except Exception: # noqa: BLE001 return None return {k: v for k, v in data.items() if v is not None} or None def ocr_image( client: OpenAI, model: str, system_prompt: str, path: Path, request_kwargs: dict, extra_body: dict, user_text: str = None, retries: int = 3, log_entry: dict | None = None, ) -> tuple[str, dict]: """OCRs (or translates) one page image. system_prompt: the part of the prompt that's IDENTICAL across every page of this run (format instructions, glossary) — sent as the `system` message so OpenAI-compatible providers can cache it instead of reprocessing it on every page. user_text: optional per-page text that varies from page to page (sliding-window context) — sent alongside the image in `user`. Never put per-page content in system_prompt, or every page invalidates the cache instead of hitting it. Returns (text, meta) where meta records what the provider reported about the accepted response (finish_reason, usage) so the caller can persist it (e.g. in the pages manifest). """ data_url = image_to_data_url(path) user_content = [{"type": "image_url", "image_url": {"url": data_url}}] if user_text: user_content.insert(0, {"type": "text", "text": user_text}) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_content}, ] if log_entry is not None: logged_user_content = [ {"type": "text", "text": user_text} if user_text else None, {"type": "image_url", "image_url": {"url": f""}}, ] log_entry["request"] = { "model": model, **request_kwargs, "extra_body": extra_body or None, "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": [c for c in logged_user_content if c]}, ], } log_entry["attempts"] = [] last_err = None for attempt in range(retries): try: response = client.chat.completions.create( model=model, messages=messages, extra_body=extra_body or None, **request_kwargs, ) choice = response.choices[0] text = (choice.message.content or "").strip() finish_reason = getattr(choice, "finish_reason", None) # A truncated or empty response must not be accepted as a # successful transcription — this file would otherwise be # written to disk and silently skipped on every future retry. if finish_reason == "length": raise RuntimeError(f"response truncated (finish_reason=length, {len(text)} chars)") if not text: raise RuntimeError(f"empty response (finish_reason={finish_reason})") if log_entry is not None: try: raw_response = response.model_dump() except Exception: # noqa: BLE001 raw_response = {"content": text} log_entry["attempts"].append({"attempt": attempt + 1, "success": True, "response": raw_response}) return text, {"finish_reason": finish_reason, "usage": usage_dict(response)} except Exception as e: # noqa: BLE001 last_err = e if log_entry is not None: log_entry["attempts"].append({"attempt": attempt + 1, "success": False, "error": str(e)}) time.sleep(3 * (attempt + 1)) raise RuntimeError(f"Failed to OCR {path.name} after {retries} attempts: {last_err}") # --------------------------------------------------------------------------- # Translation helpers: glossary + sliding-window context (--translate only) # --------------------------------------------------------------------------- def find_glossary(script_dir: Path, explicit_path: str) -> Path | None: """Resolves the glossary path: explicit --glossary if given, otherwise an auto-detected glossary.md next to the ocr scripts themselves (a stable, book-independent location — handy for a future UI to edit).""" if explicit_path: path = Path(explicit_path) if not path.exists(): print(f"Glossary file not found: {path}", file=sys.stderr) sys.exit(1) return path auto_path = script_dir / "glossary.md" return auto_path if auto_path.exists() else None def build_system_prompt(base_prompt: str, target_language: str = None, glossary: str = None) -> str: """Builds the cacheable `system` prompt: format instructions, with the target language substituted in and the glossary appended. Deliberately excludes per-page context (see build_context_block) — that changes every page and must go in `user`, or every page would invalidate the provider's cached system prefix instead of hitting it. """ prompt = base_prompt if target_language: prompt = prompt.replace("{target_language}", target_language) if glossary: prompt += ( "\n\n---\n\n" "Glossary: use these established translations for the following " "names/terms for consistency. Anything not listed here, translate " "naturally.\n\n" + glossary ) return prompt def build_context_block(past_pages: list, ahead_pages: list) -> str: """Builds the per-page `user` text carrying sliding-window context: - past_pages: the last N *translated* pages, for continuity (character voice, an ongoing exchange, pronouns/referents that only make sense given what was just said). - ahead_pages: the next N pages' raw *Japanese* transcription (never a draft translation — see the ocr/README.md design note on why), for context that only becomes clear from what's about to happen (a twist, who's actually being addressed, a joke that pays off a page later). Both are lists of (page_name, text) tuples, oldest/nearest first. Returns "" if there's nothing to add. This text is per-page by nature and must never be folded into the cacheable system prompt (see build_system_prompt) — it changes on every single page, so appending it to system would defeat prompt caching entirely instead of the intended "cache the instructions, vary only the image/context" split. """ if not past_pages and not ahead_pages: return "" parts = [] if past_pages: parts.append( "Context: translations of the immediately preceding page(s), for " "continuity only (character voice, an ongoing conversation, " "pronouns/referents). Do NOT re-translate or repeat any of this " "in your output — translate only the NEW page shown in the image.\n" ) for name, text in past_pages: parts.append(f"\n[Previous page: {name}]\n{text}") if ahead_pages: if parts: parts.append("\n\n---\n\n") parts.append( "Context: the ORIGINAL JAPANESE (not a translation) of the " "upcoming page(s) that follow the one you're translating now. " "Use this only to correctly resolve things that depend on what " "happens next — pronoun/referent gender, who a line is actually " "addressed to, a setup whose payoff lands later. Do NOT " "translate or otherwise output any of this — translate only the " "CURRENT page shown in the image.\n" ) for name, text in ahead_pages: parts.append(f"\n[Upcoming page ({name}), original Japanese]\n{text}") return "".join(parts)