#!/usr/bin/env python3 """ Extract characters, locations, and props from a screenplay PDF using OpenAI. Requirements: - OPENAI_API_KEY must be set in the environment. - Install dependencies from screenplay/requirements.txt. Example: python3 screenplay/extract_entities.py \ --input "screenplay/srd part 1 blue revision.pdf" \ --output "screenplay/srd_part_1_entities.json" """ from __future__ import annotations import argparse import hashlib import json import os import sqlite3 import sys import textwrap import time import urllib.error import urllib.request from dataclasses import dataclass from pathlib import Path from typing import Dict, Iterable, List, Sequence, Tuple DEFAULT_OPENAI_MODEL = "gpt-5.4" DEFAULT_OPENAI_FALLBACK_MODEL = "gpt-5" DEFAULT_GEMINI_MODEL = "gemini-3.1-pro-preview" DEFAULT_GEMINI_FALLBACK_MODEL = "gemini-3-flash-preview" OPENAI_API_URL = "https://api.openai.com/v1/responses" GEMINI_API_URL_TEMPLATE = "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}" API_TIMEOUT_SECONDS = 900 SCRIPT_DIR = Path(__file__).resolve().parent PROMPTS_DIR = SCRIPT_DIR / "prompts" PROMPTS_MANIFEST = PROMPTS_DIR / "manifest.json" SUPPORTED_LANGUAGE_MAP = { "ko": "Korean", "ja": "Japanese", "en": "English", } RELATION_FACT_SCHEMA: Dict[str, object] = { "type": "object", "additionalProperties": False, "properties": { "relation_family": { "type": "string", "enum": [ "identity", "kinship", "social", "conflict", "collaboration", "possession", "containment", "location", "membership", "control", "goal", "event", "state", "transformation", "other", ], }, "relation_type": {"type": "string"}, "directionality": { "type": "string", "enum": ["directed", "bidirectional", "undirected"], }, "temporal_scope": { "type": "string", "enum": ["scene", "episode", "series", "backstory", "unknown"], }, "continuity_priority": { "type": "string", "enum": ["critical", "high", "medium", "low"], }, "continuity_reason": {"type": "string"}, "participants": { "type": "array", "minItems": 2, "maxItems": 5, "items": { "type": "object", "additionalProperties": False, "properties": { "entity_name": {"type": "string"}, "entity_type": { "type": "string", "enum": ["character", "location", "prop"], }, "role": {"type": "string"}, }, "required": ["entity_name", "entity_type", "role"], }, }, "evidence": {"type": "array", "items": {"type": "string"}}, }, "required": [ "relation_family", "relation_type", "directionality", "temporal_scope", "continuity_priority", "continuity_reason", "participants", "evidence", ], } CHUNK_SCHEMA: Dict[str, object] = { "type": "object", "additionalProperties": False, "properties": { "characters": { "type": "array", "items": { "type": "object", "additionalProperties": False, "properties": { "name": {"type": "string"}, "aliases": {"type": "array", "items": {"type": "string"}}, "description": {"type": "string"}, "continuity_reason": {"type": "string"}, "visual_anchor_traits": {"type": "array", "items": {"type": "string"}}, "variant_axes": {"type": "array", "items": {"type": "string"}}, "importance": { "type": "string", "enum": ["major", "supporting", "minor", "unknown"], }, "continuity_priority": { "type": "string", "enum": ["critical", "high", "medium", "low"], }, "reference_image_priority": { "type": "string", "enum": ["required", "helpful", "not_needed"], }, "evidence": {"type": "array", "items": {"type": "string"}}, }, "required": [ "name", "aliases", "description", "continuity_reason", "visual_anchor_traits", "variant_axes", "importance", "continuity_priority", "reference_image_priority", "evidence", ], }, }, "locations": { "type": "array", "items": { "type": "object", "additionalProperties": False, "properties": { "name": {"type": "string"}, "aliases": {"type": "array", "items": {"type": "string"}}, "description": {"type": "string"}, "continuity_reason": {"type": "string"}, "visual_anchor_traits": {"type": "array", "items": {"type": "string"}}, "variant_axes": {"type": "array", "items": {"type": "string"}}, "kind": { "type": "string", "enum": ["interior", "exterior", "mixed", "unknown"], }, "importance": { "type": "string", "enum": ["major", "supporting", "minor", "unknown"], }, "continuity_priority": { "type": "string", "enum": ["critical", "high", "medium", "low"], }, "reference_image_priority": { "type": "string", "enum": ["required", "helpful", "not_needed"], }, "evidence": {"type": "array", "items": {"type": "string"}}, }, "required": [ "name", "aliases", "description", "continuity_reason", "visual_anchor_traits", "variant_axes", "kind", "importance", "continuity_priority", "reference_image_priority", "evidence", ], }, }, "props": { "type": "array", "items": { "type": "object", "additionalProperties": False, "properties": { "name": {"type": "string"}, "aliases": {"type": "array", "items": {"type": "string"}}, "description": {"type": "string"}, "continuity_reason": {"type": "string"}, "visual_anchor_traits": {"type": "array", "items": {"type": "string"}}, "variant_axes": {"type": "array", "items": {"type": "string"}}, "significance": { "type": "string", "enum": ["key", "recurring", "minor", "unknown"], }, "continuity_priority": { "type": "string", "enum": ["critical", "high", "medium", "low"], }, "reference_image_priority": { "type": "string", "enum": ["required", "helpful", "not_needed"], }, "evidence": {"type": "array", "items": {"type": "string"}}, }, "required": [ "name", "aliases", "description", "continuity_reason", "visual_anchor_traits", "variant_axes", "significance", "continuity_priority", "reference_image_priority", "evidence", ], }, }, "relation_facts": { "type": "array", "items": RELATION_FACT_SCHEMA, }, "notes": {"type": "array", "items": {"type": "string"}}, }, "required": ["characters", "locations", "props", "relation_facts", "notes"], } FINAL_SCHEMA: Dict[str, object] = { "type": "object", "additionalProperties": False, "properties": { "source_file": {"type": "string"}, "characters": CHUNK_SCHEMA["properties"]["characters"], "locations": CHUNK_SCHEMA["properties"]["locations"], "props": CHUNK_SCHEMA["properties"]["props"], "relation_facts": CHUNK_SCHEMA["properties"]["relation_facts"], "summary": {"type": "string"}, "notes": {"type": "array", "items": {"type": "string"}}, }, "required": [ "source_file", "characters", "locations", "props", "relation_facts", "summary", "notes", ], } @dataclass class TextChunk: page_start: int page_end: int text: str @dataclass class LanguageInfo: code: str name: str @dataclass class PromptBundle: version: str chunk_system: str chunk_user: str final_system: str final_user: str manifest_path: str version_path: str prompt_files: Dict[str, str] prompt_hashes: Dict[str, str] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Extract screenplay entities to JSON.") parser.add_argument( "--provider", default="openai", choices=["openai", "gemini"], help="LLM provider to use.", ) parser.add_argument( "--input", default="screenplay/srd part 1 blue revision.pdf", help="Path to the screenplay PDF.", ) parser.add_argument( "--output", default="screenplay/srd_part_1_entities.json", help="Path to write the extracted JSON.", ) parser.add_argument( "--model", default=None, help="Model to use. Defaults depend on --provider.", ) parser.add_argument( "--fallback-model", default=None, help="Fallback model to try only if the primary model is unavailable.", ) parser.add_argument( "--prompt-version", default=None, help="Prompt version to use. Defaults to the manifest current_version.", ) parser.add_argument( "--source-language", default="auto", choices=["auto", "ko", "ja", "en"], help="Force the screenplay language instead of auto detection.", ) parser.add_argument( "--mode", default="chunked", choices=["chunked", "fulltext"], help="Extraction mode: chunked multi-pass or single-call fulltext.", ) parser.add_argument( "--series-memory-json", default=None, help="Optional JSON file containing prior canon memory to inject into the prompt.", ) parser.add_argument( "--sqlite-output", default=None, help="Optional SQLite file to store extracted entities and relation facts.", ) parser.add_argument( "--episode-key", default=None, help="Optional stable episode key for SQLite storage. Defaults to the PDF stem.", ) parser.add_argument( "--chunk-chars", type=int, default=18000, help="Approximate max characters per chunk before multi-pass extraction.", ) parser.add_argument( "--temperature", type=float, default=0.2, help="Sampling temperature for extraction.", ) return parser.parse_args() def extract_pdf_pages(pdf_path: Path) -> List[str]: try: from pypdf import PdfReader except ImportError as exc: # pragma: no cover - runtime dependency guard raise SystemExit( "Missing dependency: pypdf\n" "Install it with:\n" " python3 -m pip install -r screenplay/requirements.txt" ) from exc reader = PdfReader(str(pdf_path)) pages: List[str] = [] for index, page in enumerate(reader.pages, start=1): text = page.extract_text() or "" text = text.strip() if text: pages.append(f"[PAGE {index}]\n{text}") return pages def load_prompt_bundle(version_override: str | None) -> PromptBundle: try: manifest = json.loads(PROMPTS_MANIFEST.read_text(encoding="utf-8")) except FileNotFoundError as exc: raise RuntimeError(f"Prompt manifest not found: {PROMPTS_MANIFEST}") from exc version = version_override or manifest["current_version"] version_entry = manifest["versions"].get(version) if not version_entry: raise RuntimeError(f"Unknown prompt version: {version}") version_dir = PROMPTS_DIR / version_entry["path"] prompt_files: Dict[str, str] = {} prompt_hashes: Dict[str, str] = {} def read_file(name: str) -> str: path = version_dir / name try: content = path.read_text(encoding="utf-8").strip() except FileNotFoundError as exc: raise RuntimeError(f"Prompt file not found: {path}") from exc prompt_files[name] = str(path) prompt_hashes[name] = hashlib.sha256(content.encode("utf-8")).hexdigest() return content return PromptBundle( version=version, chunk_system=read_file("chunk_system.md"), chunk_user=read_file("chunk_user.md"), final_system=read_file("final_system.md"), final_user=read_file("final_user.md"), manifest_path=str(PROMPTS_MANIFEST), version_path=str(version_dir), prompt_files=prompt_files, prompt_hashes=prompt_hashes, ) def build_chunks(page_texts: Sequence[str], max_chars: int) -> List[TextChunk]: chunks: List[TextChunk] = [] buffer: List[str] = [] page_start = 1 current_len = 0 for idx, page_text in enumerate(page_texts, start=1): page_len = len(page_text) if buffer and current_len + page_len + 2 > max_chars: chunks.append(TextChunk(page_start=page_start, page_end=idx - 1, text="\n\n".join(buffer))) buffer = [] page_start = idx current_len = 0 buffer.append(page_text) current_len += page_len + 2 if buffer: chunks.append(TextChunk(page_start=page_start, page_end=len(page_texts), text="\n\n".join(buffer))) return chunks def detect_source_language(page_texts: Sequence[str]) -> LanguageInfo: text = "\n".join(page_texts) counts = {"ko": 0, "ja": 0, "en": 0} for char in text: codepoint = ord(char) if 0xAC00 <= codepoint <= 0xD7A3: counts["ko"] += 1 elif ( 0x3040 <= codepoint <= 0x309F or 0x30A0 <= codepoint <= 0x30FF or 0x31F0 <= codepoint <= 0x31FF or 0xFF66 <= codepoint <= 0xFF9D ): counts["ja"] += 1 elif ("A" <= char <= "Z") or ("a" <= char <= "z"): counts["en"] += 1 if counts["ko"] >= max(counts["ja"], counts["en"]) and counts["ko"] > 0: return LanguageInfo(code="ko", name="Korean") if counts["ja"] >= max(counts["ko"], counts["en"]) and counts["ja"] > 0: return LanguageInfo(code="ja", name="Japanese") return LanguageInfo(code="en", name="English") def resolve_source_language(page_texts: Sequence[str], override: str) -> LanguageInfo: if override != "auto": return LanguageInfo(code=override, name=SUPPORTED_LANGUAGE_MAP[override]) return detect_source_language(page_texts) def resolve_models(provider: str, model: str | None, fallback_model: str | None) -> Tuple[str, str]: if provider == "openai": return model or DEFAULT_OPENAI_MODEL, fallback_model or DEFAULT_OPENAI_FALLBACK_MODEL return model or DEFAULT_GEMINI_MODEL, fallback_model or DEFAULT_GEMINI_FALLBACK_MODEL def load_series_memory(memory_path: str | None) -> Dict[str, object]: if not memory_path: return {} path = Path(memory_path).expanduser().resolve() try: return json.loads(path.read_text(encoding="utf-8")) except FileNotFoundError as exc: raise RuntimeError(f"Series memory JSON not found: {path}") from exc def render_series_memory_block(series_memory: Dict[str, object]) -> str: if not series_memory: return "No prior series canon memory is available for this episode." return ( "Prior series canon memory is available below. " "If a current entity is clearly the same as an earlier canon entity, reuse that established canonical name exactly.\n\n" f"{json.dumps(series_memory, ensure_ascii=False, indent=2)}" ) def response_text_from_payload(payload: Dict[str, object]) -> str: output_text = payload.get("output_text") if isinstance(output_text, str) and output_text.strip(): return output_text output = payload.get("output") if isinstance(output, list): for item in output: if not isinstance(item, dict): continue content = item.get("content") if not isinstance(content, list): continue for part in content: if isinstance(part, dict) and part.get("type") == "output_text": text = part.get("text") if isinstance(text, str) and text.strip(): return text raise RuntimeError("OpenAI response did not include output_text.") def call_openai_structured( *, model: str, instructions: str, user_input: str, schema_name: str, schema: Dict[str, object], temperature: float, max_output_tokens: int | None = None, ) -> Dict[str, object]: api_key = os.environ.get("OPENAI_API_KEY") if not api_key: raise RuntimeError("OPENAI_API_KEY is not set.") body = { "model": model, "instructions": instructions, "input": user_input, "temperature": temperature, "store": False, "text": { "format": { "type": "json_schema", "name": schema_name, "strict": True, "schema": schema, } }, } if max_output_tokens is not None: body["max_output_tokens"] = max_output_tokens request = urllib.request.Request( OPENAI_API_URL, data=json.dumps(body).encode("utf-8"), headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", }, method="POST", ) try: with urllib.request.urlopen(request, timeout=API_TIMEOUT_SECONDS) as response: payload = json.loads(response.read().decode("utf-8")) except urllib.error.HTTPError as exc: error_text = exc.read().decode("utf-8", errors="replace") raise RuntimeError(f"OpenAI API error {exc.code}: {error_text}") from exc except urllib.error.URLError as exc: raise RuntimeError(f"Failed to reach OpenAI API: {exc}") from exc text = response_text_from_payload(payload) try: return json.loads(text) except json.JSONDecodeError as exc: raise RuntimeError(f"Structured output was not valid JSON:\n{text}") from exc def response_text_from_gemini_payload(payload: Dict[str, object]) -> str: candidates = payload.get("candidates") if not isinstance(candidates, list): raise RuntimeError(f"Gemini response did not include candidates: {payload}") texts: List[str] = [] for candidate in candidates: if not isinstance(candidate, dict): continue content = candidate.get("content") if not isinstance(content, dict): continue parts = content.get("parts") if not isinstance(parts, list): continue for part in parts: if isinstance(part, dict): text = part.get("text") if isinstance(text, str) and text.strip(): texts.append(text) if texts: return "\n".join(texts) prompt_feedback = payload.get("promptFeedback") raise RuntimeError(f"Gemini response did not include text parts. promptFeedback={prompt_feedback}") def call_gemini_structured( *, model: str, instructions: str, user_input: str, schema_name: str, schema: Dict[str, object], temperature: float, ) -> Dict[str, object]: api_key = os.environ.get("GEMINI_API_KEY") if not api_key: raise RuntimeError("GEMINI_API_KEY is not set.") body = { "systemInstruction": { "parts": [{"text": instructions}], }, "contents": [ { "role": "user", "parts": [{"text": user_input}], } ], "generationConfig": { "temperature": temperature, "responseMimeType": "application/json", "responseJsonSchema": { "title": schema_name, **schema, }, }, } url = GEMINI_API_URL_TEMPLATE.format(model=model, api_key=api_key) request = urllib.request.Request( url, data=json.dumps(body).encode("utf-8"), headers={"Content-Type": "application/json"}, method="POST", ) try: with urllib.request.urlopen(request, timeout=API_TIMEOUT_SECONDS) as response: payload = json.loads(response.read().decode("utf-8")) except urllib.error.HTTPError as exc: error_text = exc.read().decode("utf-8", errors="replace") raise RuntimeError(f"Gemini API error {exc.code}: {error_text}") from exc except urllib.error.URLError as exc: raise RuntimeError(f"Failed to reach Gemini API: {exc}") from exc text = response_text_from_gemini_payload(payload) try: return json.loads(text) except json.JSONDecodeError as exc: raise RuntimeError(f"Gemini structured output was not valid JSON:\n{text}") from exc def render_prompt(template: str, **values: object) -> str: return template.format(**values) def chunk_prompt( bundle: PromptBundle, chunk: TextChunk, language: LanguageInfo, series_memory_block: str, ) -> str: return render_prompt( bundle.chunk_user, page_start=chunk.page_start, page_end=chunk.page_end, source_language_code=language.code, source_language_name=language.name, series_memory_block=series_memory_block, screenplay_text=chunk.text, ) def consolidation_prompt( bundle: PromptBundle, source_file: str, partial_results: Iterable[Dict[str, object]], language: LanguageInfo, series_memory_block: str, ) -> str: return render_prompt( bundle.final_user, source_file=source_file, source_language_code=language.code, source_language_name=language.name, series_memory_block=series_memory_block, partial_results_json=json.dumps(list(partial_results), ensure_ascii=False, indent=2), ) def is_model_unavailable_error(exc: RuntimeError) -> bool: message = str(exc).lower() return "model" in message and ( "does not exist" in message or "not found" in message or "unsupported" in message or "not available" in message or "invalid model" in message or "not found for api version" in message ) def call_structured( *, provider: str, model: str, instructions: str, user_input: str, schema_name: str, schema: Dict[str, object], temperature: float, ) -> Dict[str, object]: if provider == "openai": return call_openai_structured( model=model, instructions=instructions, user_input=user_input, schema_name=schema_name, schema=schema, temperature=temperature, ) return call_gemini_structured( model=model, instructions=instructions, user_input=user_input, schema_name=schema_name, schema=schema, temperature=temperature, ) def extract_chunk_with_fallback( *, provider: str, primary_model: str, fallback_model: str, instructions: str, user_input: str, schema_name: str, schema: Dict[str, object], temperature: float, ) -> Tuple[Dict[str, object], str]: try: return ( call_structured( provider=provider, model=primary_model, instructions=instructions, user_input=user_input, schema_name=schema_name, schema=schema, temperature=temperature, ), primary_model, ) except RuntimeError as exc: if not fallback_model or fallback_model == primary_model or not is_model_unavailable_error(exc): raise print( textwrap.dedent( f"""\ Primary model failed: {primary_model} Falling back to: {fallback_model} Reason: {exc} """ ).strip(), flush=True, ) return ( call_structured( provider=provider, model=fallback_model, instructions=instructions, user_input=user_input, schema_name=schema_name, schema=schema, temperature=temperature, ), fallback_model, ) def prune_entities(payload: Dict[str, object]) -> Dict[str, object]: def has_reference_value(item: Dict[str, object]) -> bool: return ( item["continuity_priority"] in {"critical", "high"} and item["reference_image_priority"] == "required" and len(item["visual_anchor_traits"]) >= 2 ) payload["characters"] = [ item for item in payload["characters"] if item["importance"] in {"major", "supporting"} and has_reference_value(item) ] payload["locations"] = [ item for item in payload["locations"] if item["importance"] in {"major", "supporting"} and has_reference_value(item) ] payload["props"] = [ item for item in payload["props"] if item["significance"] in {"key", "recurring"} and has_reference_value(item) ] entity_lookup = build_entity_lookup(payload) payload["relation_facts"] = merge_relation_facts( payload.get("relation_facts", []), entity_lookup, ) return payload def normalize_name(name: str) -> str: return " ".join(name.lower().split()) def merge_text(preferred: str, candidate: str) -> str: return candidate if len(candidate) > len(preferred) else preferred def entity_identifiers(item: Dict[str, object]) -> set[str]: identifiers = {normalize_name(str(item["name"]))} identifiers.update(normalize_name(alias) for alias in item["aliases"] if isinstance(alias, str) and alias.strip()) return identifiers def choose_canonical_name( current_name: str, candidate_name: str, current_aliases: Sequence[str], candidate_aliases: Sequence[str], ) -> str: current_norm = normalize_name(current_name) candidate_norm = normalize_name(candidate_name) current_alias_norms = {normalize_name(alias) for alias in current_aliases} candidate_alias_norms = {normalize_name(alias) for alias in candidate_aliases} if current_norm in candidate_alias_norms and len(candidate_name) > len(current_name): return candidate_name if candidate_norm in current_alias_norms and len(current_name) >= len(candidate_name): return current_name return candidate_name if len(candidate_name) > len(current_name) else current_name def merge_entity_lists(items: Sequence[Dict[str, object]], kind: str) -> List[Dict[str, object]]: merged: Dict[str, Dict[str, object]] = {} for item in items: match_key = None for existing_key, existing_item in merged.items(): if entity_identifiers(item) & entity_identifiers(existing_item): match_key = existing_key break key = match_key or normalize_name(str(item["name"])) if key not in merged: merged[key] = { **item, "aliases": list(dict.fromkeys(item["aliases"])), "visual_anchor_traits": list(dict.fromkeys(item["visual_anchor_traits"])), "variant_axes": list(dict.fromkeys(item["variant_axes"])), "evidence": list(dict.fromkeys(item["evidence"])), } continue current = merged[key] current["name"] = choose_canonical_name( str(current["name"]), str(item["name"]), current["aliases"], item["aliases"], ) current["description"] = merge_text(str(current["description"]), str(item["description"])) current["continuity_reason"] = merge_text( str(current["continuity_reason"]), str(item["continuity_reason"]), ) current["aliases"] = list(dict.fromkeys([*current["aliases"], *item["aliases"]])) current["visual_anchor_traits"] = list( dict.fromkeys([*current["visual_anchor_traits"], *item["visual_anchor_traits"]]) ) current["variant_axes"] = list(dict.fromkeys([*current["variant_axes"], *item["variant_axes"]])) current["evidence"] = list(dict.fromkeys([*current["evidence"], *item["evidence"]])) if kind in {"characters", "locations"}: rank = {"major": 3, "supporting": 2, "minor": 1, "unknown": 0} if rank[item["importance"]] > rank[current["importance"]]: current["importance"] = item["importance"] else: rank = {"key": 3, "recurring": 2, "minor": 1, "unknown": 0} if rank[item["significance"]] > rank[current["significance"]]: current["significance"] = item["significance"] continuity_rank = {"critical": 3, "high": 2, "medium": 1, "low": 0} if continuity_rank[item["continuity_priority"]] > continuity_rank[current["continuity_priority"]]: current["continuity_priority"] = item["continuity_priority"] reference_rank = {"required": 2, "helpful": 1, "not_needed": 0} if reference_rank[item["reference_image_priority"]] > reference_rank[current["reference_image_priority"]]: current["reference_image_priority"] = item["reference_image_priority"] for current in merged.values(): current["aliases"] = [ alias for alias in current["aliases"] if normalize_name(alias) != normalize_name(str(current["name"])) ] return list(merged.values()) def build_entity_lookup(payload: Dict[str, object]) -> Dict[Tuple[str, str], str]: lookup: Dict[Tuple[str, str], str] = {} for section, entity_type in ( ("characters", "character"), ("locations", "location"), ("props", "prop"), ): for item in payload[section]: names = [str(item["name"]), *[alias for alias in item["aliases"] if isinstance(alias, str)]] for name in names: normalized = normalize_name(name) if normalized: lookup[(entity_type, normalized)] = str(item["name"]) return lookup def normalize_relation_fact( relation: Dict[str, object], entity_lookup: Dict[Tuple[str, str], str], ) -> Dict[str, object] | None: participants: List[Dict[str, str]] = [] seen = set() for participant in relation["participants"]: entity_type = str(participant["entity_type"]) entity_name = str(participant["entity_name"]) role = str(participant["role"]).strip() or "participant" canonical_name = entity_lookup.get((entity_type, normalize_name(entity_name))) if not canonical_name: continue key = (entity_type, canonical_name, role) if key in seen: continue seen.add(key) participants.append( { "entity_type": entity_type, "entity_name": canonical_name, "role": role, } ) unique_entities = {(item["entity_type"], item["entity_name"]) for item in participants} if len(unique_entities) < 2: return None evidence = list(dict.fromkeys(str(item) for item in relation["evidence"] if str(item).strip())) if not evidence: return None return { "relation_family": str(relation["relation_family"]), "relation_type": str(relation["relation_type"]).strip(), "directionality": str(relation["directionality"]), "temporal_scope": str(relation["temporal_scope"]), "continuity_priority": str(relation["continuity_priority"]), "continuity_reason": str(relation["continuity_reason"]).strip(), "participants": participants, "evidence": evidence, } def merge_relation_facts( relation_facts: Sequence[Dict[str, object]], entity_lookup: Dict[Tuple[str, str], str], ) -> List[Dict[str, object]]: merged: Dict[str, Dict[str, object]] = {} continuity_rank = {"critical": 3, "high": 2, "medium": 1, "low": 0} for relation in relation_facts: normalized = normalize_relation_fact(relation, entity_lookup) if not normalized: continue participant_signature = tuple( sorted( f"{item['entity_type']}|{item['entity_name']}|{item['role']}" for item in normalized["participants"] ) ) signature = "||".join( [ normalized["relation_family"], normalized["relation_type"], normalized["directionality"], normalized["temporal_scope"], *participant_signature, ] ) if signature not in merged: merged[signature] = normalized continue current = merged[signature] current["continuity_reason"] = merge_text( str(current["continuity_reason"]), str(normalized["continuity_reason"]), ) current["evidence"] = list(dict.fromkeys([*current["evidence"], *normalized["evidence"]])) if continuity_rank[normalized["continuity_priority"]] > continuity_rank[current["continuity_priority"]]: current["continuity_priority"] = normalized["continuity_priority"] return list(merged.values()) def stable_id(prefix: str, *parts: str) -> str: digest = hashlib.sha1("||".join(parts).encode("utf-8")).hexdigest()[:16] return f"{prefix}_{digest}" def init_extraction_graph_sqlite(db_path: Path) -> None: db_path.parent.mkdir(parents=True, exist_ok=True) with sqlite3.connect(db_path) as conn: conn.executescript( """ PRAGMA foreign_keys = ON; CREATE TABLE IF NOT EXISTS extraction_run ( id TEXT PRIMARY KEY, source_file TEXT NOT NULL, episode_key TEXT NOT NULL, provider TEXT NOT NULL, model TEXT NOT NULL, mode TEXT NOT NULL, prompt_version TEXT NOT NULL, source_language_code TEXT NOT NULL, source_language_name TEXT NOT NULL, payload_path TEXT, created_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS extracted_entity ( id TEXT PRIMARY KEY, run_id TEXT NOT NULL REFERENCES extraction_run(id) ON DELETE CASCADE, entity_type TEXT NOT NULL CHECK (entity_type IN ('character', 'location', 'prop')), name TEXT NOT NULL, aliases_json TEXT NOT NULL DEFAULT '[]', description TEXT NOT NULL, continuity_reason TEXT NOT NULL, visual_anchor_traits_json TEXT NOT NULL DEFAULT '[]', variant_axes_json TEXT NOT NULL DEFAULT '[]', continuity_priority TEXT NOT NULL, reference_image_priority TEXT NOT NULL, kind TEXT, significance TEXT, importance TEXT, evidence_json TEXT NOT NULL DEFAULT '[]' ); CREATE INDEX IF NOT EXISTS idx_extracted_entity_run_type ON extracted_entity(run_id, entity_type, name); CREATE TABLE IF NOT EXISTS relation_fact ( id TEXT PRIMARY KEY, run_id TEXT NOT NULL REFERENCES extraction_run(id) ON DELETE CASCADE, relation_family TEXT NOT NULL, relation_type TEXT NOT NULL, directionality TEXT NOT NULL, temporal_scope TEXT NOT NULL, continuity_priority TEXT NOT NULL, continuity_reason TEXT NOT NULL, participant_count INTEGER NOT NULL, signature TEXT NOT NULL, evidence_json TEXT NOT NULL DEFAULT '[]' ); CREATE UNIQUE INDEX IF NOT EXISTS idx_relation_fact_signature ON relation_fact(run_id, signature); CREATE TABLE IF NOT EXISTS relation_participant ( id TEXT PRIMARY KEY, relation_id TEXT NOT NULL REFERENCES relation_fact(id) ON DELETE CASCADE, entity_id TEXT REFERENCES extracted_entity(id) ON DELETE SET NULL, entity_type TEXT NOT NULL CHECK (entity_type IN ('character', 'location', 'prop')), entity_name TEXT NOT NULL, role TEXT NOT NULL, participant_index INTEGER NOT NULL ); CREATE INDEX IF NOT EXISTS idx_relation_participant_relation ON relation_participant(relation_id, participant_index); """ ) def store_payload_in_sqlite( *, db_path: Path, input_path: Path, payload_path: Path | None, payload: Dict[str, object], episode_key: str, ) -> None: init_extraction_graph_sqlite(db_path) metadata = payload["extraction_metadata"] language = metadata["source_language"] run_id = stable_id( "run", str(input_path), metadata["provider"], metadata["requested_model"], metadata["mode"], metadata["prompt_version"], ) created_at = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) with sqlite3.connect(db_path) as conn: conn.execute("PRAGMA foreign_keys = ON") conn.execute("DELETE FROM extraction_run WHERE id = ?", (run_id,)) conn.execute( """ INSERT INTO extraction_run ( id, source_file, episode_key, provider, model, mode, prompt_version, source_language_code, source_language_name, payload_path, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, str(input_path), episode_key, metadata["provider"], metadata["requested_model"], metadata["mode"], metadata["prompt_version"], language["code"], language["name"], str(payload_path) if payload_path else None, created_at, ), ) entity_ids: Dict[Tuple[str, str], str] = {} for section, entity_type in ( ("characters", "character"), ("locations", "location"), ("props", "prop"), ): for item in payload[section]: entity_id = stable_id("entity", run_id, entity_type, str(item["name"])) entity_ids[(entity_type, str(item["name"]))] = entity_id conn.execute( """ INSERT INTO extracted_entity ( id, run_id, entity_type, name, aliases_json, description, continuity_reason, visual_anchor_traits_json, variant_axes_json, continuity_priority, reference_image_priority, kind, significance, importance, evidence_json ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( entity_id, run_id, entity_type, item["name"], json.dumps(item["aliases"], ensure_ascii=False), item["description"], item["continuity_reason"], json.dumps(item["visual_anchor_traits"], ensure_ascii=False), json.dumps(item["variant_axes"], ensure_ascii=False), item["continuity_priority"], item["reference_image_priority"], item.get("kind"), item.get("significance"), item.get("importance"), json.dumps(item["evidence"], ensure_ascii=False), ), ) for relation in payload.get("relation_facts", []): participant_signature = [ f"{item['entity_type']}|{item['entity_name']}|{item['role']}" for item in relation["participants"] ] signature = "||".join( [ relation["relation_family"], relation["relation_type"], relation["directionality"], relation["temporal_scope"], *sorted(participant_signature), ] ) relation_id = stable_id("rel", run_id, signature) conn.execute( """ INSERT INTO relation_fact ( id, run_id, relation_family, relation_type, directionality, temporal_scope, continuity_priority, continuity_reason, participant_count, signature, evidence_json ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( relation_id, run_id, relation["relation_family"], relation["relation_type"], relation["directionality"], relation["temporal_scope"], relation["continuity_priority"], relation["continuity_reason"], len(relation["participants"]), signature, json.dumps(relation["evidence"], ensure_ascii=False), ), ) for index, participant in enumerate(relation["participants"], start=1): participant_id = stable_id( "part", relation_id, str(index), participant["entity_type"], participant["entity_name"], participant["role"], ) conn.execute( """ INSERT INTO relation_participant ( id, relation_id, entity_id, entity_type, entity_name, role, participant_index ) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( participant_id, relation_id, entity_ids.get((participant["entity_type"], participant["entity_name"])), participant["entity_type"], participant["entity_name"], participant["role"], index, ), ) def local_consolidate(source_file: str, partials: Sequence[Dict[str, object]]) -> Dict[str, object]: characters = merge_entity_lists( [item for partial in partials for item in partial["characters"]], "characters", ) locations = merge_entity_lists( [item for partial in partials for item in partial["locations"]], "locations", ) props = merge_entity_lists( [item for partial in partials for item in partial["props"]], "props", ) entity_lookup = build_entity_lookup( { "characters": characters, "locations": locations, "props": props, } ) relation_facts = merge_relation_facts( [item for partial in partials for item in partial.get("relation_facts", [])], entity_lookup, ) notes: List[str] = [] for partial in partials: for note in partial["notes"]: if note not in notes: notes.append(note) summary = ( f"Continuity-focused extraction from {source_file}: " f"{len(characters)} characters, {len(locations)} locations, {len(props)} props, " f"{len(relation_facts)} relation facts." ) return { "source_file": source_file, "characters": characters, "locations": locations, "props": props, "relation_facts": relation_facts, "summary": summary, "notes": notes, } def localize_summary( *, source_file: str, character_count: int, location_count: int, prop_count: int, relation_count: int, language: LanguageInfo, ) -> str: if language.code == "ko": return ( f"{source_file} 연속성 앵커 추출 결과: " f"인물 {character_count}개, 배경/장소 {location_count}개, 중요 소품 {prop_count}개, 관계 {relation_count}개." ) if language.code == "ja": return ( f"{source_file} の継続性アンカー抽出結果: " f"人物 {character_count}件、背景・場所 {location_count}件、重要小道具 {prop_count}件、関係 {relation_count}件。" ) return ( f"Continuity-anchor extraction from {source_file}: " f"{character_count} characters, {location_count} locations, {prop_count} key props, {relation_count} relation facts." ) def extract_entities_for_pdf( *, input_path: Path, output_path: Path | None, sqlite_output_path: Path | None, episode_key: str, provider: str, model: str, fallback_model: str, prompt_version: str | None, source_language_override: str, mode: str, series_memory: Dict[str, object], chunk_chars: int, temperature: float, ) -> Dict[str, object]: page_texts = extract_pdf_pages(input_path) if not page_texts: raise RuntimeError("No extractable text found in PDF.") source_language = resolve_source_language(page_texts, source_language_override) if mode == "fulltext": chunks = [ TextChunk( page_start=1, page_end=len(page_texts), text="\n\n".join(page_texts), ) ] else: chunks = build_chunks(page_texts, max_chars=chunk_chars) partials: List[Dict[str, object]] = [] chunk_models_used: List[str] = [] bundle = load_prompt_bundle(prompt_version) series_memory_block = render_series_memory_block(series_memory) instructions = render_prompt( bundle.chunk_system, source_language_code=source_language.code, source_language_name=source_language.name, ) print(f"Extracting entities from {input_path.name}", flush=True) print(f"Provider: {provider}", flush=True) print(f"Model: {model}", flush=True) print(f"Mode: {mode}", flush=True) print(f"Pages with text: {len(page_texts)}", flush=True) print(f"Chunks: {len(chunks)}", flush=True) print(f"Detected language: {source_language.name} ({source_language.code})", flush=True) print(f"Prompt version: {bundle.version}", flush=True) started_at = time.time() for idx, chunk in enumerate(chunks, start=1): print(f" - chunk {idx}/{len(chunks)} pages {chunk.page_start}-{chunk.page_end}", flush=True) partial, model_used = extract_chunk_with_fallback( provider=provider, primary_model=model, fallback_model=fallback_model, instructions=instructions, user_input=chunk_prompt(bundle, chunk, source_language, series_memory_block), schema_name="screenplay_entities_chunk", schema=CHUNK_SCHEMA, temperature=temperature, ) partials.append(partial) chunk_models_used.append(model_used) final_model_used = None final_payload = local_consolidate(input_path.name, partials) final_payload = prune_entities(final_payload) final_payload["summary"] = localize_summary( source_file=input_path.name, character_count=len(final_payload["characters"]), location_count=len(final_payload["locations"]), prop_count=len(final_payload["props"]), relation_count=len(final_payload["relation_facts"]), language=source_language, ) final_payload["extraction_metadata"] = { "provider": provider, "mode": mode, "source_language": { "code": source_language.code, "name": source_language.name, }, "series_memory_entity_counts": { "characters": len(series_memory.get("characters", [])) if isinstance(series_memory, dict) else 0, "locations": len(series_memory.get("locations", [])) if isinstance(series_memory, dict) else 0, "props": len(series_memory.get("props", [])) if isinstance(series_memory, dict) else 0, "relation_facts": len(series_memory.get("relation_facts", [])) if isinstance(series_memory, dict) else 0, }, "prompt_version": bundle.version, "prompt_manifest": bundle.manifest_path, "prompt_version_path": bundle.version_path, "prompt_files": bundle.prompt_files, "prompt_hashes": bundle.prompt_hashes, "requested_model": model, "fallback_model": fallback_model, "chunk_models_used": chunk_models_used, "final_model_used": final_model_used, "elapsed_seconds": round(time.time() - started_at, 2), } if output_path is not None: output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text( json.dumps(final_payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) print(f"Saved JSON to: {output_path}", flush=True) if sqlite_output_path is not None: store_payload_in_sqlite( db_path=sqlite_output_path, input_path=input_path, payload_path=output_path, payload=final_payload, episode_key=episode_key, ) print(f"Saved SQLite graph to: {sqlite_output_path}", flush=True) return final_payload def main() -> int: args = parse_args() input_path = Path(args.input).expanduser().resolve() output_path = Path(args.output).expanduser().resolve() if not input_path.exists(): print(f"Input PDF not found: {input_path}", file=sys.stderr) return 1 model, fallback_model = resolve_models(args.provider, args.model, args.fallback_model) series_memory = load_series_memory(args.series_memory_json) sqlite_output_path = ( Path(args.sqlite_output).expanduser().resolve() if args.sqlite_output else None ) episode_key = args.episode_key or input_path.stem try: extract_entities_for_pdf( input_path=input_path, output_path=output_path, sqlite_output_path=sqlite_output_path, episode_key=episode_key, provider=args.provider, model=model, fallback_model=fallback_model, prompt_version=args.prompt_version, source_language_override=args.source_language, mode=args.mode, series_memory=series_memory, chunk_chars=args.chunk_chars, temperature=args.temperature, ) except RuntimeError as exc: print(str(exc), file=sys.stderr) return 1 return 0 if __name__ == "__main__": raise SystemExit(main())