195 lines
7.3 KiB
Python
195 lines
7.3 KiB
Python
"""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)})
|