"""엔티티 추출 모듈 — LLM 기반 캐릭터/장소/소품 추출.""" # Version: 1.2.0 — stricter extraction (v7 prompts), fewer minor entities/relations # prompt_dependency: entity_extraction/v7 # updated_at: 2026-03-16 import json from pathlib import Path from typing import Any, Dict, List, Optional from app.modules.llm.base import BaseLLMClient PROMPTS_BASE = Path(__file__).resolve().parent.parent.parent.parent / "prompts" / "_base" ENTITY_PROMPT_DIR = PROMPTS_BASE / "entity_extraction" / "v7" # Relations that affect how things LOOK — the only families extracted for T2I. VISUAL_RELATION_FAMILIES = [ "identity", # 동일인 (시각: 같은 얼굴) "transformation", # 변형 (시각: 외모 변화 — 나이, 변장, 부상) "possession", # 소유/착용 (시각: 인물이 소품을 들고/입고 있음) "containment", # 포함 (시각: 공간 내 위치) ] SUPPORTED_LANGUAGE_MAP = { "ko": "Korean", "ja": "Japanese", "en": "English", } # ── JSON Schema for entity extraction response ────────────────────── RELATION_FACT_SCHEMA: Dict[str, object] = { "type": "object", "additionalProperties": False, "properties": { "relation_family": { "type": "string", "enum": VISUAL_RELATION_FAMILIES, }, "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", ], } _ENTITY_ITEM_SCHEMA: Dict[str, object] = { "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", ], } _CHARACTER_SCHEMA = _ENTITY_ITEM_SCHEMA _LOCATION_SCHEMA: Dict[str, object] = { "type": "object", "additionalProperties": False, "properties": { **_ENTITY_ITEM_SCHEMA["properties"], "kind": { "type": "string", "enum": ["interior", "exterior", "mixed", "unknown"], }, }, "required": [*_ENTITY_ITEM_SCHEMA["required"], "kind"], } _PROP_SCHEMA: Dict[str, object] = { "type": "object", "additionalProperties": False, "properties": { **{k: v for k, v in _ENTITY_ITEM_SCHEMA["properties"].items() if k != "importance"}, "significance": { "type": "string", "enum": ["key", "recurring", "minor", "unknown"], }, }, "required": [ r if r != "importance" else "significance" for r in _ENTITY_ITEM_SCHEMA["required"] ], } EXTRACTION_SCHEMA: Dict[str, object] = { "type": "object", "additionalProperties": False, "properties": { "characters": {"type": "array", "items": _CHARACTER_SCHEMA}, "locations": {"type": "array", "items": _LOCATION_SCHEMA}, "props": {"type": "array", "items": _PROP_SCHEMA}, "relation_facts": {"type": "array", "items": RELATION_FACT_SCHEMA}, "notes": {"type": "array", "items": {"type": "string"}}, }, "required": [ "characters", "locations", "props", "relation_facts", "notes", ], } def _load_prompt(filename: str) -> str: path = ENTITY_PROMPT_DIR / filename return path.read_text(encoding="utf-8").strip() def _build_entities_list(raw: Dict[str, Any]) -> List[Dict[str, Any]]: """Flatten characters/locations/props into a unified entity list.""" entities: List[Dict[str, Any]] = [] for char in raw.get("characters", []): entities.append({ "entity_type": "character", "name": char["name"], "aliases": char.get("aliases", []), "description": char.get("description", ""), "stable_traits": json.dumps({ "visual_anchor_traits": char.get("visual_anchor_traits", []), "variant_axes": char.get("variant_axes", []), "importance": char.get("importance", "unknown"), "continuity_priority": char.get("continuity_priority", "medium"), "reference_image_priority": char.get("reference_image_priority", "helpful"), }, ensure_ascii=False), }) for loc in raw.get("locations", []): entities.append({ "entity_type": "location", "name": loc["name"], "aliases": loc.get("aliases", []), "description": loc.get("description", ""), "stable_traits": json.dumps({ "visual_anchor_traits": loc.get("visual_anchor_traits", []), "variant_axes": loc.get("variant_axes", []), "kind": loc.get("kind", "unknown"), "importance": loc.get("importance", "unknown"), "continuity_priority": loc.get("continuity_priority", "medium"), "reference_image_priority": loc.get("reference_image_priority", "helpful"), }, ensure_ascii=False), }) for prop in raw.get("props", []): entities.append({ "entity_type": "prop", "name": prop["name"], "aliases": prop.get("aliases", []), "description": prop.get("description", ""), "stable_traits": json.dumps({ "visual_anchor_traits": prop.get("visual_anchor_traits", []), "variant_axes": prop.get("variant_axes", []), "significance": prop.get("significance", "unknown"), "continuity_priority": prop.get("continuity_priority", "medium"), "reference_image_priority": prop.get("reference_image_priority", "helpful"), }, ensure_ascii=False), }) return entities def _build_relations_list(raw: Dict[str, Any]) -> List[Dict[str, Any]]: """Extract relation facts from raw LLM output.""" relations: List[Dict[str, Any]] = [] for rf in raw.get("relation_facts", []): relations.append({ "relation_family": rf["relation_family"], "relation_type": rf["relation_type"], "directionality": rf["directionality"], "temporal_scope": rf["temporal_scope"], "continuity_priority": rf["continuity_priority"], "continuity_reason": rf.get("continuity_reason", ""), "participants": [ { "entity_name": p["entity_name"], "role": p["role"], "order": idx + 1, } for idx, p in enumerate(rf.get("participants", [])) ], }) return relations def extract_entities( llm_client: BaseLLMClient, fulltext: str, language: str = "ko", source_file: str = "episode", prior_memory: Optional[str] = None, ) -> Dict[str, Any]: """Run entity extraction on screenplay fulltext. Returns: {"entities": [...], "relations": [...]} """ lang_name = SUPPORTED_LANGUAGE_MAP.get(language, "Korean") # Use chunk prompts for single-pass fulltext extraction # (final prompts are for consolidating chunked results) system_prompt = _load_prompt("chunk_system.md").format( source_language_name=lang_name, source_language_code=language, ) memory_block = prior_memory or "(없음 / None)" user_prompt = _load_prompt("chunk_user.md").format( page_start=1, page_end="end", source_language_name=lang_name, source_language_code=language, series_memory_block=memory_block, screenplay_text=fulltext, ) raw = llm_client.generate_structured( system_prompt=system_prompt, user_prompt=user_prompt, response_schema=EXTRACTION_SCHEMA, schema_name="entity_extraction", max_tokens=16000, ) return { "entities": _build_entities_list(raw), "relations": _build_relations_list(raw), }