state_manager.py 54 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. """
  4. State Manager - 状态管理模块 (v5.4)
  5. 管理 state.json 的读写操作:
  6. - 实体状态管理
  7. - 进度追踪
  8. - 关系记录
  9. v5.1 变更(v5.4 沿用):
  10. - 集成 SQLStateManager,同步写入 SQLite (index.db)
  11. - state.json 保留精简数据,大数据自动迁移到 SQLite
  12. """
  13. import json
  14. import logging
  15. import sys
  16. from copy import deepcopy
  17. from pathlib import Path
  18. from runtime_compat import enable_windows_utf8_stdio
  19. from typing import Dict, List, Optional, Any
  20. from dataclasses import dataclass, field, asdict
  21. from datetime import datetime
  22. import filelock
  23. from .config import get_config
  24. from .observability import safe_log_tool_call
  25. logger = logging.getLogger(__name__)
  26. try:
  27. # 当 scripts 目录在 sys.path 中(常见:从 scripts/ 运行)
  28. from security_utils import atomic_write_json, read_json_safe
  29. except ImportError: # pragma: no cover
  30. # 当以 `python -m scripts.data_modules...` 从仓库根目录运行
  31. from scripts.security_utils import atomic_write_json, read_json_safe
  32. @dataclass
  33. class EntityState:
  34. """实体状态"""
  35. id: str
  36. name: str
  37. type: str # 角色/地点/物品/势力
  38. tier: str = "装饰" # 核心/重要/次要/装饰
  39. aliases: List[str] = field(default_factory=list)
  40. attributes: Dict[str, Any] = field(default_factory=dict)
  41. first_appearance: int = 0
  42. last_appearance: int = 0
  43. @dataclass
  44. class Relationship:
  45. """实体关系"""
  46. from_entity: str
  47. to_entity: str
  48. type: str
  49. description: str
  50. chapter: int
  51. @dataclass
  52. class StateChange:
  53. """状态变化记录"""
  54. entity_id: str
  55. field: str
  56. old_value: Any
  57. new_value: Any
  58. reason: str
  59. chapter: int
  60. timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
  61. @dataclass
  62. class _EntityPatch:
  63. """待写入的实体增量补丁(用于锁内合并)"""
  64. entity_type: str
  65. entity_id: str
  66. replace: bool = False
  67. base_entity: Optional[Dict[str, Any]] = None # 新建实体时的完整快照(用于填充缺失字段)
  68. top_updates: Dict[str, Any] = field(default_factory=dict)
  69. current_updates: Dict[str, Any] = field(default_factory=dict)
  70. appearance_chapter: Optional[int] = None
  71. class StateManager:
  72. """状态管理器(v5.1 entities_v3 格式 + SQLite 同步,v5.4 沿用)"""
  73. # v5.0 引入的实体类型
  74. ENTITY_TYPES = ["角色", "地点", "物品", "势力", "招式"]
  75. def __init__(self, config=None, enable_sqlite_sync: bool = True):
  76. """
  77. 初始化状态管理器
  78. 参数:
  79. - config: 配置对象
  80. - enable_sqlite_sync: 是否启用 SQLite 同步 (默认 True)
  81. """
  82. self.config = config or get_config()
  83. self._state: Dict[str, Any] = {}
  84. # 与 security_utils.atomic_write_json 保持一致:state.json.lock
  85. self._lock_path = self.config.state_file.with_suffix(self.config.state_file.suffix + ".lock")
  86. # v5.1 引入: SQLite 同步
  87. self._enable_sqlite_sync = enable_sqlite_sync
  88. self._sql_state_manager = None
  89. if enable_sqlite_sync:
  90. try:
  91. from .sql_state_manager import SQLStateManager
  92. self._sql_state_manager = SQLStateManager(self.config)
  93. except ImportError:
  94. pass # SQLStateManager 不可用时静默降级
  95. # 待写入的增量(锁内重读 + 合并 + 写入)
  96. self._pending_entity_patches: Dict[tuple[str, str], _EntityPatch] = {}
  97. self._pending_alias_entries: Dict[str, List[Dict[str, str]]] = {}
  98. self._pending_state_changes: List[Dict[str, Any]] = []
  99. self._pending_structured_relationships: List[Dict[str, Any]] = []
  100. self._pending_disambiguation_warnings: List[Dict[str, Any]] = []
  101. self._pending_disambiguation_pending: List[Dict[str, Any]] = []
  102. self._pending_progress_chapter: Optional[int] = None
  103. self._pending_progress_words_delta: int = 0
  104. self._pending_chapter_meta: Dict[str, Any] = {}
  105. # v5.1 引入: 缓存待同步到 SQLite 的数据
  106. self._pending_sqlite_data: Dict[str, Any] = {
  107. "entities_appeared": [],
  108. "entities_new": [],
  109. "state_changes": [],
  110. "relationships_new": [],
  111. "chapter": None
  112. }
  113. self._load_state()
  114. def _now_progress_timestamp(self) -> str:
  115. return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
  116. def _ensure_state_schema(self, state: Dict[str, Any]) -> Dict[str, Any]:
  117. """确保 state.json 具备运行所需的关键字段(尽量不破坏既有数据)。"""
  118. if not isinstance(state, dict):
  119. state = {}
  120. state.setdefault("project_info", {})
  121. state.setdefault("progress", {})
  122. state.setdefault("protagonist_state", {})
  123. # relationships: 旧版本可能是 list(实体关系),v5.0 运行态用 dict(人物关系/重要关系)
  124. relationships = state.get("relationships")
  125. if isinstance(relationships, list):
  126. state.setdefault("structured_relationships", [])
  127. if isinstance(state.get("structured_relationships"), list):
  128. state["structured_relationships"].extend(relationships)
  129. state["relationships"] = {}
  130. elif not isinstance(relationships, dict):
  131. state["relationships"] = {}
  132. state.setdefault("world_settings", {"power_system": [], "factions": [], "locations": []})
  133. state.setdefault("plot_threads", {"active_threads": [], "foreshadowing": []})
  134. state.setdefault("review_checkpoints", [])
  135. state.setdefault("chapter_meta", {})
  136. state.setdefault(
  137. "strand_tracker",
  138. {
  139. "last_quest_chapter": 0,
  140. "last_fire_chapter": 0,
  141. "last_constellation_chapter": 0,
  142. "current_dominant": "quest",
  143. "chapters_since_switch": 0,
  144. "history": [],
  145. },
  146. )
  147. entities_v3 = state.get("entities_v3")
  148. # v5.1 引入: entities_v3, alias_index, state_changes, structured_relationships 已迁移到 index.db
  149. # 不再在 state.json 中初始化或维护这些字段
  150. if not isinstance(state.get("disambiguation_warnings"), list):
  151. state["disambiguation_warnings"] = []
  152. if not isinstance(state.get("disambiguation_pending"), list):
  153. state["disambiguation_pending"] = []
  154. # progress 基础字段
  155. progress = state["progress"]
  156. if not isinstance(progress, dict):
  157. progress = {}
  158. state["progress"] = progress
  159. progress.setdefault("current_chapter", 0)
  160. progress.setdefault("total_words", 0)
  161. progress.setdefault("last_updated", self._now_progress_timestamp())
  162. return state
  163. def _load_state(self):
  164. """加载状态文件"""
  165. if self.config.state_file.exists():
  166. self._state = read_json_safe(self.config.state_file, default={})
  167. self._state = self._ensure_state_schema(self._state)
  168. else:
  169. self._state = self._ensure_state_schema({})
  170. def save_state(self):
  171. """
  172. 保存状态文件(锁内重读 + 合并 + 原子写入)。
  173. 解决多 Agent 并行下的“读-改-写覆盖”风险:
  174. - 获取锁
  175. - 重新读取磁盘最新 state.json
  176. - 仅合并本实例产生的增量(pending_*)
  177. - 原子化写入
  178. """
  179. # 无增量时不写入,避免无意义覆盖
  180. has_pending = any(
  181. [
  182. self._pending_entity_patches,
  183. self._pending_alias_entries,
  184. self._pending_state_changes,
  185. self._pending_structured_relationships,
  186. self._pending_disambiguation_warnings,
  187. self._pending_disambiguation_pending,
  188. self._pending_chapter_meta,
  189. self._pending_progress_chapter is not None,
  190. self._pending_progress_words_delta != 0,
  191. ]
  192. )
  193. if not has_pending:
  194. return
  195. self.config.ensure_dirs()
  196. lock = filelock.FileLock(str(self._lock_path), timeout=10)
  197. try:
  198. with lock:
  199. disk_state = read_json_safe(self.config.state_file, default={})
  200. disk_state = self._ensure_state_schema(disk_state)
  201. # progress(合并为 max(chapter) + words_delta 累加)
  202. if self._pending_progress_chapter is not None or self._pending_progress_words_delta != 0:
  203. progress = disk_state.get("progress", {})
  204. if not isinstance(progress, dict):
  205. progress = {}
  206. disk_state["progress"] = progress
  207. try:
  208. current_chapter = int(progress.get("current_chapter", 0) or 0)
  209. except (TypeError, ValueError):
  210. current_chapter = 0
  211. if self._pending_progress_chapter is not None:
  212. progress["current_chapter"] = max(current_chapter, int(self._pending_progress_chapter))
  213. if self._pending_progress_words_delta:
  214. try:
  215. total_words = int(progress.get("total_words", 0) or 0)
  216. except (TypeError, ValueError):
  217. total_words = 0
  218. progress["total_words"] = total_words + int(self._pending_progress_words_delta)
  219. progress["last_updated"] = self._now_progress_timestamp()
  220. # v5.1 引入: 强制使用 SQLite 模式,移除大数据字段
  221. # 确保 state.json 中不存在这些膨胀字段
  222. for field in ["entities_v3", "alias_index", "state_changes", "structured_relationships"]:
  223. disk_state.pop(field, None)
  224. # 标记已迁移
  225. disk_state["_migrated_to_sqlite"] = True
  226. # disambiguation_warnings(追加去重 + 截断)
  227. if self._pending_disambiguation_warnings:
  228. warnings_list = disk_state.get("disambiguation_warnings")
  229. if not isinstance(warnings_list, list):
  230. warnings_list = []
  231. disk_state["disambiguation_warnings"] = warnings_list
  232. def _warn_key(w: Dict[str, Any]) -> tuple:
  233. return (
  234. w.get("chapter"),
  235. w.get("mention"),
  236. w.get("chosen_id"),
  237. w.get("confidence"),
  238. )
  239. existing_keys = {_warn_key(w) for w in warnings_list if isinstance(w, dict)}
  240. for w in self._pending_disambiguation_warnings:
  241. if not isinstance(w, dict):
  242. continue
  243. k = _warn_key(w)
  244. if k in existing_keys:
  245. continue
  246. warnings_list.append(w)
  247. existing_keys.add(k)
  248. # 只保留最近 N 条,避免文件无限增长
  249. max_keep = self.config.max_disambiguation_warnings
  250. if len(warnings_list) > max_keep:
  251. disk_state["disambiguation_warnings"] = warnings_list[-max_keep:]
  252. # disambiguation_pending(追加去重 + 截断)
  253. if self._pending_disambiguation_pending:
  254. pending_list = disk_state.get("disambiguation_pending")
  255. if not isinstance(pending_list, list):
  256. pending_list = []
  257. disk_state["disambiguation_pending"] = pending_list
  258. def _pending_key(w: Dict[str, Any]) -> tuple:
  259. return (
  260. w.get("chapter"),
  261. w.get("mention"),
  262. w.get("suggested_id"),
  263. w.get("confidence"),
  264. )
  265. existing_keys = {_pending_key(w) for w in pending_list if isinstance(w, dict)}
  266. for w in self._pending_disambiguation_pending:
  267. if not isinstance(w, dict):
  268. continue
  269. k = _pending_key(w)
  270. if k in existing_keys:
  271. continue
  272. pending_list.append(w)
  273. existing_keys.add(k)
  274. max_keep = self.config.max_disambiguation_pending
  275. if len(pending_list) > max_keep:
  276. disk_state["disambiguation_pending"] = pending_list[-max_keep:]
  277. # chapter_meta(新增:按章节号覆盖写入)
  278. if self._pending_chapter_meta:
  279. chapter_meta = disk_state.get("chapter_meta")
  280. if not isinstance(chapter_meta, dict):
  281. chapter_meta = {}
  282. disk_state["chapter_meta"] = chapter_meta
  283. chapter_meta.update(self._pending_chapter_meta)
  284. # 原子写入(锁已持有,不再二次加锁)
  285. atomic_write_json(self.config.state_file, disk_state, use_lock=False, backup=True)
  286. # v5.1 引入: 同步到 SQLite(失败时保留 pending 以便重试)
  287. sqlite_pending_snapshot = self._snapshot_sqlite_pending()
  288. sqlite_sync_ok = self._sync_to_sqlite()
  289. # 同步内存为磁盘最新快照
  290. self._state = disk_state
  291. # state.json 侧 pending 已写盘,直接清空
  292. self._pending_disambiguation_warnings.clear()
  293. self._pending_disambiguation_pending.clear()
  294. self._pending_chapter_meta.clear()
  295. self._pending_progress_chapter = None
  296. self._pending_progress_words_delta = 0
  297. # SQLite 侧 pending:成功后清空,失败则恢复快照(避免静默丢数据)
  298. if sqlite_sync_ok:
  299. self._pending_entity_patches.clear()
  300. self._pending_alias_entries.clear()
  301. self._pending_state_changes.clear()
  302. self._pending_structured_relationships.clear()
  303. self._clear_pending_sqlite_data()
  304. else:
  305. self._restore_sqlite_pending(sqlite_pending_snapshot)
  306. except filelock.Timeout:
  307. raise RuntimeError("无法获取 state.json 文件锁,请稍后重试")
  308. def _sync_to_sqlite(self) -> bool:
  309. """同步待处理数据到 SQLite(v5.1 引入,v5.4 沿用)"""
  310. if not self._sql_state_manager:
  311. return True
  312. # 方式1: 通过 process_chapter_result 收集的数据
  313. sqlite_data = self._pending_sqlite_data
  314. chapter = sqlite_data.get("chapter")
  315. # 记录已处理的 (entity_id, chapter) 组合,避免重复写入 appearances
  316. processed_appearances = set()
  317. if chapter is not None:
  318. try:
  319. self._sql_state_manager.process_chapter_entities(
  320. chapter=chapter,
  321. entities_appeared=sqlite_data.get("entities_appeared", []),
  322. entities_new=sqlite_data.get("entities_new", []),
  323. state_changes=sqlite_data.get("state_changes", []),
  324. relationships_new=sqlite_data.get("relationships_new", [])
  325. )
  326. # 标记已处理的出场记录
  327. for entity in sqlite_data.get("entities_appeared", []):
  328. if entity.get("id"):
  329. processed_appearances.add((entity.get("id"), chapter))
  330. for entity in sqlite_data.get("entities_new", []):
  331. eid = entity.get("suggested_id") or entity.get("id")
  332. if eid:
  333. processed_appearances.add((eid, chapter))
  334. except Exception as exc:
  335. logger.warning("SQLite sync failed (process_chapter_entities): %s", exc)
  336. return False
  337. # 方式2: 使用 add_entity/update_entity 收集的增量数据。
  338. # 数据缓存在 _pending_entity_patches 等变量中。
  339. return self._sync_pending_patches_to_sqlite(processed_appearances)
  340. def _sync_pending_patches_to_sqlite(self, processed_appearances: set = None) -> bool:
  341. """同步 _pending_entity_patches 等到 SQLite(v5.1 引入,v5.4 沿用)
  342. Args:
  343. processed_appearances: 已通过 process_chapter_entities 处理的 (entity_id, chapter) 集合,
  344. 用于避免重复写入 appearances 表(防止覆盖 mentions)
  345. """
  346. if not self._sql_state_manager:
  347. return True
  348. if processed_appearances is None:
  349. processed_appearances = set()
  350. # 元数据字段(不应写入 current_json)
  351. METADATA_FIELDS = {"canonical_name", "tier", "desc", "is_protagonist", "is_archived"}
  352. try:
  353. from .sql_state_manager import EntityData
  354. from .index_manager import EntityMeta
  355. # 同步实体补丁
  356. for (entity_type, entity_id), patch in self._pending_entity_patches.items():
  357. if patch.base_entity:
  358. # 新实体
  359. entity_data = EntityData(
  360. id=entity_id,
  361. type=entity_type,
  362. name=patch.base_entity.get("canonical_name", entity_id),
  363. tier=patch.base_entity.get("tier", "装饰"),
  364. desc=patch.base_entity.get("desc", ""),
  365. current=patch.base_entity.get("current", {}),
  366. aliases=[],
  367. first_appearance=patch.base_entity.get("first_appearance", 0),
  368. last_appearance=patch.base_entity.get("last_appearance", 0),
  369. is_protagonist=patch.base_entity.get("is_protagonist", False)
  370. )
  371. self._sql_state_manager.upsert_entity(entity_data)
  372. # 记录首次出场(跳过已处理的,避免覆盖 mentions)
  373. if patch.appearance_chapter is not None:
  374. if (entity_id, patch.appearance_chapter) not in processed_appearances:
  375. self._sql_state_manager._index_manager.record_appearance(
  376. entity_id=entity_id,
  377. chapter=patch.appearance_chapter,
  378. mentions=[entity_data.name],
  379. confidence=1.0,
  380. skip_if_exists=True # 关键:不覆盖已有记录
  381. )
  382. else:
  383. # 更新现有实体
  384. has_metadata_updates = bool(patch.top_updates and
  385. any(k in METADATA_FIELDS for k in patch.top_updates))
  386. # 非元数据的 top_updates 应该当作 current 更新
  387. # 例如:realm, layer, location 等状态字段
  388. non_metadata_top_updates = {
  389. k: v for k, v in patch.top_updates.items()
  390. if k not in METADATA_FIELDS
  391. } if patch.top_updates else {}
  392. # 合并 current_updates 和非元数据的 top_updates
  393. effective_current_updates = {**non_metadata_top_updates}
  394. if patch.current_updates:
  395. effective_current_updates.update(patch.current_updates)
  396. if has_metadata_updates:
  397. # 有元数据更新:使用 upsert_entity(update_metadata=True)
  398. existing = self._sql_state_manager.get_entity(entity_id)
  399. if existing:
  400. # 合并 current
  401. current = existing.get("current_json", {})
  402. if isinstance(current, str):
  403. import json
  404. current = json.loads(current) if current else {}
  405. if effective_current_updates:
  406. current.update(effective_current_updates)
  407. new_canonical_name = patch.top_updates.get("canonical_name")
  408. old_canonical_name = existing.get("canonical_name", "")
  409. entity_meta = EntityMeta(
  410. id=entity_id,
  411. type=existing.get("type", entity_type),
  412. canonical_name=new_canonical_name or old_canonical_name,
  413. tier=patch.top_updates.get("tier", existing.get("tier", "装饰")),
  414. desc=patch.top_updates.get("desc", existing.get("desc", "")),
  415. current=current,
  416. first_appearance=existing.get("first_appearance", 0),
  417. last_appearance=patch.appearance_chapter or existing.get("last_appearance", 0),
  418. is_protagonist=patch.top_updates.get("is_protagonist", existing.get("is_protagonist", False)),
  419. is_archived=patch.top_updates.get("is_archived", existing.get("is_archived", False))
  420. )
  421. self._sql_state_manager._index_manager.upsert_entity(entity_meta, update_metadata=True)
  422. # 如果 canonical_name 改名,自动注册新名字为 alias
  423. if new_canonical_name and new_canonical_name != old_canonical_name:
  424. self._sql_state_manager.register_alias(
  425. new_canonical_name, entity_id, existing.get("type", entity_type)
  426. )
  427. elif effective_current_updates:
  428. # 只有 current 更新(包括非元数据的 top_updates)
  429. self._sql_state_manager.update_entity_current(entity_id, effective_current_updates)
  430. # 更新 last_appearance 并记录出场
  431. if patch.appearance_chapter is not None:
  432. self._sql_state_manager._update_last_appearance(entity_id, patch.appearance_chapter)
  433. # 补充 appearances 记录
  434. # 使用 skip_if_exists=True 避免覆盖已有记录的 mentions
  435. if (entity_id, patch.appearance_chapter) not in processed_appearances:
  436. self._sql_state_manager._index_manager.record_appearance(
  437. entity_id=entity_id,
  438. chapter=patch.appearance_chapter,
  439. mentions=[],
  440. confidence=1.0,
  441. skip_if_exists=True # 关键:不覆盖已有记录
  442. )
  443. # 同步别名
  444. for alias, entries in self._pending_alias_entries.items():
  445. for entry in entries:
  446. entity_type = entry.get("type")
  447. entity_id = entry.get("id")
  448. if entity_type and entity_id:
  449. self._sql_state_manager.register_alias(alias, entity_id, entity_type)
  450. # 同步状态变化
  451. for change in self._pending_state_changes:
  452. self._sql_state_manager.record_state_change(
  453. entity_id=change.get("entity_id", ""),
  454. field=change.get("field", ""),
  455. old_value=change.get("old", change.get("old_value", "")),
  456. new_value=change.get("new", change.get("new_value", "")),
  457. reason=change.get("reason", ""),
  458. chapter=change.get("chapter", 0)
  459. )
  460. # 同步关系
  461. for rel in self._pending_structured_relationships:
  462. self._sql_state_manager.upsert_relationship(
  463. from_entity=rel.get("from_entity", ""),
  464. to_entity=rel.get("to_entity", ""),
  465. type=rel.get("type", "相识"),
  466. description=rel.get("description", ""),
  467. chapter=rel.get("chapter", 0)
  468. )
  469. return True
  470. except Exception as e:
  471. # SQLite 同步失败时记录警告(不中断主流程)
  472. logger.warning("SQLite sync failed: %s", e)
  473. return False
  474. def _snapshot_sqlite_pending(self) -> Dict[str, Any]:
  475. """抓取 SQLite 侧 pending 快照,用于同步失败回滚内存队列。"""
  476. return {
  477. "entity_patches": deepcopy(self._pending_entity_patches),
  478. "alias_entries": deepcopy(self._pending_alias_entries),
  479. "state_changes": deepcopy(self._pending_state_changes),
  480. "structured_relationships": deepcopy(self._pending_structured_relationships),
  481. "sqlite_data": deepcopy(self._pending_sqlite_data),
  482. }
  483. def _restore_sqlite_pending(self, snapshot: Dict[str, Any]) -> None:
  484. """恢复 SQLite 侧 pending 快照,避免同步失败后数据静默丢失。"""
  485. self._pending_entity_patches = snapshot.get("entity_patches", {})
  486. self._pending_alias_entries = snapshot.get("alias_entries", {})
  487. self._pending_state_changes = snapshot.get("state_changes", [])
  488. self._pending_structured_relationships = snapshot.get("structured_relationships", [])
  489. self._pending_sqlite_data = snapshot.get("sqlite_data", {
  490. "entities_appeared": [],
  491. "entities_new": [],
  492. "state_changes": [],
  493. "relationships_new": [],
  494. "chapter": None,
  495. })
  496. def _clear_pending_sqlite_data(self):
  497. """清空待同步的 SQLite 数据"""
  498. self._pending_sqlite_data = {
  499. "entities_appeared": [],
  500. "entities_new": [],
  501. "state_changes": [],
  502. "relationships_new": [],
  503. "chapter": None
  504. }
  505. # ==================== 进度管理 ====================
  506. def get_current_chapter(self) -> int:
  507. """获取当前章节号"""
  508. return self._state.get("progress", {}).get("current_chapter", 0)
  509. def update_progress(self, chapter: int, words: int = 0):
  510. """更新进度"""
  511. if "progress" not in self._state:
  512. self._state["progress"] = {}
  513. self._state["progress"]["current_chapter"] = chapter
  514. if words > 0:
  515. total = self._state["progress"].get("total_words", 0)
  516. self._state["progress"]["total_words"] = total + words
  517. # 记录增量:锁内合并时用 max(chapter) + words_delta 累加
  518. if self._pending_progress_chapter is None:
  519. self._pending_progress_chapter = chapter
  520. else:
  521. self._pending_progress_chapter = max(self._pending_progress_chapter, chapter)
  522. if words > 0:
  523. self._pending_progress_words_delta += int(words)
  524. # ==================== 实体管理 (v5.1 SQLite-first) ====================
  525. def get_entity(self, entity_id: str, entity_type: str = None) -> Optional[Dict]:
  526. """获取实体(v5.1 引入:优先从 SQLite 读取)"""
  527. # v5.1 引入: 优先从 SQLite 读取
  528. if self._sql_state_manager:
  529. entity = self._sql_state_manager._index_manager.get_entity(entity_id)
  530. if entity:
  531. return entity
  532. # 回退到内存 state (兼容未迁移场景)
  533. entities_v3 = self._state.get("entities_v3", {})
  534. if entity_type:
  535. return entities_v3.get(entity_type, {}).get(entity_id)
  536. # 遍历所有类型查找
  537. for type_name, entities in entities_v3.items():
  538. if entity_id in entities:
  539. return entities[entity_id]
  540. return None
  541. def get_entity_type(self, entity_id: str) -> Optional[str]:
  542. """获取实体所属类型"""
  543. # v5.1 引入: 优先从 SQLite 读取
  544. if self._sql_state_manager:
  545. entity = self._sql_state_manager._index_manager.get_entity(entity_id)
  546. if entity:
  547. return entity.get("type")
  548. # 回退到内存 state
  549. for type_name, entities in self._state.get("entities_v3", {}).items():
  550. if entity_id in entities:
  551. return type_name
  552. return None
  553. def get_all_entities(self) -> Dict[str, Dict]:
  554. """获取所有实体(扁平化视图)"""
  555. # v5.1 引入: 优先从 SQLite 读取
  556. if self._sql_state_manager:
  557. result = {}
  558. for entity_type in self.ENTITY_TYPES:
  559. entities = self._sql_state_manager._index_manager.get_entities_by_type(entity_type)
  560. for e in entities:
  561. eid = e.get("id")
  562. if eid:
  563. result[eid] = {**e, "type": entity_type}
  564. if result:
  565. return result
  566. # 回退到内存 state
  567. result = {}
  568. for type_name, entities in self._state.get("entities_v3", {}).items():
  569. for eid, e in entities.items():
  570. result[eid] = {**e, "type": type_name}
  571. return result
  572. def get_entities_by_type(self, entity_type: str) -> Dict[str, Dict]:
  573. """按类型获取实体"""
  574. # v5.1 引入: 优先从 SQLite 读取
  575. if self._sql_state_manager:
  576. entities = self._sql_state_manager._index_manager.get_entities_by_type(entity_type)
  577. if entities:
  578. return {e.get("id"): e for e in entities if e.get("id")}
  579. # 回退到内存 state
  580. return self._state.get("entities_v3", {}).get(entity_type, {})
  581. def get_entities_by_tier(self, tier: str) -> Dict[str, Dict]:
  582. """按层级获取实体"""
  583. # v5.1 引入: 优先从 SQLite 读取
  584. if self._sql_state_manager:
  585. result = {}
  586. for entity_type in self.ENTITY_TYPES:
  587. entities = self._sql_state_manager._index_manager.get_entities_by_tier(tier)
  588. for e in entities:
  589. eid = e.get("id")
  590. if eid and e.get("type") == entity_type:
  591. result[eid] = {**e, "type": entity_type}
  592. if result:
  593. return result
  594. # 回退到内存 state
  595. result = {}
  596. for type_name, entities in self._state.get("entities_v3", {}).items():
  597. for eid, e in entities.items():
  598. if e.get("tier") == tier:
  599. result[eid] = {**e, "type": type_name}
  600. return result
  601. def add_entity(self, entity: EntityState) -> bool:
  602. """添加新实体(v5.0 entities_v3 格式,v5.4 沿用)"""
  603. entity_type = entity.type
  604. if entity_type not in self.ENTITY_TYPES:
  605. entity_type = "角色"
  606. if "entities_v3" not in self._state:
  607. self._state["entities_v3"] = {t: {} for t in self.ENTITY_TYPES}
  608. if entity_type not in self._state["entities_v3"]:
  609. self._state["entities_v3"][entity_type] = {}
  610. # 检查是否已存在
  611. if entity.id in self._state["entities_v3"][entity_type]:
  612. return False
  613. # 转换为 v3 格式
  614. v3_entity = {
  615. "canonical_name": entity.name,
  616. "tier": entity.tier,
  617. "desc": "",
  618. "current": entity.attributes,
  619. "first_appearance": entity.first_appearance,
  620. "last_appearance": entity.last_appearance,
  621. "history": []
  622. }
  623. self._state["entities_v3"][entity_type][entity.id] = v3_entity
  624. # 记录实体补丁(新建:仅填充缺失字段,避免覆盖并发写入)
  625. patch = self._pending_entity_patches.get((entity_type, entity.id))
  626. if patch is None:
  627. patch = _EntityPatch(entity_type=entity_type, entity_id=entity.id)
  628. self._pending_entity_patches[(entity_type, entity.id)] = patch
  629. patch.replace = True
  630. patch.base_entity = v3_entity
  631. # v5.1 引入: 注册别名到 index.db (通过 SQLStateManager)
  632. if self._sql_state_manager:
  633. self._sql_state_manager._index_manager.register_alias(entity.name, entity.id, entity_type)
  634. for alias in entity.aliases:
  635. if alias:
  636. self._sql_state_manager._index_manager.register_alias(alias, entity.id, entity_type)
  637. return True
  638. def _register_alias_internal(self, entity_id: str, entity_type: str, alias: str):
  639. """内部方法:注册别名到 index.db(v5.1 引入)"""
  640. if not alias:
  641. return
  642. # v5.1 引入: 直接写入 SQLite
  643. if self._sql_state_manager:
  644. self._sql_state_manager._index_manager.register_alias(alias, entity_id, entity_type)
  645. def update_entity(self, entity_id: str, updates: Dict[str, Any], entity_type: str = None) -> bool:
  646. """更新实体属性(v5.0 引入,v5.4 沿用)"""
  647. # 查找实体
  648. if entity_type:
  649. if entity_id not in self._state.get("entities_v3", {}).get(entity_type, {}):
  650. return False
  651. entity = self._state["entities_v3"][entity_type][entity_id]
  652. else:
  653. entity_type = self.get_entity_type(entity_id)
  654. if not entity_type:
  655. return False
  656. entity = self._state["entities_v3"][entity_type][entity_id]
  657. for key, value in updates.items():
  658. if key == "attributes" and isinstance(value, dict):
  659. # v5.0 引入: attributes 存在 current 字段
  660. if "current" not in entity:
  661. entity["current"] = {}
  662. entity["current"].update(value)
  663. # 记录补丁(current 增量)
  664. patch = self._pending_entity_patches.get((entity_type, entity_id))
  665. if patch is None:
  666. patch = _EntityPatch(entity_type=entity_type, entity_id=entity_id)
  667. self._pending_entity_patches[(entity_type, entity_id)] = patch
  668. patch.current_updates.update(value)
  669. elif key == "current" and isinstance(value, dict):
  670. if "current" not in entity:
  671. entity["current"] = {}
  672. entity["current"].update(value)
  673. patch = self._pending_entity_patches.get((entity_type, entity_id))
  674. if patch is None:
  675. patch = _EntityPatch(entity_type=entity_type, entity_id=entity_id)
  676. self._pending_entity_patches[(entity_type, entity_id)] = patch
  677. patch.current_updates.update(value)
  678. else:
  679. entity[key] = value
  680. patch = self._pending_entity_patches.get((entity_type, entity_id))
  681. if patch is None:
  682. patch = _EntityPatch(entity_type=entity_type, entity_id=entity_id)
  683. self._pending_entity_patches[(entity_type, entity_id)] = patch
  684. patch.top_updates[key] = value
  685. return True
  686. def update_entity_appearance(self, entity_id: str, chapter: int, entity_type: str = None):
  687. """更新实体出场章节"""
  688. if not entity_type:
  689. entity_type = self.get_entity_type(entity_id)
  690. if not entity_type:
  691. return
  692. entities_v3 = self._state.get("entities_v3")
  693. if not isinstance(entities_v3, dict):
  694. entities_v3 = {t: {} for t in self.ENTITY_TYPES}
  695. self._state["entities_v3"] = entities_v3
  696. entities_v3.setdefault(entity_type, {})
  697. entity = entities_v3[entity_type].get(entity_id)
  698. if entity:
  699. if entity.get("first_appearance", 0) == 0:
  700. entity["first_appearance"] = chapter
  701. entity["last_appearance"] = chapter
  702. # 记录补丁:锁内应用 first=min(non-zero), last=max
  703. patch = self._pending_entity_patches.get((entity_type, entity_id))
  704. if patch is None:
  705. patch = _EntityPatch(entity_type=entity_type, entity_id=entity_id)
  706. self._pending_entity_patches[(entity_type, entity_id)] = patch
  707. if patch.appearance_chapter is None:
  708. patch.appearance_chapter = chapter
  709. else:
  710. patch.appearance_chapter = max(int(patch.appearance_chapter), int(chapter))
  711. # ==================== 状态变化记录 ====================
  712. def record_state_change(
  713. self,
  714. entity_id: str,
  715. field: str,
  716. old_value: Any,
  717. new_value: Any,
  718. reason: str,
  719. chapter: int
  720. ):
  721. """记录状态变化"""
  722. if "state_changes" not in self._state:
  723. self._state["state_changes"] = []
  724. change = StateChange(
  725. entity_id=entity_id,
  726. field=field,
  727. old_value=old_value,
  728. new_value=new_value,
  729. reason=reason,
  730. chapter=chapter
  731. )
  732. change_dict = asdict(change)
  733. self._state["state_changes"].append(change_dict)
  734. self._pending_state_changes.append(change_dict)
  735. # 同时更新实体属性
  736. self.update_entity(entity_id, {"attributes": {field: new_value}})
  737. def get_state_changes(self, entity_id: Optional[str] = None) -> List[Dict]:
  738. """获取状态变化历史"""
  739. changes = self._state.get("state_changes", [])
  740. if entity_id:
  741. changes = [c for c in changes if c.get("entity_id") == entity_id]
  742. return changes
  743. # ==================== 关系管理 ====================
  744. def add_relationship(
  745. self,
  746. from_entity: str,
  747. to_entity: str,
  748. rel_type: str,
  749. description: str,
  750. chapter: int
  751. ):
  752. """添加关系"""
  753. rel = Relationship(
  754. from_entity=from_entity,
  755. to_entity=to_entity,
  756. type=rel_type,
  757. description=description,
  758. chapter=chapter
  759. )
  760. # v5.0 引入: 实体关系存入 structured_relationships,避免与 relationships(人物关系字典) 冲突
  761. if "structured_relationships" not in self._state:
  762. self._state["structured_relationships"] = []
  763. rel_dict = asdict(rel)
  764. self._state["structured_relationships"].append(rel_dict)
  765. self._pending_structured_relationships.append(rel_dict)
  766. def get_relationships(self, entity_id: Optional[str] = None) -> List[Dict]:
  767. """获取关系列表"""
  768. rels = self._state.get("structured_relationships", [])
  769. if entity_id:
  770. rels = [
  771. r for r in rels
  772. if r.get("from_entity") == entity_id or r.get("to_entity") == entity_id
  773. ]
  774. return rels
  775. # ==================== 批量操作 ====================
  776. def _record_disambiguation(self, chapter: int, uncertain_items: Any) -> List[str]:
  777. """
  778. 记录消歧反馈到 state.json,便于 Writer/Context Agent 感知风险。
  779. 约定:
  780. - >= extraction_confidence_medium:写入 disambiguation_warnings(采用但警告)
  781. - < extraction_confidence_medium:写入 disambiguation_pending(需人工确认)
  782. """
  783. if not isinstance(uncertain_items, list) or not uncertain_items:
  784. return []
  785. warnings: List[str] = []
  786. now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
  787. for item in uncertain_items:
  788. if not isinstance(item, dict):
  789. continue
  790. mention = str(item.get("mention", "") or "").strip()
  791. if not mention:
  792. continue
  793. raw_conf = item.get("confidence", 0.0)
  794. try:
  795. confidence = float(raw_conf)
  796. except (TypeError, ValueError):
  797. confidence = 0.0
  798. # 候选:支持 [{"type","id"}...] 或 ["id1","id2"] 两种形式
  799. candidates_raw = item.get("candidates", [])
  800. candidates: List[Dict[str, str]] = []
  801. if isinstance(candidates_raw, list):
  802. for c in candidates_raw:
  803. if isinstance(c, dict):
  804. cid = str(c.get("id", "") or "").strip()
  805. ctype = str(c.get("type", "") or "").strip()
  806. entry: Dict[str, str] = {}
  807. if ctype:
  808. entry["type"] = ctype
  809. if cid:
  810. entry["id"] = cid
  811. if entry:
  812. candidates.append(entry)
  813. else:
  814. cid = str(c).strip()
  815. if cid:
  816. candidates.append({"id": cid})
  817. entity_type = str(item.get("type", "") or "").strip()
  818. suggested_id = str(item.get("suggested", "") or "").strip()
  819. adopted_raw = item.get("adopted", None)
  820. chosen_id = ""
  821. if isinstance(adopted_raw, str):
  822. chosen_id = adopted_raw.strip()
  823. elif adopted_raw is True:
  824. chosen_id = suggested_id
  825. else:
  826. # 兼容字段名:entity_id / chosen_id
  827. chosen_id = str(item.get("entity_id") or item.get("chosen_id") or "").strip() or suggested_id
  828. context = str(item.get("context", "") or "").strip()
  829. note = str(item.get("warning", "") or "").strip()
  830. record: Dict[str, Any] = {
  831. "chapter": int(chapter),
  832. "mention": mention,
  833. "type": entity_type,
  834. "suggested_id": suggested_id,
  835. "chosen_id": chosen_id,
  836. "confidence": confidence,
  837. "candidates": candidates,
  838. "context": context,
  839. "note": note,
  840. "created_at": now,
  841. }
  842. if confidence >= float(self.config.extraction_confidence_medium):
  843. self._state.setdefault("disambiguation_warnings", []).append(record)
  844. self._pending_disambiguation_warnings.append(record)
  845. chosen_part = f" → {chosen_id}" if chosen_id else ""
  846. warnings.append(f"消歧警告: {mention}{chosen_part} (confidence: {confidence:.2f})")
  847. else:
  848. self._state.setdefault("disambiguation_pending", []).append(record)
  849. self._pending_disambiguation_pending.append(record)
  850. warnings.append(f"消歧需人工确认: {mention} (confidence: {confidence:.2f})")
  851. return warnings
  852. def process_chapter_result(self, chapter: int, result: Dict) -> List[str]:
  853. """
  854. 处理 Data Agent 的章节处理结果(v5.1 引入,v5.4 沿用)
  855. 输入格式:
  856. - entities_appeared: 出场实体列表
  857. - entities_new: 新实体列表
  858. - state_changes: 状态变化列表
  859. - relationships_new: 新关系列表
  860. 返回警告列表
  861. """
  862. warnings = []
  863. # v5.1 引入: 记录章节号用于 SQLite 同步
  864. self._pending_sqlite_data["chapter"] = chapter
  865. # 处理出场实体
  866. for entity in result.get("entities_appeared", []):
  867. entity_id = entity.get("id")
  868. entity_type = entity.get("type")
  869. if entity_id:
  870. self.update_entity_appearance(entity_id, chapter, entity_type)
  871. # v5.1 引入: 缓存用于 SQLite 同步
  872. self._pending_sqlite_data["entities_appeared"].append(entity)
  873. # 处理新实体
  874. for entity in result.get("entities_new", []):
  875. entity_id = entity.get("suggested_id") or entity.get("id")
  876. if entity_id and entity_id != "NEW":
  877. new_entity = EntityState(
  878. id=entity_id,
  879. name=entity.get("name", ""),
  880. type=entity.get("type", "角色"),
  881. tier=entity.get("tier", "装饰"),
  882. aliases=entity.get("mentions", []),
  883. first_appearance=chapter,
  884. last_appearance=chapter
  885. )
  886. if not self.add_entity(new_entity):
  887. warnings.append(f"实体已存在: {entity_id}")
  888. # v5.1 引入: 缓存用于 SQLite 同步
  889. self._pending_sqlite_data["entities_new"].append(entity)
  890. # 处理状态变化
  891. for change in result.get("state_changes", []):
  892. self.record_state_change(
  893. entity_id=change.get("entity_id", ""),
  894. field=change.get("field", ""),
  895. old_value=change.get("old"),
  896. new_value=change.get("new"),
  897. reason=change.get("reason", ""),
  898. chapter=chapter
  899. )
  900. # v5.1 引入: 缓存用于 SQLite 同步
  901. self._pending_sqlite_data["state_changes"].append(change)
  902. # 处理关系
  903. for rel in result.get("relationships_new", []):
  904. self.add_relationship(
  905. from_entity=rel.get("from", ""),
  906. to_entity=rel.get("to", ""),
  907. rel_type=rel.get("type", ""),
  908. description=rel.get("description", ""),
  909. chapter=chapter
  910. )
  911. # v5.1 引入: 缓存用于 SQLite 同步
  912. self._pending_sqlite_data["relationships_new"].append(rel)
  913. # 处理消歧不确定项(不影响实体写入,但必须对 Writer 可见)
  914. warnings.extend(self._record_disambiguation(chapter, result.get("uncertain", [])))
  915. # 写入 chapter_meta(钩子/模式/结束状态)
  916. chapter_meta = result.get("chapter_meta")
  917. if isinstance(chapter_meta, dict):
  918. meta_key = f"{int(chapter):04d}"
  919. self._state.setdefault("chapter_meta", {})
  920. self._state["chapter_meta"][meta_key] = chapter_meta
  921. self._pending_chapter_meta[meta_key] = chapter_meta
  922. # 更新进度
  923. self.update_progress(chapter)
  924. # 同步主角状态(entities_v3 → protagonist_state)
  925. self.sync_protagonist_from_entity()
  926. return warnings
  927. # ==================== 导出 ====================
  928. def export_for_context(self) -> Dict:
  929. """导出用于上下文的精简版状态(v5.0 引入,v5.4 沿用)"""
  930. # 从 entities_v3 构建精简视图
  931. entities_flat = {}
  932. for type_name, entities in self._state.get("entities_v3", {}).items():
  933. for eid, e in entities.items():
  934. entities_flat[eid] = {
  935. "name": e.get("canonical_name", eid),
  936. "type": type_name,
  937. "tier": e.get("tier", "装饰"),
  938. "current": e.get("current", {})
  939. }
  940. return {
  941. "progress": self._state.get("progress", {}),
  942. "entities": entities_flat,
  943. # v5.1 引入: alias_index 已迁移到 index.db,这里返回空(兼容性)
  944. "alias_index": {},
  945. "recent_changes": [], # v5.1 引入: 从 index.db 查询
  946. "disambiguation": {
  947. "warnings": self._state.get("disambiguation_warnings", [])[-self.config.export_disambiguation_slice:],
  948. "pending": self._state.get("disambiguation_pending", [])[-self.config.export_disambiguation_slice:],
  949. },
  950. }
  951. # ==================== 主角同步 ====================
  952. def get_protagonist_entity_id(self) -> Optional[str]:
  953. """获取主角实体 ID(通过 is_protagonist 标记或 SQLite 查询)"""
  954. # 方式1: 通过 SQLStateManager 查询 (v5.1)
  955. if self._sql_state_manager:
  956. protagonist = self._sql_state_manager.get_protagonist()
  957. if protagonist:
  958. return protagonist.get("id")
  959. # 方式2: 通过 protagonist_state.name 查找别名
  960. protag_name = self._state.get("protagonist_state", {}).get("name")
  961. if protag_name and self._sql_state_manager:
  962. entities = self._sql_state_manager._index_manager.get_entities_by_alias(protag_name)
  963. for entry in entities:
  964. if entry.get("type") == "角色":
  965. return entry.get("id")
  966. return None
  967. def sync_protagonist_from_entity(self, entity_id: str = None):
  968. """
  969. 将主角实体的状态同步到 protagonist_state (v5.1: 从 SQLite 读取)
  970. 用于确保 consistency-checker 等依赖 protagonist_state 的组件获取最新数据
  971. """
  972. if entity_id is None:
  973. entity_id = self.get_protagonist_entity_id()
  974. if entity_id is None:
  975. return
  976. entity = self.get_entity(entity_id, "角色")
  977. if not entity:
  978. return
  979. current = entity.get("current")
  980. if not isinstance(current, dict):
  981. current = entity.get("current_json", {})
  982. if isinstance(current, str):
  983. try:
  984. current = json.loads(current) if current else {}
  985. except (json.JSONDecodeError, TypeError):
  986. current = {}
  987. if not isinstance(current, dict):
  988. current = {}
  989. protag = self._state.setdefault("protagonist_state", {})
  990. # 同步境界
  991. if "realm" in current:
  992. power = protag.setdefault("power", {})
  993. power["realm"] = current["realm"]
  994. if "layer" in current:
  995. power["layer"] = current["layer"]
  996. # 同步位置
  997. if "location" in current:
  998. loc = protag.setdefault("location", {})
  999. loc["current"] = current["location"]
  1000. if "last_chapter" in current:
  1001. loc["last_chapter"] = current["last_chapter"]
  1002. def sync_protagonist_to_entity(self, entity_id: str = None):
  1003. """
  1004. 将 protagonist_state 同步到 entities_v3 中的主角实体
  1005. 用于初始化或手动编辑 protagonist_state 后保持一致性
  1006. """
  1007. if entity_id is None:
  1008. entity_id = self.get_protagonist_entity_id()
  1009. if entity_id is None:
  1010. return
  1011. protag = self._state.get("protagonist_state", {})
  1012. if not protag:
  1013. return
  1014. updates = {}
  1015. # 同步境界
  1016. power = protag.get("power", {})
  1017. if power.get("realm"):
  1018. updates["realm"] = power["realm"]
  1019. if power.get("layer"):
  1020. updates["layer"] = power["layer"]
  1021. # 同步位置
  1022. loc = protag.get("location", {})
  1023. if loc.get("current"):
  1024. updates["location"] = loc["current"]
  1025. if updates:
  1026. self.update_entity(entity_id, updates, "角色")
  1027. # ==================== CLI 接口 ====================
  1028. def main():
  1029. import argparse
  1030. from pydantic import ValidationError
  1031. from .cli_output import print_success, print_error
  1032. from .schemas import validate_data_agent_output, format_validation_error, normalize_data_agent_output
  1033. from .index_manager import IndexManager
  1034. parser = argparse.ArgumentParser(description="State Manager CLI (v5.4)")
  1035. parser.add_argument("--project-root", type=str, help="项目根目录")
  1036. subparsers = parser.add_subparsers(dest="command")
  1037. # 读取进度
  1038. subparsers.add_parser("get-progress")
  1039. # 获取实体
  1040. get_entity_parser = subparsers.add_parser("get-entity")
  1041. get_entity_parser.add_argument("--id", required=True)
  1042. # 列出实体
  1043. list_parser = subparsers.add_parser("list-entities")
  1044. list_parser.add_argument("--type", help="按类型过滤")
  1045. list_parser.add_argument("--tier", help="按层级过滤")
  1046. # 处理章节结果
  1047. process_parser = subparsers.add_parser("process-chapter")
  1048. process_parser.add_argument("--chapter", type=int, required=True, help="章节号")
  1049. process_parser.add_argument("--data", required=True, help="JSON 格式的处理结果")
  1050. args = parser.parse_args()
  1051. # 初始化
  1052. config = None
  1053. if args.project_root:
  1054. from .config import DataModulesConfig
  1055. config = DataModulesConfig.from_project_root(args.project_root)
  1056. manager = StateManager(config)
  1057. logger = IndexManager(config)
  1058. tool_name = f"state_manager:{args.command or 'unknown'}"
  1059. def emit_success(data=None, message: str = "ok"):
  1060. print_success(data, message=message)
  1061. safe_log_tool_call(logger, tool_name=tool_name, success=True)
  1062. def emit_error(code: str, message: str, suggestion: str | None = None):
  1063. print_error(code, message, suggestion=suggestion)
  1064. safe_log_tool_call(
  1065. logger,
  1066. tool_name=tool_name,
  1067. success=False,
  1068. error_code=code,
  1069. error_message=message,
  1070. )
  1071. if args.command == "get-progress":
  1072. emit_success(manager._state.get("progress", {}), message="progress")
  1073. elif args.command == "get-entity":
  1074. entity = manager.get_entity(args.id)
  1075. if entity:
  1076. emit_success(entity, message="entity")
  1077. else:
  1078. emit_error("NOT_FOUND", f"未找到实体: {args.id}")
  1079. elif args.command == "list-entities":
  1080. if args.type:
  1081. entities = manager.get_entities_by_type(args.type)
  1082. elif args.tier:
  1083. entities = manager.get_entities_by_tier(args.tier)
  1084. else:
  1085. entities = manager.get_all_entities()
  1086. payload = [{"id": eid, **e} for eid, e in entities.items()]
  1087. emit_success(payload, message="entities")
  1088. elif args.command == "process-chapter":
  1089. data = json.loads(args.data)
  1090. validated = None
  1091. last_exc = None
  1092. for _ in range(3):
  1093. try:
  1094. validated = validate_data_agent_output(data)
  1095. break
  1096. except ValidationError as exc:
  1097. last_exc = exc
  1098. data = normalize_data_agent_output(data)
  1099. if validated is None:
  1100. err = format_validation_error(last_exc) if last_exc else {
  1101. "code": "SCHEMA_VALIDATION_FAILED",
  1102. "message": "数据结构校验失败",
  1103. "details": {"errors": []},
  1104. "suggestion": "请检查 data-agent 输出字段是否完整且类型正确",
  1105. }
  1106. emit_error(err["code"], err["message"], suggestion=err.get("suggestion"))
  1107. return
  1108. warnings = manager.process_chapter_result(args.chapter, validated.model_dump(by_alias=True))
  1109. manager.save_state()
  1110. emit_success({"chapter": args.chapter, "warnings": warnings}, message="chapter_processed")
  1111. else:
  1112. emit_error("UNKNOWN_COMMAND", "未指定有效命令", suggestion="请查看 --help")
  1113. if __name__ == "__main__":
  1114. if sys.platform == "win32":
  1115. enable_windows_utf8_stdio()
  1116. main()