From f870198922ec9ac29e9c030caa2c3c93e71df84e Mon Sep 17 00:00:00 2001 From: WillLewis Date: Thu, 2 Jul 2026 23:42:34 -0400 Subject: [PATCH] Add P0 extraction workstream package Serves PLAN Phase 0 R2/R3: structured extraction, entity resolution, and confidence gate JSON outputs. --- extract/README.md | 90 ++ extract/__init__.py | 9 + extract/__main__.py | 3 + extract/cli.py | 310 +++++++ extract/examples/greyharbor_aliases.json | 109 +++ .../greyharbor_ep101_ep102_scenes.json | 104 +++ extract/examples/marcus_aliases.json | 14 + extract/examples/marcus_candidates.json | 58 ++ extract/extraction.py | 185 ++++ extract/gate.py | 151 ++++ extract/llm.py | 66 ++ extract/models.py | 114 +++ extract/prompts.py | 63 ++ extract/prompts/entity_resolution_system.md | 13 + extract/prompts/extraction_system.md | 57 ++ extract/resolution.py | 789 ++++++++++++++++++ extract/schema.py | 122 +++ 17 files changed, 2257 insertions(+) create mode 100644 extract/README.md create mode 100644 extract/__init__.py create mode 100644 extract/__main__.py create mode 100644 extract/cli.py create mode 100644 extract/examples/greyharbor_aliases.json create mode 100644 extract/examples/greyharbor_ep101_ep102_scenes.json create mode 100644 extract/examples/marcus_aliases.json create mode 100644 extract/examples/marcus_candidates.json create mode 100644 extract/extraction.py create mode 100644 extract/gate.py create mode 100644 extract/llm.py create mode 100644 extract/models.py create mode 100644 extract/prompts.py create mode 100644 extract/prompts/entity_resolution_system.md create mode 100644 extract/prompts/extraction_system.md create mode 100644 extract/resolution.py create mode 100644 extract/schema.py diff --git a/extract/README.md b/extract/README.md new file mode 100644 index 0000000..1e3965c --- /dev/null +++ b/extract/README.md @@ -0,0 +1,90 @@ +# P0-EXTRACT + +Standalone extraction workstream for Canon AI Phase 0. + +This directory is JSON-in / JSON-out. It does not parse scripts, write to +Postgres, create migrations, or build UI. + +## Input Scene JSON + +Accepted shapes: + +```json +{ + "world": "greyharbor", + "scenes": [ + { + "scene_id": "E101/sc1", + "work_title": "E101 - The Ledger", + "scene_index": 1, + "slug": "INT. HARBORMASTER'S OFFICE - DAY", + "story_position": 1, + "is_flashback": false, + "raw_text": "INT. HARBORMASTER'S OFFICE - DAY\n\n..." + } + ] +} +``` + +`slug`, `story_position`, `is_flashback`, and `raw_text` are required. `scene_id` +is strongly preferred because it is the citation handle before database IDs +exist. + +## Commands + +Dry-run prompts, with no API call: + +```sh +python -m extract extract extract/examples/greyharbor_ep101_ep102_scenes.json --dry-run --limit 1 +``` + +Run LLM extraction: + +```sh +python -m extract extract extract/examples/greyharbor_ep101_ep102_scenes.json \ + --world greyharbor \ + --out out/candidates.json +``` + +Resolve entities from candidate assertions: + +```sh +python -m extract resolve out/candidates.json \ + --aliases extract/examples/greyharbor_aliases.json \ + --state out/extract-state.json \ + --out out/resolved.json +``` + +Apply the confidence gate: + +```sh +python -m extract gate \ + --state out/extract-state.json \ + --out out/resolved-gated.json \ + --eval-out out/assertions.json +``` + +Review human queues: + +```sh +python -m extract confirm-entities --state out/extract-state.json --out out/resolved.json +python -m extract confirm-assertions --state out/extract-state.json --out out/resolved.json +``` + +Human-only merge: + +```sh +python -m extract merge --state out/extract-state.json --keep "Marcus Hale" --drop "MARCUS" --reason "same cast-listed character" +``` + +## Status Rules + +- candidate extraction drops any assertion without a verbatim scene quote. +- predicates are a closed set from `docs/extraction.md`. +- `confidence >= 0.85` becomes `draft`. +- lower-confidence assertions become `needs_confirmation` queue items. +- accepted/edited queue items become `canon`; rejected items are excluded from + eval output. +- entity disambiguation may assign one referring expression to one existing + entity, but it never merges two existing entities. Merges only happen through + the human CLI. diff --git a/extract/__init__.py b/extract/__init__.py new file mode 100644 index 0000000..120a880 --- /dev/null +++ b/extract/__init__.py @@ -0,0 +1,9 @@ +"""P0-EXTRACT workstream package. + +This package is intentionally JSON-in / JSON-out. It does not parse scripts and +does not write to Postgres; ingestion and storage remain separate workstreams. +""" + +from .schema import AUTO_ACCEPT_CONFIDENCE, PREDICATES + +__all__ = ["AUTO_ACCEPT_CONFIDENCE", "PREDICATES"] diff --git a/extract/__main__.py b/extract/__main__.py new file mode 100644 index 0000000..eb53e2f --- /dev/null +++ b/extract/__main__.py @@ -0,0 +1,3 @@ +from .cli import main + +raise SystemExit(main()) diff --git a/extract/cli.py b/extract/cli.py new file mode 100644 index 0000000..47c46ea --- /dev/null +++ b/extract/cli.py @@ -0,0 +1,310 @@ +"""Command line interface for the P0-EXTRACT workstream.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +from . import extraction as extraction_mod +from . import gate as gate_mod +from . import resolution as resolution_mod +from .llm import DEFAULT_EFFORT, DEFAULT_MODEL, has_credentials, make_client +from .models import load_scenes +from .resolution import state_from_dict, state_to_dict +from .schema import AUTO_ACCEPT_CONFIDENCE + + +def _load_json(path: str | Path) -> dict[str, Any]: + return json.loads(Path(path).read_text(encoding="utf-8")) + + +def _write_json(path: str | Path, payload: dict[str, Any]) -> None: + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8") + + +def cmd_extract(args: argparse.Namespace) -> int: + try: + input_world, scenes = load_scenes(args.scenes) + except (OSError, ValueError, json.JSONDecodeError) as exc: + print(f"error: could not load scenes: {exc}", file=sys.stderr) + return 2 + + world = args.world or input_world + if args.dry_run or not has_credentials(): + if not has_credentials() and not args.dry_run: + print( + "note: no Anthropic credentials set; showing prompts instead of calling the API.\n", + file=sys.stderr, + ) + print(extraction_mod.render_dry_run(scenes, model=args.model, effort=args.effort, limit=args.limit)) + return 0 + + try: + client = make_client() + except RuntimeError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + def on_result(result) -> None: + drops = sum(result.dropped.values()) + suffix = f" ({drops} dropped)" if drops else "" + print(f" {result.scene_id}: {len(result.assertions)} kept{suffix}", file=sys.stderr) + + try: + results = extraction_mod.run_extraction( + client, + scenes, + model=args.model, + effort=args.effort, + thinking=not args.no_thinking, + verify_quotes=not args.no_verify_quotes, + limit=args.limit, + on_result=on_result, + ) + except (RuntimeError, json.JSONDecodeError) as exc: + print(f"error: extraction failed: {exc}", file=sys.stderr) + return 1 + + payload = extraction_mod.to_candidate_json(world, args.model, results) + if args.out: + _write_json(args.out, payload) + total = sum(len(scene.assertions) for scene in results) + print(f"wrote {args.out}: {len(results)} scene(s), {total} candidate assertion(s).") + else: + print(json.dumps(payload, indent=2, ensure_ascii=False)) + return 0 + + +def cmd_resolve(args: argparse.Namespace) -> int: + try: + candidates = _load_json(args.candidates) + alias_table = _load_json(args.aliases) if args.aliases else None + except (OSError, json.JSONDecodeError) as exc: + print(f"error: could not load input JSON: {exc}", file=sys.stderr) + return 2 + + client = None + if not args.no_llm: + if has_credentials(): + try: + client = make_client() + except RuntimeError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + else: + print( + "note: no Anthropic credentials; using exact/fuzzy resolution and queueing the rest.\n", + file=sys.stderr, + ) + + try: + state = resolution_mod.resolve_candidates( + candidates, + world=args.world or "", + alias_table=alias_table, + client=client, + model=args.model, + effort=args.effort, + thinking=not args.no_thinking, + conf_auto=args.conf_auto, + ) + except RuntimeError as exc: + print(f"error: resolution failed: {exc}", file=sys.stderr) + return 1 + + if args.state: + _write_json(args.state, state_to_dict(state)) + print(f"wrote {args.state}") + if args.out: + _write_json(args.out, resolution_mod.to_resolved_assertions(state)) + print(f"wrote {args.out}: {len(state.assertions)} resolved assertion(s)") + if not args.state and not args.out: + print(resolution_mod.render_summary(state)) + if state.entity_queue: + print(f"{len(state.entity_queue)} entity item(s) need confirmation.", file=sys.stderr) + return 0 + + +def _write_state_and_outputs(state, args) -> None: + _write_json(args.state, state_to_dict(state)) + print(f"wrote {args.state}") + if getattr(args, "out", None): + _write_json(args.out, resolution_mod.to_resolved_assertions(state)) + print(f"wrote {args.out}: {len(state.assertions)} assertion(s)") + if getattr(args, "eval_out", None): + _write_json(args.eval_out, gate_mod.to_eval_assertions(state)) + print(f"wrote {args.eval_out}: eval assertion contract") + + +def cmd_gate(args: argparse.Namespace) -> int: + try: + state = state_from_dict(_load_json(args.state)) + except (OSError, json.JSONDecodeError, TypeError, KeyError) as exc: + print(f"error: could not load state: {exc}", file=sys.stderr) + return 2 + gate_mod.gate_state(state, threshold=args.threshold) + _write_state_and_outputs(state, args) + if state.assertion_queue: + print(f"{len(state.assertion_queue)} assertion(s) need confirmation.", file=sys.stderr) + return 0 + + +def _prompt_entity(item, state) -> str: + print("\n" + "-" * 70) + print(f'unresolved entity: "{item.surface}" [{item.kind_hint}] ({item.reason})') + for occurrence in item.occurrences: + pred = f", {occurrence['predicate']}" if occurrence.get("predicate") else "" + print(f" {occurrence['scene']} ({occurrence['role']}{pred})") + if occurrence.get("quote"): + print(f" quote: {occurrence['quote']}") + if item.candidates: + for idx, candidate in enumerate(item.candidates, start=1): + print(f" [{idx}] {candidate}") + known = [entity.name for entity in state.registry.entities if not entity.provisional] + if known: + print(f" existing: {', '.join(known)}") + print(" actions: pick | = assign | new | | keep | skip") + try: + return input("> ").strip() + except EOFError: + return "skip" + + +def cmd_confirm_entities(args: argparse.Namespace) -> int: + try: + state = state_from_dict(_load_json(args.state)) + except (OSError, json.JSONDecodeError, TypeError, KeyError) as exc: + print(f"error: could not load state: {exc}", file=sys.stderr) + return 2 + if not state.entity_queue: + print("entity queue is empty.") + return 0 + resolved = resolution_mod.walk_entity_queue(state, _prompt_entity) + _write_state_and_outputs(state, args) + print(f"resolved {resolved}; {len(state.entity_queue)} entity item(s) still queued.") + return 0 + + +def _prompt_assertion(item, _state) -> str: + print("\n" + "-" * 70) + print(f"low-confidence assertion #{item['assertion_index']} ({item['confidence']:.2f})") + print(f" scene: {item.get('scene') or item.get('scene_id')}") + print(f" assertion: {item['summary']}") + print(f" quote: {item.get('supporting_quote')}") + print(' actions: accept | reject | edit {"field": "value"} | skip') + try: + return input("> ").strip() + except EOFError: + return "skip" + + +def cmd_confirm_assertions(args: argparse.Namespace) -> int: + try: + state = state_from_dict(_load_json(args.state)) + except (OSError, json.JSONDecodeError, TypeError, KeyError) as exc: + print(f"error: could not load state: {exc}", file=sys.stderr) + return 2 + if not state.assertion_queue: + print("assertion queue is empty.") + return 0 + resolved = gate_mod.walk_assertion_queue(state, _prompt_assertion) + _write_state_and_outputs(state, args) + print(f"resolved {resolved}; {len(state.assertion_queue)} assertion(s) still queued.") + return 0 + + +def cmd_merge(args: argparse.Namespace) -> int: + try: + state = state_from_dict(_load_json(args.state)) + except (OSError, json.JSONDecodeError, TypeError, KeyError) as exc: + print(f"error: could not load state: {exc}", file=sys.stderr) + return 2 + if not resolution_mod.merge_entities(state, args.keep, args.drop, args.reason): + print(f"error: could not merge keep={args.keep!r} drop={args.drop!r}", file=sys.stderr) + return 1 + _write_state_and_outputs(state, args) + print(f"merged {args.drop!r} into {args.keep!r}.") + return 0 + + +def cmd_summary(args: argparse.Namespace) -> int: + try: + state = state_from_dict(_load_json(args.state)) + except (OSError, json.JSONDecodeError, TypeError, KeyError) as exc: + print(f"error: could not load state: {exc}", file=sys.stderr) + return 2 + print(resolution_mod.render_summary(state)) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="python -m extract", description="Canon AI P0 extraction workstream") + sub = parser.add_subparsers(dest="command", required=True) + + ext = sub.add_parser("extract", help="scene JSON -> candidate assertion JSON via Anthropic structured outputs") + ext.add_argument("scenes", help="scene records JSON") + ext.add_argument("--world", default=None, help="world name (default: from scene JSON)") + ext.add_argument("--model", default=DEFAULT_MODEL) + ext.add_argument("--effort", default=DEFAULT_EFFORT, choices=["low", "medium", "high", "xhigh", "max"]) + ext.add_argument("--no-thinking", action="store_true") + ext.add_argument("--no-verify-quotes", action="store_true") + ext.add_argument("--limit", type=int, default=None) + ext.add_argument("--out", default=None) + ext.add_argument("--dry-run", action="store_true") + ext.set_defaults(func=cmd_extract) + + res = sub.add_parser("resolve", help="candidate assertion JSON -> resolved entity/assertion state") + res.add_argument("candidates", help="candidate JSON from extract") + res.add_argument("--aliases", default=None, help="optional seeded entity/alias table JSON") + res.add_argument("--world", default=None) + res.add_argument("--model", default=DEFAULT_MODEL) + res.add_argument("--effort", default=DEFAULT_EFFORT, choices=["low", "medium", "high", "xhigh", "max"]) + res.add_argument("--no-thinking", action="store_true") + res.add_argument("--no-llm", action="store_true", help="exact/fuzzy only; queue ambiguous role refs") + res.add_argument("--conf-auto", type=float, default=AUTO_ACCEPT_CONFIDENCE) + res.add_argument("--state", default=None, help="write full resolution state JSON") + res.add_argument("--out", default=None, help="write resolved assertions JSON") + res.set_defaults(func=cmd_resolve) + + gate = sub.add_parser("gate", help="apply confidence gate to a resolution state") + gate.add_argument("--state", required=True) + gate.add_argument("--threshold", type=float, default=AUTO_ACCEPT_CONFIDENCE) + gate.add_argument("--out", default=None, help="write resolved assertions JSON") + gate.add_argument("--eval-out", default=None, help="write eval/run_eval.py assertion contract") + gate.set_defaults(func=cmd_gate) + + ce = sub.add_parser("confirm-entities", help="walk entity confirmation queue") + ce.add_argument("--state", required=True) + ce.add_argument("--out", default=None) + ce.add_argument("--eval-out", default=None) + ce.set_defaults(func=cmd_confirm_entities) + + ca = sub.add_parser("confirm-assertions", help="walk low-confidence assertion queue") + ca.add_argument("--state", required=True) + ca.add_argument("--out", default=None) + ca.add_argument("--eval-out", default=None) + ca.set_defaults(func=cmd_confirm_assertions) + + merge = sub.add_parser("merge", help="human-only entity merge") + merge.add_argument("--state", required=True) + merge.add_argument("--keep", required=True) + merge.add_argument("--drop", required=True) + merge.add_argument("--reason", default="manual merge") + merge.add_argument("--out", default=None) + merge.add_argument("--eval-out", default=None) + merge.set_defaults(func=cmd_merge) + + summary = sub.add_parser("summary", help="summarize a resolution state") + summary.add_argument("--state", required=True) + summary.set_defaults(func=cmd_summary) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + return args.func(args) diff --git a/extract/examples/greyharbor_aliases.json b/extract/examples/greyharbor_aliases.json new file mode 100644 index 0000000..f039e89 --- /dev/null +++ b/extract/examples/greyharbor_aliases.json @@ -0,0 +1,109 @@ +{ + "entities": [ + { + "name": "Mara Voss", + "kind": "character", + "dossier": "Mara Voss, Tobias's niece, searches for her brother Danny and reads the ledger.", + "aliases": [ + {"alias": "MARA", "kind": "name_variant"}, + {"alias": "Mara", "kind": "name_variant"} + ] + }, + { + "name": "Tobias Voss", + "kind": "character", + "dossier": "Tobias Voss is Mara's uncle and Greyharbor's harbormaster.", + "aliases": [ + {"alias": "TOBIAS", "kind": "name_variant"}, + {"alias": "Tobias", "kind": "name_variant"}, + {"alias": "your uncle", "kind": "role_reference"}, + {"alias": "uncle", "kind": "role_reference"}, + {"alias": "harbormaster", "kind": "role_reference"} + ] + }, + { + "name": "Edda Quinn", + "kind": "character", + "dossier": "Edda Quinn is the harbor cook and warns Mara about Tobias's numbers.", + "aliases": [ + {"alias": "EDDA", "kind": "name_variant"}, + {"alias": "Edda", "kind": "name_variant"} + ] + }, + { + "name": "Cole Brannigan", + "kind": "character", + "dossier": "Deputy Cole Brannigan is a deputy in Greyharbor.", + "aliases": [ + {"alias": "COLE", "kind": "name_variant"}, + {"alias": "Cole", "kind": "name_variant"}, + {"alias": "DEPUTY COLE BRANNIGAN", "kind": "name_variant"}, + {"alias": "the deputy", "kind": "role_reference"}, + {"alias": "deputy", "kind": "role_reference"}, + {"alias": "C.B.", "kind": "nickname"} + ] + }, + { + "name": "Danny Voss", + "kind": "character", + "dossier": "Danny Voss is Mara's missing brother.", + "provisional": true, + "aliases": [ + {"alias": "Danny", "kind": "name_variant"}, + {"alias": "her brother", "kind": "role_reference"}, + {"alias": "brother", "kind": "role_reference"} + ] + }, + { + "name": "Chapel on the Point", + "kind": "location", + "dossier": "The old chapel on the point has a loose floor stone where Tobias hides the leather ledger.", + "aliases": [ + {"alias": "Chapel", "kind": "name_variant"}, + {"alias": "chapel", "kind": "name_variant"}, + {"alias": "the old chapel on the point", "kind": "name_variant"}, + {"alias": "old chapel", "kind": "name_variant"} + ] + }, + { + "name": "Harbormaster's Office", + "kind": "location", + "dossier": "Tobias's office at the harbor.", + "aliases": [ + {"alias": "office", "kind": "name_variant"}, + {"alias": "harbormaster's office", "kind": "name_variant"} + ] + }, + { + "name": "Greyharbor Docks", + "kind": "location", + "dossier": "The docks in Greyharbor.", + "aliases": [ + {"alias": "docks", "kind": "name_variant"}, + {"alias": "Greyharbor Docks", "kind": "name_variant"} + ] + }, + { + "name": "Leather Ledger", + "kind": "object", + "dossier": "The real ledger of Thursday numbers hidden by Tobias.", + "aliases": [ + {"alias": "LEATHER LEDGER", "kind": "name_variant"}, + {"alias": "Leather Ledger", "kind": "name_variant"}, + {"alias": "ledger", "kind": "name_variant"}, + {"alias": "real ledger", "kind": "name_variant"}, + {"alias": "Thursday numbers", "kind": "nickname"}, + {"alias": "that book", "kind": "role_reference"} + ] + }, + { + "name": "Danny-Lee", + "kind": "object", + "dossier": "Danny's skiff.", + "aliases": [ + {"alias": "DANNY-LEE", "kind": "name_variant"}, + {"alias": "skiff", "kind": "name_variant"} + ] + } + ] +} diff --git a/extract/examples/greyharbor_ep101_ep102_scenes.json b/extract/examples/greyharbor_ep101_ep102_scenes.json new file mode 100644 index 0000000..2d3487f --- /dev/null +++ b/extract/examples/greyharbor_ep101_ep102_scenes.json @@ -0,0 +1,104 @@ +{ + "world": "greyharbor", + "scenes": [ + { + "scene_id": "E101/sc1", + "work_title": "E101 - The Ledger", + "scene_index": 1, + "slug": "INT. HARBORMASTER'S OFFICE - DAY", + "story_position": 1, + "is_flashback": false, + "raw_text": "INT. HARBORMASTER'S OFFICE - DAY\n\nMARA VOSS, 30s, salt-bitten and unsentimental, slaps a tide chart on the desk. Her uncle TOBIAS VOSS, 60s, doesn't look up from his ledger.\n\nMARA\nThree boats out past the shoals last night. None of them fishing.\n\nTOBIAS\nYou count boats now?\n\nMARA\nI count everything. You taught me that.\n\nTOBIAS\nI taught you to drive a truck, too, and look how that ended.\n\nMARA\nOne fence. I haven't touched a wheel since and you know I can't. License is gone for good.\n\nTobias finally looks up. Closes the ledger.\n\nTOBIAS\nLeave it alone, Mara." + }, + { + "scene_id": "E101/sc2", + "work_title": "E101 - The Ledger", + "scene_index": 2, + "slug": "EXT. GREYHARBOR DOCKS - NIGHT", + "story_position": 2, + "is_flashback": false, + "raw_text": "EXT. GREYHARBOR DOCKS - NIGHT\n\nMara walks the planks. EDDA QUINN, 50s, harbor cook, smoking by the gutting tables.\n\nEDDA\nYour uncle keeps two sets of numbers. The real one's not in that office.\n\nMARA\nYou've seen it?\n\nEDDA\nI've seen where he goes Thursday nights. The old chapel on the point." + }, + { + "scene_id": "E101/sc3", + "work_title": "E101 - The Ledger", + "scene_index": 3, + "slug": "INT. CHAPEL ON THE POINT - NIGHT", + "story_position": 3, + "is_flashback": false, + "raw_text": "INT. CHAPEL ON THE POINT - NIGHT\n\nCandle stubs. Tobias kneels at a loose floor stone, slides a LEATHER LEDGER beneath it.\n\nTOBIAS\n(to himself)\nForgive me the arithmetic." + }, + { + "scene_id": "E101/sc4", + "work_title": "E101 - The Ledger", + "scene_index": 4, + "slug": "EXT. GREYHARBOR DOCKS - NIGHT", + "story_position": 4, + "is_flashback": false, + "raw_text": "EXT. GREYHARBOR DOCKS - NIGHT\n\nMara finds her brother's skiff, the DANNY-LEE, tied off and empty. A coast guard NOTICE stapled to the mast: MISSING.\n\nMARA\n(quiet)\nI'll find you, Danny. Whatever it takes." + }, + { + "scene_id": "E101/sc5", + "work_title": "E101 - The Ledger", + "scene_index": 5, + "slug": "INT. HARBORMASTER'S OFFICE - NIGHT", + "story_position": 5, + "is_flashback": false, + "raw_text": "INT. HARBORMASTER'S OFFICE - NIGHT\n\nTobias at his desk. The door opens behind him. We don't see who enters.\n\nTOBIAS\nYou. I told them I needed more time --\n\nA gunshot. Tobias slumps over the ledger of fake numbers. Dead." + }, + { + "scene_id": "E101/sc6", + "work_title": "E101 - The Ledger", + "scene_index": 6, + "slug": "EXT. CHAPEL ON THE POINT - NIGHT", + "story_position": 6, + "is_flashback": false, + "raw_text": "EXT. CHAPEL ON THE POINT - NIGHT\n\nFlames climb the steeple. The chapel burns to its stones. No one comes to put it out.\n\nFADE OUT." + }, + { + "scene_id": "E102/sc1", + "work_title": "E102 - Thursday Numbers", + "scene_index": 1, + "slug": "EXT. GREYHARBOR DOCKS - DAY", + "story_position": 7, + "is_flashback": false, + "raw_text": "EXT. GREYHARBOR DOCKS - DAY\n\nPolice tape on the harbormaster's door. Mara stares at it. DEPUTY COLE BRANNIGAN, 40s, approaches, hat in hand.\n\nCOLE\nWe're saying robbery. Till was open.\n\nMARA\nThe till's always open. Nobody steals from a man who knows every hull in the harbor.\n\nCOLE\nWhoever it was, they didn't find the real ledger. It's still under the chapel floor stone where your uncle hid it.\n\nMARA\nHow would you know where he hid anything?\n\nCole stiffens. Tips his hat. Walks off." + }, + { + "scene_id": "E102/sc2", + "work_title": "E102 - Thursday Numbers", + "scene_index": 2, + "slug": "INT. CHAPEL ON THE POINT - DAY", + "story_position": 8, + "is_flashback": false, + "raw_text": "INT. CHAPEL ON THE POINT - DAY\n\nMara pries up the loose floor stone. The LEATHER LEDGER, intact. She flips pages: names, weights, payments.\n\nMARA\nForgive me the arithmetic." + }, + { + "scene_id": "E102/sc3", + "work_title": "E102 - Thursday Numbers", + "scene_index": 3, + "slug": "INT. HARBORMASTER'S OFFICE - DAY", + "story_position": 9, + "is_flashback": false, + "raw_text": "INT. HARBORMASTER'S OFFICE - DAY\n\nMara spreads the ledger pages on the desk. TOBIAS stands at the window, watching the boats.\n\nTOBIAS\nNow you've counted everything. Happy?\n\nMARA\nWho were you protecting?\n\nTOBIAS\nWrong question. Ask who was protecting me." + }, + { + "scene_id": "E102/sc4", + "work_title": "E102 - Thursday Numbers", + "scene_index": 4, + "slug": "EXT. COAST ROAD - DUSK", + "story_position": 10, + "is_flashback": false, + "raw_text": "EXT. COAST ROAD - DUSK\n\nMara behind the wheel of Tobias's pickup, driving hard along the cliff road, ledger on the passenger seat." + }, + { + "scene_id": "E102/sc5", + "work_title": "E102 - Thursday Numbers", + "scene_index": 5, + "slug": "EXT. EDDA'S SHACK - NIGHT", + "story_position": 11, + "is_flashback": false, + "raw_text": "EXT. EDDA'S SHACK - NIGHT\n\nEdda opens the door before Mara knocks.\n\nEDDA\nYou found the Thursday numbers.\n\nMARA\nCole knew where they were. Before I did.\n\nEDDA\nThen the deputy's name is in that book. Look under the weights nobody could land alone.\n\nMara opens the ledger under the porch light.\n\nMARA\n(reading)\n\"C.B. -- four hundred, the shoals run.\"\n\nShe looks up. Somewhere out on the water, an engine starts.\n\nFADE OUT." + } + ] +} diff --git a/extract/examples/marcus_aliases.json b/extract/examples/marcus_aliases.json new file mode 100644 index 0000000..f73f722 --- /dev/null +++ b/extract/examples/marcus_aliases.json @@ -0,0 +1,14 @@ +{ + "entities": [ + { + "name": "Marcus Hale", + "kind": "character", + "dossier": "Marcus Hale is the husband referred to in the scene.", + "aliases": [ + {"alias": "MARCUS", "kind": "name_variant"}, + {"alias": "Marcus", "kind": "name_variant"}, + {"alias": "her husband", "kind": "role_reference"} + ] + } + ] +} diff --git a/extract/examples/marcus_candidates.json b/extract/examples/marcus_candidates.json new file mode 100644 index 0000000..1f32090 --- /dev/null +++ b/extract/examples/marcus_candidates.json @@ -0,0 +1,58 @@ +{ + "world": "alias-test", + "scenes": [ + { + "scene_id": "TST/sc1", + "work_title": "TST - Alias Test", + "scene_index": 1, + "slug": "INT. KITCHEN - NIGHT", + "story_position": 1, + "is_flashback": false, + "assertions": [ + { + "subject": "MARCUS", + "predicate": "occupation", + "object_entity": null, + "object_value": "doctor", + "object_fact_ref": null, + "polarity": true, + "starts_here": true, + "ends_here": false, + "supporting_quote": "MARCUS, still in his hospital badge, enters.", + "confidence": 0.94, + "notes": null + }, + { + "subject": "Marcus Hale", + "predicate": "possesses", + "object_entity": "watch", + "object_value": null, + "object_fact_ref": null, + "polarity": true, + "starts_here": true, + "ends_here": false, + "supporting_quote": "Marcus Hale sets the watch on the counter.", + "confidence": 0.91, + "notes": null + }, + { + "subject": "her husband", + "predicate": "knows", + "object_entity": null, + "object_value": "safe code", + "object_fact_ref": null, + "polarity": true, + "starts_here": true, + "ends_here": false, + "supporting_quote": "Her husband knows the safe code.", + "confidence": 0.8, + "notes": null + } + ], + "scene_presence": ["MARCUS", "Marcus Hale", "her husband"], + "deaths": [], + "destructions": [], + "open_questions": [] + } + ] +} diff --git a/extract/extraction.py b/extract/extraction.py new file mode 100644 index 0000000..a6e69c8 --- /dev/null +++ b/extract/extraction.py @@ -0,0 +1,185 @@ +"""Per-scene LLM extraction and quote gates.""" + +from __future__ import annotations + +from dataclasses import asdict +from typing import Any, Callable + +import json +import re + +from .llm import DEFAULT_EFFORT, DEFAULT_MODEL, MAX_EXTRACTION_TOKENS, first_text, structured_request +from .models import SceneExtraction, SceneRecord +from .prompts import build_synopsis, build_user_prompt, extraction_system_prompt +from .schema import INTRANSITIVE, PREDICATES, candidate_schema + + +def _normalize(text: str) -> str: + text = ( + (text or "") + .replace("'", "'") + .replace("\u2018", "'") + .replace("\u2019", "'") + .replace("\u201c", '"') + .replace("\u201d", '"') + .replace("\u2013", "-") + .replace("\u2014", "-") + ) + return " ".join(text.split()).casefold() + + +def quote_in_text(quote: str, text: str) -> bool: + return bool(quote and _normalize(quote) in _normalize(text)) + + +def clamp_confidence(value: Any) -> float: + try: + conf = float(value) + except (TypeError, ValueError): + conf = 0.0 + return max(0.0, min(1.0, conf)) + + +def post_process(scene: SceneRecord, payload: dict[str, Any], verify_quotes: bool = True) -> tuple[list[dict], dict]: + """Apply the non-negotiable extraction quality gates.""" + + dropped = {"no_quote": 0, "bad_predicate": 0, "no_object": 0, "quote_unverified": 0} + kept: list[dict[str, Any]] = [] + for assertion in payload.get("assertions") or []: + quote = (assertion.get("supporting_quote") or "").strip() + predicate = assertion.get("predicate") + has_object = bool( + assertion.get("object_entity") + or assertion.get("object_value") + or assertion.get("object_fact_ref") + ) + if not quote: + dropped["no_quote"] += 1 + continue + if predicate not in PREDICATES: + dropped["bad_predicate"] += 1 + continue + if not has_object and predicate not in INTRANSITIVE: + dropped["no_object"] += 1 + continue + if verify_quotes and not quote_in_text(quote, scene.raw_text): + dropped["quote_unverified"] += 1 + continue + + row = dict(assertion) + row["supporting_quote"] = quote + row["confidence"] = clamp_confidence(row.get("confidence")) + row["polarity"] = bool(row.get("polarity", True)) + row["starts_here"] = bool(row.get("starts_here", True)) + row["ends_here"] = bool(row.get("ends_here", False)) + row["scene_id"] = scene.scene_id + row["scene"] = scene.label + row["story_position"] = scene.story_position + kept.append(row) + return kept, dropped + + +def extract_scene( + client, + scene: SceneRecord, + prior: list[SceneExtraction], + *, + model: str = DEFAULT_MODEL, + effort: str = DEFAULT_EFFORT, + thinking: bool = True, + verify_quotes: bool = True, +) -> SceneExtraction: + request = structured_request( + model=model, + system=extraction_system_prompt(), + user=build_user_prompt(scene, build_synopsis(prior)), + schema=candidate_schema(), + max_tokens=MAX_EXTRACTION_TOKENS, + effort=effort, + thinking=thinking, + ) + response = client.messages.create(**request) + if getattr(response, "stop_reason", None) == "refusal": + raise RuntimeError(f"model refused scene {scene.scene_id} ({scene.slug})") + text = first_text(response) + if text is None: + raise RuntimeError(f"no text block in response for scene {scene.scene_id}") + payload = json.loads(text) + assertions, dropped = post_process(scene, payload, verify_quotes=verify_quotes) + return SceneExtraction( + scene_id=scene.scene_id, + work_title=scene.work_title, + scene_index=scene.scene_index, + slug=scene.slug, + story_position=scene.story_position, + is_flashback=scene.is_flashback, + assertions=assertions, + scene_presence=payload.get("scene_presence") or [], + deaths=payload.get("deaths") or [], + destructions=payload.get("destructions") or [], + open_questions=payload.get("open_questions") or [], + dropped=dropped, + ) + + +def run_extraction( + client, + scenes: list[SceneRecord], + *, + model: str = DEFAULT_MODEL, + effort: str = DEFAULT_EFFORT, + thinking: bool = True, + verify_quotes: bool = True, + limit: int | None = None, + on_result: Callable[[SceneExtraction], None] | None = None, +) -> list[SceneExtraction]: + if limit is not None: + scenes = scenes[:limit] + results: list[SceneExtraction] = [] + for scene in scenes: + result = extract_scene( + client, + scene, + results, + model=model, + effort=effort, + thinking=thinking, + verify_quotes=verify_quotes, + ) + results.append(result) + if on_result: + on_result(result) + return results + + +def to_candidate_json(world: str, model: str, scenes: list[SceneExtraction]) -> dict[str, Any]: + return {"world": world, "model": model, "scenes": [asdict(scene) for scene in scenes]} + + +def render_dry_run( + scenes: list[SceneRecord], + *, + model: str = DEFAULT_MODEL, + effort: str = DEFAULT_EFFORT, + limit: int | None = None, +) -> str: + shown = scenes[: (limit or 1)] + lines = [ + f"model: {model} effort: {effort} thinking: adaptive", + f"scenes to extract: {len(scenes)}", + "", + "=== SYSTEM PROMPT ===", + extraction_system_prompt(), + "", + ] + prior: list[SceneExtraction] = [] + for scene in shown: + lines.append(f"=== USER PROMPT - {scene.scene_id} | {scene.slug} ===") + lines.append(build_user_prompt(scene, build_synopsis(prior))) + lines.append("") + if len(scenes) > len(shown): + lines.append( + f"... {len(scenes) - len(shown)} more scene(s); each receives a rolling synopsis " + "built from prior scene extractions." + ) + return "\n".join(lines) diff --git a/extract/gate.py b/extract/gate.py new file mode 100644 index 0000000..2db56b8 --- /dev/null +++ b/extract/gate.py @@ -0,0 +1,151 @@ +"""Confidence gate and assertion confirmation queue.""" + +from __future__ import annotations + +from typing import Any + +import json + +from .schema import AUTO_ACCEPT_CONFIDENCE, PREDICATES +from .resolution import ResolutionState + + +def _confidence(value: Any) -> float: + try: + return max(0.0, min(1.0, float(value))) + except (TypeError, ValueError): + return 0.0 + + +def assertion_summary(assertion: dict[str, Any]) -> str: + obj = assertion.get("object_entity") or assertion.get("object_value") or "" + neg = "" if assertion.get("polarity", True) else "NOT " + return f"{neg}{assertion.get('subject')} {assertion.get('predicate')} {obj}".strip() + + +def gate_state(state: ResolutionState, threshold: float = AUTO_ACCEPT_CONFIDENCE) -> ResolutionState: + """Apply Stage 4 confidence gate in-place and return state. + + `confidence >= threshold` becomes `draft`. Lower confidence assertions are + queued for human accept/edit/reject. Human-confirmed assertions later become + `canon`. + """ + + queue: list[dict[str, Any]] = [] + for idx, assertion in enumerate(state.assertions): + if assertion.get("status") in {"canon", "rejected"}: + continue + conf = _confidence(assertion.get("confidence")) + assertion["confidence"] = conf + if conf >= threshold: + assertion["status"] = "draft" + assertion["confirmed_by_human"] = bool(assertion.get("confirmed_by_human", False)) + continue + assertion["status"] = "needs_confirmation" + queue.append( + { + "assertion_index": idx, + "reason": "low_confidence", + "confidence": conf, + "scene": assertion.get("scene"), + "scene_id": assertion.get("scene_id"), + "subject": assertion.get("subject"), + "predicate": assertion.get("predicate"), + "object_entity": assertion.get("object_entity"), + "object_value": assertion.get("object_value"), + "supporting_quote": assertion.get("supporting_quote"), + "summary": assertion_summary(assertion), + } + ) + state.assertion_queue = queue + return state + + +def apply_assertion_decision(state: ResolutionState, item: dict[str, Any], decision: str) -> bool: + """Apply one human assertion decision. + + Grammar: `accept`, `reject`, `edit {"field": "value"}`, or `skip`. + Accepted or edited assertions become `canon` with `confirmed_by_human`. + """ + + choice = (decision or "").strip() + if not choice or choice.lower() in {"skip", "s"}: + return False + + idx = int(item["assertion_index"]) + if not (0 <= idx < len(state.assertions)): + return True + assertion = state.assertions[idx] + low = choice.lower() + + if low in {"accept", "a"}: + assertion["status"] = "canon" + assertion["confirmed_by_human"] = True + return True + + if low in {"reject", "r"}: + assertion["status"] = "rejected" + assertion["confirmed_by_human"] = True + assertion["rejection_reason"] = "human rejected from confidence queue" + return True + + if low.startswith("edit ") or low.startswith("e "): + payload = choice.split(" ", 1)[1].strip() + try: + patch = json.loads(payload) + except json.JSONDecodeError: + return False + if patch.get("predicate") and patch["predicate"] not in PREDICATES: + return False + for key, value in patch.items(): + if key in { + "subject", + "predicate", + "object_entity", + "object_value", + "polarity", + "starts_here", + "ends_here", + "supporting_quote", + "confidence", + "notes", + }: + assertion[key] = value + assertion["status"] = "canon" + assertion["confirmed_by_human"] = True + assertion["human_edit"] = True + return True + + return False + + +def walk_assertion_queue(state: ResolutionState, prompt_fn) -> int: + remaining: list[dict[str, Any]] = [] + resolved = 0 + for item in state.assertion_queue: + if apply_assertion_decision(state, item, prompt_fn(item, state)): + resolved += 1 + else: + remaining.append(item) + state.assertion_queue = remaining + return resolved + + +def to_eval_assertions(state: ResolutionState) -> dict[str, list[dict[str, Any]]]: + """Flat eval contract, excluding rejected assertions.""" + + return { + "assertions": [ + { + "subject": assertion.get("subject"), + "predicate": assertion.get("predicate"), + "object_value": assertion.get("object_value"), + "object_entity": assertion.get("object_entity"), + "scene": assertion.get("scene"), + "confidence": assertion.get("confidence"), + "status": assertion.get("status"), + } + for assertion in state.assertions + if assertion.get("status") != "rejected" + ] + } diff --git a/extract/llm.py b/extract/llm.py new file mode 100644 index 0000000..45d5af3 --- /dev/null +++ b/extract/llm.py @@ -0,0 +1,66 @@ +"""Anthropic client helpers for structured-output calls.""" + +from __future__ import annotations + +import os +from typing import Any + +DEFAULT_MODEL = "claude-opus-4-8" +DEFAULT_EFFORT = "high" +MAX_EXTRACTION_TOKENS = 16000 +MAX_RESOLUTION_TOKENS = 4000 + + +def has_credentials() -> bool: + return bool(os.environ.get("ANTHROPIC_API_KEY") or os.environ.get("ANTHROPIC_AUTH_TOKEN")) + + +def make_client(): + try: + import anthropic + except ModuleNotFoundError as exc: + raise RuntimeError( + "anthropic is not installed. Install requirements or run with --dry-run/--no-llm." + ) from exc + # No per-request training flag exists in the SDK. The configured key must + # belong to a no-training / zero-data-retention workspace. + return anthropic.Anthropic() + + +def first_text(response: Any) -> str | None: + for block in getattr(response, "content", None) or []: + if getattr(block, "type", None) == "text": + return block.text + return None + + +def structured_request( + *, + model: str, + system: str, + user: str, + schema: dict, + max_tokens: int, + effort: str = DEFAULT_EFFORT, + thinking: bool = True, +) -> dict: + """Build a native Claude Messages API request. + + Intentionally omits sampling params such as `temperature`; current Opus + models reject non-default sampling parameters, and the extraction contract is + enforced by schema plus conservative gates. + """ + + output_config: dict[str, Any] = {"format": {"type": "json_schema", "schema": schema}} + if effort: + output_config["effort"] = effort + request: dict[str, Any] = { + "model": model, + "max_tokens": max_tokens, + "system": system, + "messages": [{"role": "user", "content": user}], + "output_config": output_config, + } + if thinking: + request["thinking"] = {"type": "adaptive"} + return request diff --git a/extract/models.py b/extract/models.py new file mode 100644 index 0000000..340b7db --- /dev/null +++ b/extract/models.py @@ -0,0 +1,114 @@ +"""JSON contracts and dataclasses for P0 extraction.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +import json + + +@dataclass +class SceneRecord: + """Scene JSON consumed by this workstream. + + Required JSON fields are `slug`, `story_position`, `is_flashback`, and + `raw_text`. `scene_id`, `work_title`, and `scene_index` are strongly + recommended because they make citations stable before the DB exists. + """ + + scene_id: str + work_title: str + scene_index: int + slug: str + story_position: int + is_flashback: bool + raw_text: str + source_file: str | None = None + + @property + def label(self) -> str: + if self.scene_id: + return self.scene_id + head = (self.work_title or "?").split()[0] + return f"{head}/sc{self.scene_index}" + + +@dataclass +class SceneExtraction: + scene_id: str + work_title: str + scene_index: int + slug: str + story_position: int + is_flashback: bool + assertions: list[dict[str, Any]] = field(default_factory=list) + scene_presence: list[str] = field(default_factory=list) + deaths: list[str] = field(default_factory=list) + destructions: list[str] = field(default_factory=list) + open_questions: list[str] = field(default_factory=list) + dropped: dict[str, int] = field(default_factory=dict) + + +def _coerce_int(value: Any, fallback: int, field_name: str) -> int: + if value is None or value == "": + return fallback + try: + return int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field_name} must be an integer, got {value!r}") from exc + + +def scene_from_dict(row: dict[str, Any], index: int) -> SceneRecord: + slug = str(row.get("slug") or row.get("scene_slug") or "").strip() + raw_text = str(row.get("raw_text") or row.get("text") or "").strip() + if not slug: + raise ValueError(f"scene {index}: missing slug") + if not raw_text: + raise ValueError(f"scene {index}: missing raw_text") + + story_position = _coerce_int(row.get("story_position"), index, "story_position") + scene_index = _coerce_int( + row.get("scene_index") or row.get("scene_number"), + index, + "scene_index", + ) + work_title = str(row.get("work_title") or row.get("work") or row.get("episode") or "").strip() + scene_id = str(row.get("scene_id") or row.get("id") or "").strip() + if not scene_id: + head = (work_title or "?").split()[0] + scene_id = f"{head}/sc{scene_index}" + return SceneRecord( + scene_id=scene_id, + work_title=work_title, + scene_index=scene_index, + slug=slug, + story_position=story_position, + is_flashback=bool(row.get("is_flashback", False)), + raw_text=raw_text, + source_file=row.get("source_file"), + ) + + +def load_scenes(path: str | Path) -> tuple[str, list[SceneRecord]]: + """Load scene records from either `{world, scenes}` or a bare scene list.""" + + data = json.loads(Path(path).read_text(encoding="utf-8")) + if isinstance(data, dict): + world = str(data.get("world") or "") + rows = data.get("scenes") or [] + elif isinstance(data, list): + world = "" + rows = data + else: + raise ValueError("scene input must be a JSON object with scenes or a list") + if not isinstance(rows, list): + raise ValueError("scenes must be a list") + scenes = [scene_from_dict(row, i) for i, row in enumerate(rows, start=1)] + scenes.sort(key=lambda s: (s.story_position, s.scene_id)) + return world, scenes + + +def extraction_to_dict(ex: SceneExtraction) -> dict[str, Any]: + return asdict(ex) diff --git a/extract/prompts.py b/extract/prompts.py new file mode 100644 index 0000000..65f0801 --- /dev/null +++ b/extract/prompts.py @@ -0,0 +1,63 @@ +"""Prompt assembly from versioned prompt files.""" + +from __future__ import annotations + +from pathlib import Path + +from .models import SceneExtraction, SceneRecord +from .schema import PREDICATES + +PROMPT_DIR = Path(__file__).resolve().parent / "prompts" + + +def read_prompt(name: str) -> str: + return (PROMPT_DIR / name).read_text(encoding="utf-8") + + +def extraction_system_prompt() -> str: + return read_prompt("extraction_system.md").format(predicate_list=", ".join(PREDICATES)) + + +def resolution_system_prompt() -> str: + return read_prompt("entity_resolution_system.md") + + +def build_synopsis(prior: list[SceneExtraction], facts_per_scene: int = 12) -> str: + """Compact rolling context from prior extracted facts. + + This is deterministic indexing text, not story generation. + """ + + lines: list[str] = [] + for ex in prior: + facts: list[str] = [] + for a in ex.assertions: + obj = a.get("object_entity") or a.get("object_value") or a.get("object_fact_ref") or "" + neg = "" if a.get("polarity", True) else "NOT " + facts.append(f"{neg}{a.get('subject', '?')} {a.get('predicate', '?')} {obj}".strip()) + line = f"[{ex.scene_id} | {ex.slug}] present: {', '.join(ex.scene_presence)}" + if facts: + line += " | " + "; ".join(facts[:facts_per_scene]) + lines.append(line) + return "\n".join(lines) + + +def build_user_prompt(scene: SceneRecord, synopsis: str) -> str: + prior = ( + "PRIOR SCENES (context only - do NOT re-extract these):\n" + synopsis + if synopsis + else "PRIOR SCENES: none - this is the first scene." + ) + flashback = " [FLASHBACK]" if scene.is_flashback else "" + return ( + f"{prior}\n\n" + "CURRENT SCENE - extract assertions for THIS scene only.\n" + f"scene_id: {scene.scene_id}\n" + f"work: {scene.work_title}\n" + f"scene: {scene.slug}{flashback}\n" + f"story_position: {scene.story_position}\n" + "---\n" + f"{scene.raw_text}\n" + "---\n\n" + "Return the structured candidate assertions for this scene." + ) diff --git a/extract/prompts/entity_resolution_system.md b/extract/prompts/entity_resolution_system.md new file mode 100644 index 0000000..0ee4cd8 --- /dev/null +++ b/extract/prompts/entity_resolution_system.md @@ -0,0 +1,13 @@ +You are the entity-resolution assistant for Canon AI, a continuity index for fiction. + +You never write, invent, continue, or improve story content. You resolve references using only the provided occurrences, aliases, and entity dossiers. + +Given one referring expression and the entities already known in this world, decide whether the expression denotes one known entity, a brand-new entity, or is unclear. + +Rules: +- Choose at most one known entity. +- Never merge two known entities. Merging is a human-only CLI action. +- Prefer an existing entity when a supplied alias, dossier, role reference, or occurrence makes the match clear. +- If the reference is ambiguous or under-supported, return `unsure`. +- If the reference is a role reference such as `her husband`, `the deputy`, or `your uncle`, use the provided scene occurrence and known dossiers. Do not guess relationship facts not present in the prompt. +- Confidence must reflect evidence quality, not plausibility. diff --git a/extract/prompts/extraction_system.md b/extract/prompts/extraction_system.md new file mode 100644 index 0000000..2775502 --- /dev/null +++ b/extract/prompts/extraction_system.md @@ -0,0 +1,57 @@ +You are the extraction engine for Canon AI, a continuity and canon system for serialized fiction. + +YOUR ONE HARD RULE: you never write, invent, continue, or paraphrase story content. You do not speculate about what might happen or fill gaps. You record only assertions the scene text explicitly states or directly implies, each tied to a verbatim quote. You are an indexer, not an author. + +For the CURRENT scene only, extract candidate continuity assertions as structured output. + +PREDICATES - choose exactly one from this closed set. Never invent one: +{predicate_list} + +Use `fact` with a free-text object_value only for true, citable facts that no other predicate fits. `alive` is implicit at a character's first appearance only when the scene directly establishes that person is alive or active. Do not emit `alive` for every name mentioned. Emit `dies` when a character dies. + +EVIDENCE: +- Every assertion MUST include `supporting_quote`: text copied verbatim from the current scene. If you cannot quote it from THIS scene, do not emit the assertion. +- Extract only what the text states or directly implies. No inference chains. Inference is the checks layer's job. +- Direct implication includes: acting on information, titles or signage implying occupation, and an object being placed or found somewhere implying `located_at`. + +OBJECT VALUES ARE CANONICAL HANDLES: +- `object_value` is a short handle, usually 2-5 words, not a sentence: `ledger location`, `two sets of numbers`, `drive`, `find Danny`. +- The same fact gets the same handle everywhere. Before creating a handle, scan PRIOR SCENES and copy the existing handle character-for-character when it is the same fact. +- Use these handle shapes where applicable: ` location`; bare infinitive verbs for capabilities; `find ` for search goals; ` payment` for money entries. +- Reuse a handle only for the same fact. Related but distinct facts get distinct handles. +- For `dies` and `destroyed`, leave object fields null and put manner/detail in `notes`. + +KNOWS VS BELIEVES: +- `knows`: first-hand or directly supplied information. The character saw it, did it, read it, was told it on screen, or acts on it. +- `believes`: suspicion, inference, hearsay, or an unverifiable claim. +- A character's dialogue claim about the world is `believes(speaker, X)`, not a world fact, unless action lines or independent sources corroborate it. Characters can lie. +- A first-person statement about something the speaker directly did/saw can establish the epistemic assertion `knows(speaker, X)`. + +CAPABILITY VIOLATIONS AND ACTIVITIES: +- When a character performs a notable physical activity on screen (driving, swimming, shooting, riding), emit `fact(subject, )`. +- If PRIOR SCENES establish `cannot(X, handle)` and X now performs that same action, also emit `cannot(X, same handle)` with `polarity: false` quoting the action line. + +POLARITY AND NEGATION: +- `polarity: true` means the assertion holds. +- Use `polarity: false` for explicit negation of a relational predicate. +- For inability, use predicate `cannot` with `polarity: true` and the capability in `object_value`. + +TIME: +- Default `starts_here: true` and `ends_here: false`. +- BACKDATING: if dialogue establishes that a fact began earlier, set `starts_here: false` and explain in `notes`. Do not guess the earlier position. + +OBJECT FIELDS: +- `object_entity`: use when the object is another named entity. +- `object_value`: use for literal values and canonical fact handles. +- `object_fact_ref`: use only when the assertion is explicitly about another fact and `object_value` would lose necessary structure. +- Set unused object fields to null. + +ALSO REPORT: +- `scene_presence`: every character physically present or speaking in the scene. +- `deaths`: characters who die in this scene. +- `destructions`: locations or objects destroyed in this scene. +- `open_questions`: setups, promises, or unresolved questions raised by the scene. + +CONFIDENCE: +- Return 0..1 calibrated confidence that the assertion is correct and citable. +- Be conservative. A missed borderline assertion is cheaper than a wrong confident one. diff --git a/extract/resolution.py b/extract/resolution.py new file mode 100644 index 0000000..e8ba400 --- /dev/null +++ b/extract/resolution.py @@ -0,0 +1,789 @@ +"""Entity resolution for candidate assertions. + +Stage contract: exact alias match -> fuzzy match -> LLM disambiguation with +dossiers. Unresolved role references become provisional entities and queue +items. Two existing entities are never merged automatically. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import json +import re + +from .llm import DEFAULT_EFFORT, DEFAULT_MODEL, MAX_RESOLUTION_TOKENS, first_text, structured_request +from .prompts import resolution_system_prompt +from .schema import ALIAS_KINDS, AUTO_ACCEPT_CONFIDENCE, ENTITY_KINDS, resolution_schema + +ROLE_WORDS = frozenset( + { + "uncle", + "aunt", + "brother", + "sister", + "father", + "mother", + "mom", + "dad", + "son", + "daughter", + "cousin", + "nephew", + "niece", + "husband", + "wife", + "widow", + "deputy", + "sheriff", + "captain", + "cook", + "boss", + "stranger", + "man", + "woman", + "boy", + "girl", + "kid", + "the man", + "the woman", + } +) +PRONOUNS = frozenset({"he", "she", "they", "him", "her", "them", "it", "his", "hers", "its", "their"}) +INITIALS_RE = re.compile(r"^([A-Za-z]\.){2,}$") + +CHARACTER_SUBJECT_PREDS = frozenset( + { + "knows", + "believes", + "cannot", + "dies", + "occupation", + "married_to", + "parent_of", + "sibling_of", + "romantic_with", + "allied_with", + "enemy_of", + "member_of", + "possesses", + "promised", + "goal", + "secret_of", + "trait", + "alive", + "present_in_scene", + } +) +OBJECT_KIND_BY_PRED = { + "located_at": "location", + "possesses": "object", + "member_of": "faction", + "married_to": "character", + "parent_of": "character", + "sibling_of": "character", + "romantic_with": "character", + "allied_with": "character", + "enemy_of": "character", + "created": "object", + "secret_of": "character", +} +KIND_PRIORITY = {k: i for i, k in enumerate(("character", "location", "object", "faction", "event", "rule", "other"))} + + +def norm(text: str | None) -> str: + return re.sub(r"[^a-z0-9 ]", "", (text or "").lower()).strip() + + +def tokens(text: str | None) -> set[str]: + return set(norm(text).split()) + + +def display_name(text: str) -> str: + text = (text or "").strip() + if text and text.isupper() and any(c.isalpha() for c in text): + return text.title() + return text + + +def is_initials(text: str) -> bool: + return bool(INITIALS_RE.match((text or "").strip())) + + +def strip_articles(text: str) -> str: + parts = (text or "").strip().split() + while parts and parts[0].lower() in ("the", "a", "an"): + parts = parts[1:] + return " ".join(parts) + + +def is_role_reference(surface: str) -> bool: + if is_initials(surface): + return True + n = norm(strip_articles(surface)) + return (not n) or n in ROLE_WORDS or n in PRONOUNS + + +def occ_kind(role: str, predicate: str | None) -> str: + if role in ("presence", "death"): + return "character" + if role == "destruction": + return "location" + if role == "subject": + if predicate in ("located_at", "created"): + return "object" + if predicate == "destroyed": + return "location" + if predicate in CHARACTER_SUBJECT_PREDS: + return "character" + return "other" + if role == "object": + return OBJECT_KIND_BY_PRED.get(predicate, "other") + return "other" + + +def scene_label(scene: dict[str, Any]) -> str: + if scene.get("scene_id"): + return scene["scene_id"] + head = (scene.get("work_title") or "?").split()[0] + return f"{head}/sc{scene.get('scene_index', 0)}" + + +@dataclass +class Mention: + surface: str + role: str + predicate: str | None + story_position: int + scene: str + slug: str + quote: str | None + + +@dataclass +class SurfaceInfo: + display: str + norm: str + kind_hint: str + first_position: int + is_role: bool + occurrences: list[Mention] = field(default_factory=list) + + +@dataclass(eq=False) +class Entity: + local_id: int + name: str + kind: str + dossier: str | None + provisional: bool + seeded: bool = False + aliases: list[list[str]] = field(default_factory=list) + + def alias_norms(self) -> set[str]: + return {norm(self.name)} | {norm(alias) for alias, _kind in self.aliases} + + def token_sets(self) -> list[set[str]]: + return [tokens(self.name)] + [tokens(alias) for alias, _kind in self.aliases] + + +class Registry: + def __init__(self) -> None: + self.entities: list[Entity] = [] + self._next = 1 + + def add( + self, + name: str, + kind: str, + dossier: str | None, + provisional: bool, + *, + alias_kind: str = "name_variant", + seeded: bool = False, + ) -> Entity: + entity = Entity( + local_id=self._next, + name=display_name(name), + kind=kind if kind in ENTITY_KINDS else "other", + dossier=dossier, + provisional=provisional, + seeded=seeded, + aliases=[[name, alias_kind if alias_kind in ALIAS_KINDS else "name_variant"]], + ) + self._next += 1 + self.entities.append(entity) + return entity + + def add_alias(self, entity: Entity, surface: str, alias_kind: str = "name_variant", upgrade: bool = True) -> None: + alias_kind = alias_kind if alias_kind in ALIAS_KINDS else "name_variant" + if norm(surface) not in entity.alias_norms(): + entity.aliases.append([surface, alias_kind]) + if ( + upgrade + and alias_kind == "name_variant" + and (len(tokens(surface)), len(surface)) > (len(tokens(entity.name)), len(entity.name)) + ): + entity.name = display_name(surface) + + def by_name(self, name: str | None) -> Entity | None: + n = norm(name) + if not n: + return None + return next((entity for entity in self.entities if n in entity.alias_norms()), None) + + def match(self, surface: str) -> tuple[str, Entity | list[Entity] | None]: + n = norm(surface) + ts = tokens(surface) + if not n: + return "none", None + exact = [entity for entity in self.entities if n in entity.alias_norms()] + if len(exact) == 1: + return "exact", exact[0] + if len(exact) > 1: + return "ambiguous", exact + + candidates: list[Entity] = [] + for entity in self.entities: + for known_tokens in entity.token_sets(): + if not ts or not known_tokens or ts == known_tokens: + continue + if ts <= known_tokens or known_tokens <= ts: + candidates.append(entity) + break + if len(ts & known_tokens) / len(ts | known_tokens) >= 0.5: + candidates.append(entity) + break + candidates = list(dict.fromkeys(candidates)) + if len(candidates) == 1: + return "fuzzy", candidates[0] + if len(candidates) > 1: + return "ambiguous", candidates + return "none", None + + def merge(self, keep: Entity, drop: Entity, reason: str) -> None: + """Human-only operation. Call only from confirm/merge CLI paths.""" + + if keep is drop: + return + for alias, alias_kind in ([[drop.name, "name_variant"]] + drop.aliases): + self.add_alias(keep, alias, alias_kind, upgrade=False) + keep.provisional = keep.provisional and drop.provisional + keep.seeded = keep.seeded or drop.seeded + if not keep.dossier and drop.dossier: + keep.dossier = drop.dossier + self.entities = [entity for entity in self.entities if entity is not drop] + + +def registry_from_alias_table(alias_table: dict[str, Any] | None) -> Registry: + reg = Registry() + if not alias_table: + return reg + + if isinstance(alias_table.get("entities"), list): + for row in alias_table["entities"]: + name = row.get("canonical_name") or row.get("name") + if not name: + continue + entity = reg.add( + name, + row.get("kind") or "other", + row.get("dossier"), + bool(row.get("provisional", False)), + seeded=True, + ) + for alias_row in row.get("aliases") or []: + if isinstance(alias_row, str): + reg.add_alias(entity, alias_row, "name_variant", upgrade=False) + elif alias_row.get("alias"): + reg.add_alias(entity, alias_row["alias"], alias_row.get("kind") or "name_variant", upgrade=False) + + if isinstance(alias_table.get("aliases"), list): + for row in alias_table["aliases"]: + canonical = row.get("canonical_name") or row.get("entity") or row.get("name") + alias = row.get("alias") + if not canonical or not alias: + continue + entity = reg.by_name(canonical) + if entity is None: + entity = reg.add(canonical, row.get("entity_kind") or "other", None, False, seeded=True) + reg.add_alias(entity, alias, row.get("kind") or "name_variant", upgrade=False) + return reg + + +def scenes_from_candidates(candidates: dict[str, Any] | list[dict[str, Any]]) -> list[dict[str, Any]]: + if isinstance(candidates, dict): + return candidates.get("scenes") or [] + return list(candidates or []) + + +def collect_mentions(candidates: dict[str, Any] | list[dict[str, Any]]) -> list[Mention]: + mentions: list[Mention] = [] + for scene in scenes_from_candidates(candidates): + label = scene_label(scene) + story_position = int(scene.get("story_position") or 0) + slug = scene.get("slug") or "" + + def add(surface: Any, role: str, predicate: str | None = None, quote: str | None = None) -> None: + if surface and str(surface).strip(): + mentions.append(Mention(str(surface).strip(), role, predicate, story_position, label, slug, quote)) + + for assertion in scene.get("assertions") or []: + add(assertion.get("subject"), "subject", assertion.get("predicate"), assertion.get("supporting_quote")) + if assertion.get("object_entity"): + add(assertion.get("object_entity"), "object", assertion.get("predicate"), assertion.get("supporting_quote")) + for name in scene.get("scene_presence") or []: + add(name, "presence") + for name in scene.get("deaths") or []: + add(name, "death") + for name in scene.get("destructions") or []: + add(name, "destruction") + return mentions + + +def vote_kind(occurrences: list[Mention]) -> str: + votes: dict[str, int] = {} + for mention in occurrences: + kind = occ_kind(mention.role, mention.predicate) + votes[kind] = votes.get(kind, 0) + 1 + return min(votes, key=lambda k: (-votes[k], KIND_PRIORITY.get(k, 99))) if votes else "other" + + +def group_surfaces(mentions: list[Mention]) -> list[SurfaceInfo]: + groups: dict[str, SurfaceInfo] = {} + for mention in mentions: + n = norm(mention.surface) + if not n: + continue + info = groups.get(n) + if info is None: + info = SurfaceInfo( + display=mention.surface, + norm=n, + kind_hint="other", + first_position=mention.story_position, + is_role=is_role_reference(mention.surface), + ) + groups[n] = info + info.occurrences.append(mention) + info.first_position = min(info.first_position, mention.story_position) + if (len(tokens(mention.surface)), len(mention.surface)) > (len(tokens(info.display)), len(info.display)): + info.display = mention.surface + + for info in groups.values(): + info.display = display_name(info.display) + info.kind_hint = vote_kind(info.occurrences) + return sorted(groups.values(), key=lambda s: (s.first_position, s.is_role, s.display)) + + +def build_disambiguation_prompt(surface: SurfaceInfo, registry: Registry) -> str: + known = ( + "\n".join( + f"- {entity.name} [{entity.kind}]" + f"{' (provisional)' if entity.provisional else ''} - {entity.dossier or '(no dossier)'}" + f" aliases: {', '.join(alias for alias, _kind in entity.aliases)}" + for entity in registry.entities + ) + if registry.entities + else "(none yet)" + ) + occurrences = "\n".join( + f"- {mention.scene} ({mention.role}" + + (f", predicate {mention.predicate}" if mention.predicate else "") + + ")" + + (f' quote: "{mention.quote}"' if mention.quote else "") + for mention in surface.occurrences[:8] + ) + return ( + f"KNOWN ENTITIES:\n{known}\n\n" + f'REFERENCE TO RESOLVE: "{surface.display}"\n' + f"kind hint: {surface.kind_hint}\n" + f"occurrences:\n{occurrences}\n\n" + "Decide whether this reference is one of the KNOWN ENTITIES, a NEW entity, or UNSURE.\n" + "- existing: entity_name must be the exact canonical name from the known list.\n" + "- new: provide canonical_name, kind, and a one-line dossier.\n" + "- unsure: use when the evidence is not enough.\n" + "Return confidence 0..1 and a brief reason." + ) + + +def disambiguate( + client, + surface: SurfaceInfo, + registry: Registry, + *, + model: str = DEFAULT_MODEL, + effort: str = DEFAULT_EFFORT, + thinking: bool = True, +) -> dict[str, Any]: + request = structured_request( + model=model, + system=resolution_system_prompt(), + user=build_disambiguation_prompt(surface, registry), + schema=resolution_schema(), + max_tokens=MAX_RESOLUTION_TOKENS, + effort=effort, + thinking=thinking, + ) + response = client.messages.create(**request) + text = first_text(response) + if text is None: + raise RuntimeError(f"no text block resolving {surface.display!r}") + return json.loads(text) + + +@dataclass +class EntityQueueItem: + surface: str + kind_hint: str + reason: str + candidates: list[str] + entity_name: str + occurrences: list[dict[str, Any]] + + +@dataclass +class ResolutionState: + world: str + registry: Registry + assertions: list[dict[str, Any]] + scene_presence: list[dict[str, Any]] + entity_queue: list[EntityQueueItem] + assertion_queue: list[dict[str, Any]] = field(default_factory=list) + merges: list[dict[str, str]] = field(default_factory=list) + + +def _resolve_surface( + registry: Registry, + surface: SurfaceInfo, + client, + *, + model: str, + effort: str, + thinking: bool, + conf_auto: float, +) -> tuple[Entity, str | None, list[str]]: + status, found = registry.match(surface.display) + if status in ("exact", "fuzzy"): + entity = found + assert isinstance(entity, Entity) + registry.add_alias(entity, surface.display, "role_reference" if surface.is_role else "name_variant") + return entity, None, [] + + candidates = [entity.name for entity in found] if status == "ambiguous" and isinstance(found, list) else [] + needs_llm = status == "ambiguous" or (status == "none" and surface.is_role) + if needs_llm and client is not None: + decision = disambiguate(client, surface, registry, model=model, effort=effort, thinking=thinking) + confidence = float(decision.get("confidence") or 0.0) + if decision.get("decision") == "existing" and confidence >= conf_auto: + entity = registry.by_name(decision.get("entity_name")) + if entity is not None: + registry.add_alias(entity, surface.display, decision.get("alias_kind") or "role_reference") + return entity, None, [] + if decision.get("decision") == "new" and confidence >= conf_auto: + entity = registry.add( + decision.get("canonical_name") or surface.display, + decision.get("kind") or surface.kind_hint, + decision.get("dossier"), + provisional=False, + ) + return entity, None, [] + + if status == "ambiguous": + entity = registry.add(surface.display, surface.kind_hint, None, provisional=True) + return entity, "ambiguous", candidates + + if not surface.is_role: + entity = registry.add(surface.display, surface.kind_hint, None, provisional=False) + return entity, None, [] + + entity = registry.add(surface.display, surface.kind_hint, None, provisional=True, alias_kind="role_reference") + return entity, "role_reference", [] + + +def resolve_candidates( + candidates: dict[str, Any], + *, + world: str = "", + alias_table: dict[str, Any] | None = None, + client=None, + model: str = DEFAULT_MODEL, + effort: str = DEFAULT_EFFORT, + thinking: bool = True, + conf_auto: float = AUTO_ACCEPT_CONFIDENCE, +) -> ResolutionState: + if isinstance(candidates, dict) and not world: + world = candidates.get("world") or "" + registry = registry_from_alias_table(alias_table) + mentions = collect_mentions(candidates) + surfaces = group_surfaces(mentions) + + surface_to_entity: dict[str, Entity] = {} + entity_queue: list[EntityQueueItem] = [] + for surface in surfaces: + entity, reason, cand_names = _resolve_surface( + registry, + surface, + client, + model=model, + effort=effort, + thinking=thinking, + conf_auto=conf_auto, + ) + surface_to_entity[surface.norm] = entity + if reason: + entity_queue.append( + EntityQueueItem( + surface=surface.display, + kind_hint=surface.kind_hint, + reason=reason, + candidates=cand_names, + entity_name=entity.name, + occurrences=[ + { + "scene": mention.scene, + "role": mention.role, + "predicate": mention.predicate, + "quote": mention.quote, + } + for mention in surface.occurrences[:8] + ], + ) + ) + + presence = _build_scene_presence(candidates, surface_to_entity) + present_norms = {norm(name) for row in presence for name in row["entities"]} + for entity in registry.entities: + if entity.seeded: + continue + if entity.kind in ("character", "other") and norm(entity.name) not in present_norms: + entity.provisional = True + + assertions = _build_assertions(candidates, surface_to_entity) + return ResolutionState(world, registry, assertions, presence, entity_queue) + + +def _entity_for(surface_to_entity: dict[str, Entity], surface: Any) -> Entity | None: + if not surface or not str(surface).strip(): + return None + return surface_to_entity.get(norm(str(surface))) + + +def _build_assertions(candidates: dict[str, Any], surface_to_entity: dict[str, Entity]) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + ordinal = 1 + for scene in scenes_from_candidates(candidates): + label = scene_label(scene) + for assertion in scene.get("assertions") or []: + subject = _entity_for(surface_to_entity, assertion.get("subject")) + obj = _entity_for(surface_to_entity, assertion.get("object_entity")) if assertion.get("object_entity") else None + out.append( + { + "assertion_id": assertion.get("assertion_id") or f"A{ordinal}", + "subject": subject.name if subject else assertion.get("subject"), + "predicate": assertion.get("predicate"), + "object_entity": obj.name if obj else None, + "object_value": assertion.get("object_value") or assertion.get("object_fact_ref"), + "polarity": assertion.get("polarity", True), + "starts_here": assertion.get("starts_here", True), + "ends_here": assertion.get("ends_here", False), + "supporting_quote": assertion.get("supporting_quote"), + "confidence": assertion.get("confidence"), + "notes": assertion.get("notes"), + "scene": label, + "scene_id": scene.get("scene_id") or label, + "story_position": scene.get("story_position"), + "subject_provisional": bool(subject.provisional) if subject else True, + "object_provisional": bool(obj.provisional) if obj else False, + "status": assertion.get("status") or "candidate", + "confirmed_by_human": bool(assertion.get("confirmed_by_human", False)), + } + ) + ordinal += 1 + return out + + +def _build_scene_presence(candidates: dict[str, Any], surface_to_entity: dict[str, Entity]) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + for scene in scenes_from_candidates(candidates): + names: list[str] = [] + for name in scene.get("scene_presence") or []: + entity = _entity_for(surface_to_entity, name) + if entity and entity.name not in names: + names.append(entity.name) + out.append( + { + "scene": scene_label(scene), + "scene_id": scene.get("scene_id") or scene_label(scene), + "story_position": scene.get("story_position"), + "entities": names, + } + ) + return out + + +def repoint_assertions(state: ResolutionState, old_name: str, new_name: str) -> None: + old = norm(old_name) + for assertion in state.assertions: + if norm(assertion.get("subject")) == old: + assertion["subject"] = new_name + if assertion.get("object_entity") and norm(assertion["object_entity"]) == old: + assertion["object_entity"] = new_name + for row in state.scene_presence: + row["entities"] = [new_name if norm(name) == old else name for name in row["entities"]] + + +def merge_entities(state: ResolutionState, keep_name: str, drop_name: str, reason: str) -> bool: + keep = state.registry.by_name(keep_name) + drop = state.registry.by_name(drop_name) + if keep is None or drop is None or keep is drop: + return False + old = drop.name + state.registry.merge(keep, drop, reason) + repoint_assertions(state, old, keep.name) + for assertion in state.assertions: + if norm(assertion.get("subject")) == norm(keep.name): + assertion["subject_provisional"] = keep.provisional + if assertion.get("object_entity") and norm(assertion["object_entity"]) == norm(keep.name): + assertion["object_provisional"] = keep.provisional + state.entity_queue = [item for item in state.entity_queue if norm(item.entity_name) != norm(old)] + state.merges.append({"keep": keep.name, "drop": old, "reason": reason}) + return True + + +def apply_entity_decision(state: ResolutionState, item: EntityQueueItem, decision: str) -> bool: + """Apply one human entity decision. + + Grammar: ``, `= Existing Name`, `new Name | kind`, `keep`, or `skip`. + """ + + choice = (decision or "").strip() + if not choice or choice.lower() == "skip": + return False + if choice.lower() == "keep": + return True + if choice.isdigit(): + idx = int(choice) - 1 + if 0 <= idx < len(item.candidates): + return merge_entities(state, item.candidates[idx], item.entity_name, "confirm: picked candidate") + return False + if choice.startswith("="): + return merge_entities(state, choice[1:].strip(), item.entity_name, "confirm: assigned to existing") + if choice.lower().startswith("new "): + rest = choice[4:].strip() + name, _sep, kind = rest.partition("|") + name = name.strip() + kind = kind.strip() or "other" + entity = state.registry.by_name(item.entity_name) + if entity is None or not name: + return False + old = entity.name + entity.name = display_name(name) + entity.kind = kind if kind in ENTITY_KINDS else entity.kind + entity.provisional = False + state.registry.add_alias(entity, old, "name_variant", upgrade=False) + repoint_assertions(state, old, entity.name) + return True + return False + + +def walk_entity_queue(state: ResolutionState, prompt_fn) -> int: + remaining: list[EntityQueueItem] = [] + resolved = 0 + for item in state.entity_queue: + if apply_entity_decision(state, item, prompt_fn(item, state)): + resolved += 1 + else: + remaining.append(item) + state.entity_queue = remaining + return resolved + + +def state_to_dict(state: ResolutionState) -> dict[str, Any]: + return { + "world": state.world, + "entities": [ + { + "local_id": entity.local_id, + "name": entity.name, + "kind": entity.kind, + "dossier": entity.dossier, + "provisional": entity.provisional, + "seeded": entity.seeded, + "aliases": [{"alias": alias, "kind": kind} for alias, kind in entity.aliases], + } + for entity in state.registry.entities + ], + "assertions": state.assertions, + "scene_presence": state.scene_presence, + "entity_queue": [item.__dict__ for item in state.entity_queue], + "assertion_queue": state.assertion_queue, + "merges": state.merges, + } + + +def state_from_dict(data: dict[str, Any]) -> ResolutionState: + registry = Registry() + for row in data.get("entities") or []: + entity = Entity( + local_id=int(row["local_id"]), + name=row["name"], + kind=row.get("kind") or "other", + dossier=row.get("dossier"), + provisional=bool(row.get("provisional", False)), + seeded=bool(row.get("seeded", False)), + aliases=[[alias["alias"], alias["kind"]] for alias in row.get("aliases") or []], + ) + registry.entities.append(entity) + registry._next = max(registry._next, entity.local_id + 1) + return ResolutionState( + world=data.get("world") or "", + registry=registry, + assertions=data.get("assertions") or [], + scene_presence=data.get("scene_presence") or [], + entity_queue=[EntityQueueItem(**item) for item in data.get("entity_queue") or data.get("queue") or []], + assertion_queue=data.get("assertion_queue") or [], + merges=data.get("merges") or [], + ) + + +def to_resolved_assertions(state: ResolutionState, *, include_rejected: bool = False) -> dict[str, Any]: + assertions = [ + assertion + for assertion in state.assertions + if include_rejected or assertion.get("status") != "rejected" + ] + return {"world": state.world, "assertions": assertions, "scene_presence": state.scene_presence} + + +def render_summary(state: ResolutionState) -> str: + provisional_count = sum(1 for entity in state.registry.entities if entity.provisional) + lines = [ + f"world: {state.world}", + f"entities: {len(state.registry.entities)} ({provisional_count} provisional)", + f"resolved assertions: {len(state.assertions)}", + f"entity queue: {len(state.entity_queue)}", + f"assertion queue: {len(state.assertion_queue)}", + "", + ] + for entity in state.registry.entities: + tag = " [provisional]" if entity.provisional else "" + aliases = ", ".join(alias for alias, _kind in entity.aliases if norm(alias) != norm(entity.name)) + lines.append(f" {entity.name} [{entity.kind}]{tag}" + (f" (aka {aliases})" if aliases else "")) + if state.entity_queue: + lines.append("") + lines.append("entity queue:") + for item in state.entity_queue: + candidates = f" candidates: {', '.join(item.candidates)}" if item.candidates else "" + lines.append(f" - {item.surface!r} [{item.kind_hint}] - {item.reason}{candidates}") + if state.assertion_queue: + lines.append("") + lines.append("assertion queue:") + for item in state.assertion_queue[:10]: + lines.append( + f" - #{item['assertion_index']} {item['subject']} {item['predicate']} " + f"{item.get('object_value') or item.get('object_entity') or ''} " + f"conf={item.get('confidence')}" + ) + return "\n".join(lines) diff --git a/extract/schema.py b/extract/schema.py new file mode 100644 index 0000000..597c092 --- /dev/null +++ b/extract/schema.py @@ -0,0 +1,122 @@ +"""Closed extraction schemas from docs/extraction.md.""" + +from __future__ import annotations + +AUTO_ACCEPT_CONFIDENCE = 0.85 + +# Closed predicate vocabulary. Do not add predicates here without updating the +# product decision log and extraction spec. +PREDICATES = ( + "alive", + "dies", + "located_at", + "present_in_scene", + "member_of", + "possesses", + "married_to", + "parent_of", + "sibling_of", + "romantic_with", + "allied_with", + "enemy_of", + "knows", + "believes", + "secret_of", + "destroyed", + "created", + "occupation", + "trait", + "cannot", + "goal", + "promised", + "fact", +) + +INTRANSITIVE = frozenset({"alive", "dies", "destroyed"}) + +ENTITY_KINDS = ("character", "location", "object", "faction", "event", "rule", "other") +ALIAS_KINDS = ("name_variant", "nickname", "role_reference") + + +def _nullable_str() -> dict: + return {"anyOf": [{"type": "string"}, {"type": "null"}]} + + +def candidate_schema() -> dict: + """Anthropic structured-output JSON schema for one scene.""" + + assertion = { + "type": "object", + "additionalProperties": False, + "properties": { + "subject": {"type": "string"}, + "predicate": {"type": "string", "enum": list(PREDICATES)}, + "object_entity": _nullable_str(), + "object_value": _nullable_str(), + "object_fact_ref": _nullable_str(), + "polarity": {"type": "boolean"}, + "starts_here": {"type": "boolean"}, + "ends_here": {"type": "boolean"}, + "supporting_quote": {"type": "string"}, + "confidence": {"type": "number"}, + "notes": _nullable_str(), + }, + "required": [ + "subject", + "predicate", + "object_entity", + "object_value", + "object_fact_ref", + "polarity", + "starts_here", + "ends_here", + "supporting_quote", + "confidence", + "notes", + ], + } + str_array = {"type": "array", "items": {"type": "string"}} + return { + "type": "object", + "additionalProperties": False, + "properties": { + "assertions": {"type": "array", "items": assertion}, + "scene_presence": str_array, + "deaths": str_array, + "destructions": str_array, + "open_questions": str_array, + }, + "required": ["assertions", "scene_presence", "deaths", "destructions", "open_questions"], + } + + +def resolution_schema() -> dict: + """Structured-output schema for LLM entity disambiguation.""" + + def nullable(t: dict) -> dict: + return {"anyOf": [t, {"type": "null"}]} + + return { + "type": "object", + "additionalProperties": False, + "properties": { + "decision": {"type": "string", "enum": ["existing", "new", "unsure"]}, + "entity_name": nullable({"type": "string"}), + "canonical_name": nullable({"type": "string"}), + "kind": nullable({"type": "string", "enum": list(ENTITY_KINDS)}), + "alias_kind": nullable({"type": "string", "enum": list(ALIAS_KINDS)}), + "dossier": nullable({"type": "string"}), + "confidence": {"type": "number"}, + "reason": {"type": "string"}, + }, + "required": [ + "decision", + "entity_name", + "canonical_name", + "kind", + "alias_kind", + "dossier", + "confidence", + "reason", + ], + }