lizhilun/AI_NovelGenerator
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

使用新的生成逻辑

Browse Source
This commit is contained in:
YILING0013
2025-02-05 21:55:48 +08:00
parent d230d4ba23
commit dd78666071
5 changed files with 1859 additions and 1209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
novel_generator.py
+225 -184
View File
@@ -15,7 +15,6 @@ from langchain.docstore.document import Document
# nltk、sentence_transformers 及文本处理相关
import nltk
import math
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
@@ -27,26 +26,26 @@ from utils import (
# prompt模板
from prompt_definitions import (
# 设定相关
set_prompt, character_prompt, dark_lines_prompt,
finalize_setting_prompt, novel_directory_prompt,
# 写作流程相关
summary_prompt, update_character_state_prompt,
chapter_outline_prompt, chapter_write_prompt
core_seed_prompt,
character_dynamics_prompt,
world_building_prompt,
plot_architecture_prompt,
chapter_blueprint_prompt,
summary_prompt,
update_character_state_prompt,
scene_dynamics_prompt
)
# Ollama嵌入 (如使用Ollama时需要)
from embedding_ollama import OllamaEmbeddings
# 用于目录解析章节标题/简介
from chapter_directory_parser import get_chapter_info_from_directory
from chapter_directory_parser import get_chapter_info_from_blueprint
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
# ============ 帮助函数 ============
# ============ 基础工具 ============
def remove_think_tags(text: str) -> str:
"""移除 <think>...</think> 包裹的内容"""
return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
@@ -333,8 +332,8 @@ def get_relevant_context_from_vector_store(
return combined
# ============ 1. 生成小说“设定” (Novel_setting.txt) ============
def Novel_setting_generate(
# ========== 1) 生成总体架构 (Novel_architecture.txt) ==========
def Novel_architecture_generate(
api_key: str,
base_url: str,
llm_model: str,
@@ -345,8 +344,15 @@ def Novel_setting_generate(
filepath: str,
temperature: float = 0.7
) -> None:
"""
依次调用：
1. core_seed_prompt
2. character_dynamics_prompt
3. world_building_prompt
4. plot_architecture_prompt
将结果整合为“Novel_architecture.txt”。
"""
os.makedirs(filepath, exist_ok=True)
model = ChatOpenAI(
model=llm_model,
api_key=api_key,
@@ -354,58 +360,95 @@ def Novel_setting_generate(
temperature=temperature
)
# Step1: 基础设定
prompt_base = set_prompt.format(
# 1) 核心种子
prompt_core = core_seed_prompt.format(
topic=topic,
genre=genre,
number_of_chapters=number_of_chapters,
word_number=word_number
)
base_setting = invoke_with_cleaning(model, prompt_base)
core_seed_result = invoke_with_cleaning(model, prompt_core)
core_seed_text = core_seed_result.strip()
# Step2: 角色设定
prompt_char = character_prompt.format(
novel_setting=base_setting
# 2) 角色动力学
prompt_character = character_dynamics_prompt.format(core_seed=core_seed_text)
character_dynamics_result = invoke_with_cleaning(model, prompt_character)
character_dynamics_text = character_dynamics_result.strip()
# 3) 世界观
prompt_world = world_building_prompt.format(core_seed=core_seed_text)
world_building_result = invoke_with_cleaning(model, prompt_world)
world_building_text = world_building_result.strip()
# 4) 三幕式情节架构
prompt_plot = plot_architecture_prompt.format(
core_seed=core_seed_text,
character_dynamics=character_dynamics_text,
world_building=world_building_text
)
character_setting = invoke_with_cleaning(model, prompt_char)
plot_arch_result = invoke_with_cleaning(model, prompt_plot)
plot_arch_text = plot_arch_result.strip()
# Step3: 暗线/雷点
prompt_dark = dark_lines_prompt.format(
character_info=character_setting
# 整合并写入 Novel_architecture.txt
final_content = (
"#=== 1) 核心种子 ===\n"
f"{core_seed_text}\n\n"
"#=== 2) 角色动力学 ===\n"
f"{character_dynamics_text}\n\n"
"#=== 3) 世界观 ===\n"
f"{world_building_text}\n\n"
"#=== 4) 三幕式情节架构 ===\n"
f"{plot_arch_text}\n"
)
dark_lines = invoke_with_cleaning(model, prompt_dark)
# Step4: 最终整合
prompt_final = finalize_setting_prompt.format(
novel_setting_base=base_setting,
character_setting=character_setting,
dark_lines=dark_lines
)
final_novel_setting = invoke_with_cleaning(model, prompt_final)
arch_file = os.path.join(filepath, "Novel_architecture.txt")
clear_file_content(arch_file)
save_string_to_txt(final_content, arch_file)
filename_set = os.path.join(filepath, "Novel_setting.txt")
clear_file_content(filename_set)
final_novel_setting_cleaned = final_novel_setting.replace('#', '').replace('*', '')
save_string_to_txt(final_novel_setting_cleaned, filename_set)
logging.info("Novel_setting.txt has been generated successfully.")
logging.info("Novel_architecture.txt has been generated successfully.")
# ============ 2. 生成小说目录 (Novel_directory.txt) ============
def Novel_directory_generate(
# ========== 2) 生成章节蓝图 (Novel_directory.txt) ==========
def Chapter_blueprint_generate(
api_key: str,
base_url: str,
llm_model: str,
number_of_chapters: int,
filepath: str,
temperature: float = 0.7
) -> None:
filename_set = os.path.join(filepath, "Novel_setting.txt")
final_novel_setting = read_file(filename_set).strip()
if not final_novel_setting:
logging.warning("Novel_setting.txt 内容为空，请先生成小说设定。")
"""
基于“Novel_architecture.txt”中的三幕式情节架构，调用 chapter_blueprint_prompt，
生成章节蓝图并写入 Novel_directory.txt。
"""
arch_file = os.path.join(filepath, "Novel_architecture.txt")
if not os.path.exists(arch_file):
logging.warning("Novel_architecture.txt not found. Please generate architecture first.")
return
architecture_text = read_file(arch_file).strip()
if not architecture_text:
logging.warning("Novel_architecture.txt is empty.")
return
# 从内容中尽量提取 number_of_chapters
# 如果之前已经存储了 number_of_chapters，可以在外面传入，这里做简化：
# 这里用正则或者其他逻辑提取，但演示时直接写 10 也可
match_chaps = re.search(r'约(\d+)章', architecture_text)
if match_chaps:
number_of_chapters = int(match_chaps.group(1))
else:
number_of_chapters = 10 # fallback
# 提取三幕式文本
# 在写入时，我们将 4) 三幕式情节架构 作为传给 prompt 的核心
# 这里做一个简易匹配
plot_arch_text = ""
# 假设 "#=== 4) 三幕式情节架构 ===" 是分隔点
pat_plot = r'#=== 4\) 三幕式情节架构 ===\n([\s\S]+)$'
m = re.search(pat_plot, architecture_text)
if m:
plot_arch_text = m.group(1).strip()
model = ChatOpenAI(
model=llm_model,
api_key=api_key,
@@ -413,22 +456,20 @@ def Novel_directory_generate(
temperature=temperature
)
prompt_dir = novel_directory_prompt.format(
final_novel_setting=final_novel_setting,
prompt = chapter_blueprint_prompt.format(
plot_architecture=plot_arch_text,
number_of_chapters=number_of_chapters
)
final_novel_directory = invoke_with_cleaning(model, prompt_dir)
if not final_novel_directory.strip():
logging.warning("Novel_directory生成结果为空。")
blueprint_text = invoke_with_cleaning(model, prompt)
if not blueprint_text.strip():
logging.warning("Chapter blueprint generation result is empty.")
return
filename_dir = os.path.join(filepath, "Novel_directory.txt")
clear_file_content(filename_dir)
save_string_to_txt(blueprint_text, filename_dir)
final_novel_directory_cleaned = final_novel_directory.replace('#', '').replace('*', '')
save_string_to_txt(final_novel_directory_cleaned, filename_dir)
logging.info("Novel_directory.txt has been generated successfully.")
logging.info("Novel_directory.txt (chapter blueprint) has been generated successfully.")
# ============ 获取最近 N 章内容，生成短期摘要 ============
@@ -514,58 +555,114 @@ def update_plot_arcs(
return arcs_text
# ============ 生成章节草稿 ============
# ========== 3) 生成章节草稿 ==========
def generate_chapter_draft(
novel_settings: str,
global_summary: str,
character_state: str,
recent_chapters_summary: str,
user_guidance: str,
api_key: str,
base_url: str,
model_name: str,
filepath: str,
novel_number: int,
word_number: int,
temperature: float,
novel_novel_directory: str,
filepath: str,
interface_format: str,
embedding_model_name: str,
embedding_base_url: str,
embedding_retrieval_k: int = 4
user_guidance: str,
characters_involved: str,
key_items: str,
scene_location: str,
time_constraint: str,
embedding_retrieval_k: int = 2
) -> str:
# 1) 根据目录解析标题、简介
chapter_info = get_chapter_info_from_directory(novel_novel_directory, novel_number)
"""
根据 scene_dynamics_prompt，生成本章草稿。
- novel_architecture 取自 Novel_architecture.txt
- blueprint 取自 Novel_directory.txt
- global_summary, character_state 分别取自全局摘要、角色状态文件
- 向量库检索上下文
- 用户还可以额外提供四个可选元素：核心人物、关键道具、空间坐标、时间压力
"""
# 1) 读取相关文件
arch_file = os.path.join(filepath, "Novel_architecture.txt")
novel_architecture_text = read_file(arch_file)
directory_file = os.path.join(filepath, "Novel_directory.txt")
blueprint_text = read_file(directory_file)
global_summary_file = os.path.join(filepath, "global_summary.txt")
global_summary_text = read_file(global_summary_file)
character_state_file = os.path.join(filepath, "character_state.txt")
character_state_text = read_file(character_state_file)
# 2) 解析 blueprint，得到本章所需的字段
chapter_info = get_chapter_info_from_blueprint(blueprint_text, novel_number)
chapter_title = chapter_info["chapter_title"]
chapter_brief = chapter_info["chapter_brief"]
chapter_role = chapter_info["chapter_role"]
chapter_purpose = chapter_info["chapter_purpose"]
suspense_level = chapter_info["suspense_level"]
foreshadowing = chapter_info["foreshadowing"]
plot_twist_level = chapter_info["plot_twist_level"]
chapter_summary = chapter_info["chapter_summary"]
# 合并要检索的文本（用户指导 + 章节简介 + 最近摘要）
combined_query_parts = []
if user_guidance.strip():
combined_query_parts.append(user_guidance)
if chapter_brief.strip():
combined_query_parts.append(chapter_brief)
if recent_chapters_summary.strip():
combined_query_parts.append(recent_chapters_summary)
# 额外加一个关键字
combined_query_parts.append("回顾剧情")
# 3) 取最近3章文本，拼成查询语句 => 用于向量库检索
chapters_dir = os.path.join(filepath, "chapters")
recent_3_texts = get_last_n_chapters_text(chapters_dir, novel_number, n=3)
merged_query_str = "回顾剧情：\n" + "\n".join(recent_3_texts) + "\n" + user_guidance
merged_query_str = "\n".join(combined_query_parts)
# 2) 从向量库检索上下文
# 4) 检索向量库上下文
relevant_context = get_relevant_context_from_vector_store(
api_key=api_key,
base_url=embedding_base_url if embedding_base_url else base_url,
base_url=base_url,
query=merged_query_str,
interface_format=interface_format,
embedding_model_name=embedding_model_name,
embedding_model_name=model_name,
filepath=filepath,
k=embedding_retrieval_k
)
if not relevant_context.strip():
relevant_context = "暂无相关内容。"
# 3) 生成本章大纲
if not relevant_context.strip():
relevant_context = "（无检索到的上下文）"
# 5) 构造prompt，调用 scene_dynamics_prompt
# 在这里，我们拆分架构文本，以便给模型提供：
# - “世界观”与“小说设定”可以从 arch_file 中的相应片段读取
# 这里为了简化，直接把 novel_architecture_text 整体塞入 novel_setting
# 也可更精细地拆分 "#=== 3) 世界观 ===" 片段给 world_building
# 下方仅作示例。
world_building_text = ""
match_world = re.search(r'#=== 3\) 世界观 ===\n([\s\S]+?)\n#===', novel_architecture_text)
if match_world:
world_building_text = match_world.group(1).strip()
else:
world_building_text = "暂无世界观信息"
novel_setting_text = novel_architecture_text # 整份当做“小说设定”参考
prompt_text = scene_dynamics_prompt.format(
novel_number=novel_number,
chapter_title=chapter_title,
chapter_role=chapter_role,
chapter_purpose=chapter_purpose,
suspense_level=suspense_level,
foreshadowing=foreshadowing,
plot_twist_level=plot_twist_level,
chapter_summary=chapter_summary,
characters_involved=characters_involved,
key_items=key_items,
scene_location=scene_location,
time_constraint=time_constraint,
world_building=world_building_text,
novel_setting=novel_setting_text,
global_summary=global_summary_text,
character_state=character_state_text
)
# 因为我们还想让模型了解向量库检索到的上下文，可以合并到最后
prompt_text += f"\n\n【检索到的上下文】\n{relevant_context}"
# 也可合并用户指导
prompt_text += f"\n\n【用户指导】\n{user_guidance}\n"
model = ChatOpenAI(
model=model_name,
api_key=api_key,
@@ -573,44 +670,15 @@ def generate_chapter_draft(
temperature=temperature
)
outline_prompt_text = chapter_outline_prompt.format(
novel_setting=novel_settings,
character_state=character_state + "\n\n【检索到的上下文】\n" + relevant_context,
global_summary=global_summary,
novel_number=novel_number,
chapter_title=chapter_title,
chapter_brief=chapter_brief
)
outline_prompt_text += f"\n\n【最近几章摘要】\n{recent_chapters_summary}"
outline_prompt_text += f"\n\n【用户指导】\n{user_guidance if user_guidance else '（无）'}"
chapter_outline = invoke_with_cleaning(model, outline_prompt_text)
outlines_dir = os.path.join(filepath, "outlines")
os.makedirs(outlines_dir, exist_ok=True)
outline_file = os.path.join(outlines_dir, f"outline_{novel_number}.txt")
clear_file_content(outline_file)
save_string_to_txt(chapter_outline, outline_file)
# 4) 生成正文草稿
writing_prompt_text = chapter_write_prompt.format(
novel_setting=novel_settings,
character_state=character_state + "\n\n【检索到的上下文】\n" + relevant_context,
global_summary=global_summary,
chapter_outline=chapter_outline,
word_number=word_number,
novel_number=novel_number,
chapter_title=chapter_title,
chapter_brief=chapter_brief
)
writing_prompt_text += f"\n\n【最近几章摘要】\n{recent_chapters_summary}"
writing_prompt_text += f"\n\n【用户指导】\n{user_guidance if user_guidance else '（无）'}"
chapter_content = invoke_with_cleaning(model, writing_prompt_text)
chapter_content = invoke_with_cleaning(model, prompt_text)
if not chapter_content.strip():
logging.warning("Generated chapter draft is empty.")
# 6) 写入 chapters 目录
chapters_dir = os.path.join(filepath, "chapters")
os.makedirs(chapters_dir, exist_ok=True)
chapter_file = os.path.join(chapters_dir, f"chapter_{novel_number}.txt")
clear_file_content(chapter_file)
save_string_to_txt(chapter_content, chapter_file)
@@ -618,20 +686,20 @@ def generate_chapter_draft(
return chapter_content
# ============ 定稿章节 ============
# ========== 4) 定稿章节 ==========
def finalize_chapter(
novel_number: int,
word_number: int,
api_key: str,
base_url: str,
interface_format: str,
embedding_model_name: str,
model_name: str,
temperature: float,
filepath: str,
embedding_base_url: str,
embedding_api_key: str
embedding_model_name: str
):
"""
定稿：更新全局摘要、角色状态，并将本章文本插入向量库。
"""
chapters_dir = os.path.join(filepath, "chapters")
chapter_file = os.path.join(chapters_dir, f"chapter_{novel_number}.txt")
chapter_text = read_file(chapter_file).strip()
@@ -639,82 +707,55 @@ def finalize_chapter(
logging.warning(f"Chapter {novel_number} is empty, cannot finalize.")
return
character_state_file = os.path.join(filepath, "character_state.txt")
global_summary_file = os.path.join(filepath, "global_summary.txt")
plot_arcs_file = os.path.join(filepath, "plot_arcs.txt")
old_char_state = read_file(character_state_file)
old_global_summary = read_file(global_summary_file)
old_plot_arcs = read_file(plot_arcs_file)
# 篇幅不足，二次扩写
if len(chapter_text) < 0.8 * word_number:
logging.info("Chapter text is shorter than 80% of desired length. Enriching...")
chapter_text = enrich_chapter_text(
chapter_text=chapter_text,
word_number=word_number,
api_key=api_key,
base_url=base_url,
model_name=model_name,
temperature=temperature
)
# 如果长度比目标少很多，可考虑在此扩写
if len(chapter_text) < 0.6 * word_number:
chapter_text = enrich_chapter_text(chapter_text, word_number, api_key, base_url, model_name, temperature)
clear_file_content(chapter_file)
save_string_to_txt(chapter_text, chapter_file)
# 更新全局摘要
# 读取全局摘要、角色状态
global_summary_file = os.path.join(filepath, "global_summary.txt")
old_global_summary = read_file(global_summary_file)
character_state_file = os.path.join(filepath, "character_state.txt")
old_character_state = read_file(character_state_file)
# 1) 更新全局摘要
model = ChatOpenAI(
model=model_name,
api_key=api_key,
base_url=ensure_openai_base_url_has_v1(base_url),
temperature=temperature
)
def update_global_summary(chapter_text: str, old_summary: str) -> str:
prompt = summary_prompt.format(
chapter_text=chapter_text,
global_summary=old_summary
)
return invoke_with_cleaning(model, prompt) or old_summary
new_global_summary = update_global_summary(chapter_text, old_global_summary)
# 更新角色状态
def update_character_state(chapter_text: str, old_state: str) -> str:
prompt = update_character_state_prompt.format(
chapter_text=chapter_text,
old_state=old_state
)
return invoke_with_cleaning(model, prompt) or old_state
new_char_state = update_character_state(chapter_text, old_char_state)
# 更新剧情要点
new_plot_arcs = update_plot_arcs(
prompt_summary = summary_prompt.format(
chapter_text=chapter_text,
old_plot_arcs=old_plot_arcs,
api_key=api_key,
base_url=base_url,
model_name=model_name,
temperature=temperature
global_summary=old_global_summary
)
new_global_summary = invoke_with_cleaning(model, prompt_summary)
if not new_global_summary.strip():
new_global_summary = old_global_summary
# 2) 更新角色状态
prompt_char_state = update_character_state_prompt.format(
chapter_text=chapter_text,
old_state=old_character_state
)
new_char_state = invoke_with_cleaning(model, prompt_char_state)
if not new_char_state.strip():
new_char_state = old_character_state
# 写回文件
clear_file_content(character_state_file)
save_string_to_txt(new_char_state, character_state_file)
clear_file_content(global_summary_file)
save_string_to_txt(new_global_summary, global_summary_file)
clear_file_content(plot_arcs_file)
save_string_to_txt(new_plot_arcs, plot_arcs_file)
clear_file_content(character_state_file)
save_string_to_txt(new_char_state, character_state_file)
# 更新向量库（此时用 embedding_api_key/embedding_base_url）
# 3) 更新向量库
update_vector_store(
api_key=embedding_api_key,
base_url=embedding_base_url if embedding_base_url else base_url,
api_key=api_key,
base_url=base_url,
new_chapter=chapter_text,
interface_format=interface_format,
embedding_model_name=embedding_model_name,
model_name=embedding_model_name, # 用于embedding
filepath=filepath
)