#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ IndexChapterMixin extracted from IndexManager. """ from __future__ import annotations import json from datetime import datetime from typing import Any, Dict, List, Optional class IndexChapterMixin: def add_chapter(self, meta: ChapterMeta): """添加/更新章节元数据""" with self._get_conn() as conn: cursor = conn.cursor() cursor.execute( """ INSERT OR REPLACE INTO chapters (chapter, title, location, word_count, characters, summary) VALUES (?, ?, ?, ?, ?, ?) """, ( meta.chapter, meta.title, meta.location, meta.word_count, json.dumps(meta.characters, ensure_ascii=False), meta.summary, ), ) conn.commit() def get_chapter(self, chapter: int) -> Optional[Dict]: """获取章节元数据""" with self._get_conn() as conn: cursor = conn.cursor() cursor.execute("SELECT * FROM chapters WHERE chapter = ?", (chapter,)) row = cursor.fetchone() if row: return self._row_to_dict(row, parse_json=["characters"]) return None def get_recent_chapters(self, limit: int = None) -> List[Dict]: """获取最近章节""" if limit is None: limit = self.config.query_recent_chapters_limit with self._get_conn() as conn: cursor = conn.cursor() cursor.execute( """ SELECT * FROM chapters ORDER BY chapter DESC LIMIT ? """, (limit,), ) return [ self._row_to_dict(row, parse_json=["characters"]) for row in cursor.fetchall() ] # ==================== 场景操作 ==================== def add_scenes(self, chapter: int, scenes: List[SceneMeta]): """添加章节场景""" with self._get_conn() as conn: cursor = conn.cursor() # 先删除该章节旧场景 cursor.execute("DELETE FROM scenes WHERE chapter = ?", (chapter,)) # 插入新场景 for scene in scenes: cursor.execute( """ INSERT INTO scenes (chapter, scene_index, start_line, end_line, location, summary, characters) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( scene.chapter, scene.scene_index, scene.start_line, scene.end_line, scene.location, scene.summary, json.dumps(scene.characters, ensure_ascii=False), ), ) conn.commit() def get_scenes(self, chapter: int) -> List[Dict]: """获取章节场景""" with self._get_conn() as conn: cursor = conn.cursor() cursor.execute( """ SELECT * FROM scenes WHERE chapter = ? ORDER BY scene_index """, (chapter,), ) return [ self._row_to_dict(row, parse_json=["characters"]) for row in cursor.fetchall() ] def search_scenes_by_location(self, location: str, limit: int = None) -> List[Dict]: """按地点搜索场景""" if limit is None: limit = self.config.query_scenes_by_location_limit with self._get_conn() as conn: cursor = conn.cursor() cursor.execute( """ SELECT * FROM scenes WHERE location LIKE ? ORDER BY chapter DESC LIMIT ? """, (f"%{location}%", limit), ) return [ self._row_to_dict(row, parse_json=["characters"]) for row in cursor.fetchall() ] # ==================== 出场记录操作 ==================== def record_appearance( self, entity_id: str, chapter: int, mentions: List[str], confidence: float = 1.0, skip_if_exists: bool = False, ): """记录实体出场 Args: entity_id: 实体ID chapter: 章节号 mentions: 提及列表 confidence: 置信度 skip_if_exists: 如果为True,当记录已存在时跳过(避免覆盖已有mentions) """ with self._get_conn() as conn: cursor = conn.cursor() if skip_if_exists: # 先检查是否已存在 cursor.execute( "SELECT 1 FROM appearances WHERE entity_id = ? AND chapter = ?", (entity_id, chapter), ) if cursor.fetchone(): return # 已存在,跳过 cursor.execute( """ INSERT OR REPLACE INTO appearances (entity_id, chapter, mentions, confidence) VALUES (?, ?, ?, ?) """, ( entity_id, chapter, json.dumps(mentions, ensure_ascii=False), confidence, ), ) conn.commit() def get_entity_appearances(self, entity_id: str, limit: int = None) -> List[Dict]: """获取实体出场记录""" if limit is None: limit = self.config.query_entity_appearances_limit with self._get_conn() as conn: cursor = conn.cursor() cursor.execute( """ SELECT * FROM appearances WHERE entity_id = ? ORDER BY chapter DESC LIMIT ? """, (entity_id, limit), ) return [ self._row_to_dict(row, parse_json=["mentions"]) for row in cursor.fetchall() ] def get_recent_appearances(self, limit: int = None) -> List[Dict]: """获取最近出场的实体""" if limit is None: limit = self.config.query_recent_appearances_limit with self._get_conn() as conn: cursor = conn.cursor() cursor.execute( """ SELECT entity_id, MAX(chapter) as last_chapter, COUNT(*) as total FROM appearances GROUP BY entity_id ORDER BY last_chapter DESC LIMIT ? """, (limit,), ) return [dict(row) for row in cursor.fetchall()] def get_chapter_appearances(self, chapter: int) -> List[Dict]: """获取某章所有出场实体""" with self._get_conn() as conn: cursor = conn.cursor() cursor.execute( """ SELECT * FROM appearances WHERE chapter = ? ORDER BY confidence DESC """, (chapter,), ) return [ self._row_to_dict(row, parse_json=["mentions"]) for row in cursor.fetchall() ] # ==================== v5.1 实体操作 ==================== def process_chapter_data( self, chapter: int, title: str, location: str, word_count: int, entities: List[Dict], scenes: List[Dict], ) -> Dict[str, int]: """ 处理章节数据,批量写入索引 返回写入统计 """ from .index_manager import ChapterMeta, SceneMeta stats = {"chapters": 0, "scenes": 0, "appearances": 0} # 提取出场角色 characters = [e.get("id") for e in entities if e.get("type") == "角色"] # 写入章节元数据 self.add_chapter( ChapterMeta( chapter=chapter, title=title, location=location, word_count=word_count, characters=characters, summary="", # 可后续由 Data Agent 生成 ) ) stats["chapters"] = 1 # 写入场景 scene_metas = [] for s in scenes: scene_metas.append( SceneMeta( chapter=chapter, scene_index=s.get("index", 0), start_line=s.get("start_line", 0), end_line=s.get("end_line", 0), location=s.get("location", ""), summary=s.get("summary", ""), characters=s.get("characters", []), ) ) self.add_scenes(chapter, scene_metas) stats["scenes"] = len(scene_metas) # 写入出场记录 for entity in entities: entity_id = entity.get("id") if entity_id and entity_id != "NEW": self.record_appearance( entity_id=entity_id, chapter=chapter, mentions=entity.get("mentions", []), confidence=entity.get("confidence", 1.0), ) stats["appearances"] += 1 return stats # ==================== 辅助方法 ====================