"""Chapter generation service: context assembly + streaming chapter writing. Pure functions — no Flask imports. Context assembly follows the 拆书→写书 pipeline: recent chapters full-text, earlier chapters as short summaries, relevant characters, open foreshadowing, novel meta. Total prompt size is bounded by ``config.CTX_BUDGET_MAX``. """ import json import logging import config from repositories import chapter_repo, character_repo, foreshadowing_repo from services import llm logger = logging.getLogger(__name__) SYSTEM_PROMPT = ( "你是一位专业网文写手, 擅长按大纲续写章节, 文风贴合用户指定风格。" "只输出章节正文, 不要输出标题、解释或任何元信息。" ) _EARLIER_SUMMARY_CHARS = 200 # per earlier-chapter heuristic summary _PREFIX_TAIL_CHARS = 500 # how much of a resume-prefix to echo back def _truncate(text, limit): if not isinstance(text, str): return text if len(text) <= limit: return text return text[:limit] + "..." def build_chapter_context(db, novel, chapter, ctx_prev=None): """Assemble the context dict for writing ``chapter`` of ``novel``. Budget enforcement: drop earlier summaries first, then truncate prev chapter contents, to stay under ``config.CTX_BUDGET_MAX``. """ novel_id = novel["id"] prev_limit = ctx_prev or config.CTX_PREV_CHAPTERS prev = chapter_repo.get_previous_chapters( db, novel_id, chapter.get("volume") or 1, chapter.get("chapter_number") or 1, limit=prev_limit, ) prev_chapters = [ {"title": c.get("title"), "volume": c.get("volume"), "chapter_number": c.get("chapter_number"), "content": c.get("content") or ""} for c in reversed(prev) # chronological order ] prev_ids = {c["id"] for c in prev} done = chapter_repo.list_done_chapters( db, novel_id, limit=config.CTX_SUMMARY_CHAPTERS ) earlier_summaries = [ {"title": c.get("title"), "volume": c.get("volume"), "chapter_number": c.get("chapter_number"), "summary": _truncate(c.get("content") or "", _EARLIER_SUMMARY_CHARS)} for c in done if c["id"] not in prev_ids ] outline = chapter.get("outline") or "" characters = character_repo.find_characters_in_text(db, novel_id, outline) if not characters: characters = character_repo.list_characters(db, novel_id) context = { "novel_meta": { "title": novel.get("title"), "genre": novel.get("genre"), "style": novel.get("style"), }, "prev_chapters": prev_chapters, "earlier_summaries": earlier_summaries, "characters": characters, "open_foreshadowing": foreshadowing_repo.list_open_foreshadowing(db, novel_id), } serialized = json.dumps(context, ensure_ascii=False, default=str) # Budget pass 1: drop earlier summaries. if len(serialized) > config.CTX_BUDGET_MAX and context["earlier_summaries"]: logger.info("chapter context over budget, dropping earlier summaries") context["earlier_summaries"] = [] serialized = json.dumps(context, ensure_ascii=False, default=str) # Budget pass 2: truncate prev chapter bodies proportionally. if len(serialized) > config.CTX_BUDGET_MAX and context["prev_chapters"]: per = max(1000, config.CTX_BUDGET_MAX // (2 * len(context["prev_chapters"]))) for c in context["prev_chapters"]: c["content"] = _truncate(c["content"], per) serialized = json.dumps(context, ensure_ascii=False, default=str) logger.info("chapter context over budget, truncated prev chapters to %d chars each", per) context["total_chars"] = len(serialized) if len(serialized) < config.CTX_BUDGET_MIN: logger.info( "chapter context small (%d chars < %d): early chapter or sparse data", len(serialized), config.CTX_BUDGET_MIN, ) return context def build_chapter_prompts(context, chapter, prefix=None): """Assemble system/user messages for chapter generation.""" meta = context["novel_meta"] parts = [ f"小说: 《{meta['title']}》 类型: {meta['genre']} 文风: {meta['style']}", "", ] if context["earlier_summaries"]: parts.append("【更早章节摘要】") for s in context["earlier_summaries"]: parts.append( f"第{s['volume']}卷 第{s['chapter_number']}章 {s.get('title') or ''}: {s['summary']}" ) parts.append("") if context["prev_chapters"]: parts.append("【最近章节正文】") for c in context["prev_chapters"]: parts.append( f"--- 第{c['volume']}卷 第{c['chapter_number']}章 {c.get('title') or ''} ---" ) parts.append(c["content"]) parts.append("") if context["characters"]: parts.append("【相关角色】") for ch in context["characters"]: parts.append( f"{ch.get('name')} ({ch.get('role') or '角色'}, 状态: {ch.get('status')}): " f"{_truncate(ch.get('description') or '', 300)}" ) parts.append("") if context["open_foreshadowing"]: parts.append("【未回收伏笔】") for f in context["open_foreshadowing"]: parts.append(f"- {f.get('description')} (埋设: 第{f.get('planted_chapter') or '?'}章)") parts.append("") parts.append("【当前章节大纲】") parts.append( f"第{chapter.get('volume') or 1}卷 第{chapter.get('chapter_number') or 1}章 " f"{chapter.get('title') or ''}" ) parts.append(chapter.get("outline") or "(无大纲, 请自由发挥但承接前文)") parts.append("") if prefix: parts.append("【续写指令】以下为本章已生成的开头, 请无缝续写, 不要重复已有内容:") parts.append(prefix[-_PREFIX_TAIL_CHARS:]) parts.append("") parts.append("要求: 写 3000-8000 字正文, 只输出正文, 不要标题不要解释。") user_prompt = "\n".join(parts) return [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ] def stream_chapter(db, novel, chapter, prefix=None): """Generator yielding SSE frames: context → token… → done / error. NOTE: ``db`` is only used here to build the context; callers must build everything inside the request view and may pass a connection that dies with the request — so we build context eagerly before streaming. """ context = build_chapter_context(db, novel, chapter) summary = ( f"已注入上下文: 前 {len(context['prev_chapters'])} 章正文, " f"{len(context['earlier_summaries'])} 条更早摘要, " f"{len(context['characters'])} 个角色, " f"{len(context['open_foreshadowing'])} 条未回收伏笔, " f"共约 {context['total_chars']} 字符" + (f"; 续写前缀 {len(prefix)} 字符" if prefix else "") ) yield llm._sse("context", {"summary": summary}) messages = build_chapter_prompts(context, chapter, prefix=prefix) parts = [] def _on_token(t): parts.append(t) yield from llm.make_token_stream(messages, _on_token, max_tokens=8192) full_text = (prefix or "") + "".join(parts) yield llm._sse("done", {"content": full_text, "word_count": len(full_text)})