318 lines
12 KiB
Python
318 lines
12 KiB
Python
import os
|
|
from typing_extensions import TypedDict
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import StateGraph, START, END
|
|
from typing import Dict
|
|
|
|
from utils import (
|
|
read_file, append_text_to_file, clear_file_content, save_string_to_txt
|
|
)
|
|
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
|
|
)
|
|
|
|
# 向量检索相关 (以Chroma为例)，需要安装 langchain, chromadb 等
|
|
from langchain.embeddings.openai import OpenAIEmbeddings
|
|
from langchain.vectorstores import Chroma
|
|
from langchain.docstore.document import Document
|
|
|
|
# 默认用此目录存放向量库
|
|
VECTOR_STORE_DIR = "vectorstore"
|
|
|
|
# =============== 多步生成：设置 & 目录 ===============
|
|
class OverallState(TypedDict):
|
|
topic: str
|
|
genre: str
|
|
number_of_chapters: int
|
|
word_number: int
|
|
novel_setting_base: str
|
|
character_setting: str
|
|
dark_lines: str
|
|
final_novel_setting: str
|
|
novel_directory: str
|
|
|
|
def Novel_novel_directory_generate(
|
|
api_key: str,
|
|
base_url: str,
|
|
llm_model: str,
|
|
topic: str,
|
|
genre: str,
|
|
number_of_chapters: int,
|
|
word_number: int,
|
|
filepath: str
|
|
):
|
|
"""
|
|
使用多步流程，生成 Novel_setting.txt 与 Novel_directory.txt
|
|
"""
|
|
model = ChatOpenAI(
|
|
model=llm_model,
|
|
api_key=api_key,
|
|
base_url=base_url
|
|
)
|
|
|
|
def generate_base_setting(state: OverallState):
|
|
prompt = set_prompt.format(
|
|
topic=state["topic"],
|
|
genre=state["genre"],
|
|
number_of_chapters=state["number_of_chapters"],
|
|
word_number=state["word_number"],
|
|
)
|
|
response = model.invoke(prompt)
|
|
if not response:
|
|
return {"novel_setting_base": ""}
|
|
return {"novel_setting_base": response.content.strip()}
|
|
|
|
def generate_character_setting(state: OverallState):
|
|
prompt = character_prompt.format(novel_setting=state["novel_setting_base"])
|
|
response = model.invoke(prompt)
|
|
if not response:
|
|
return {"character_setting": ""}
|
|
return {"character_setting": response.content.strip()}
|
|
|
|
def generate_dark_lines(state: OverallState):
|
|
prompt = dark_lines_prompt.format(character_info=state["character_setting"])
|
|
response = model.invoke(prompt)
|
|
if not response:
|
|
return {"dark_lines": ""}
|
|
return {"dark_lines": response.content.strip()}
|
|
|
|
def finalize_novel_setting(state: OverallState):
|
|
prompt = finalize_setting_prompt.format(
|
|
novel_setting_base=state["novel_setting_base"],
|
|
character_setting=state["character_setting"],
|
|
dark_lines=state["dark_lines"]
|
|
)
|
|
response = model.invoke(prompt)
|
|
if not response:
|
|
return {"final_novel_setting": ""}
|
|
return {"final_novel_setting": response.content.strip()}
|
|
|
|
def generate_novel_directory(state: OverallState):
|
|
prompt = novel_directory_prompt.format(
|
|
final_novel_setting=state["final_novel_setting"],
|
|
number_of_chapters=state["number_of_chapters"]
|
|
)
|
|
response = model.invoke(prompt)
|
|
if not response:
|
|
return {"novel_directory": ""}
|
|
return {"novel_directory": response.content.strip()}
|
|
|
|
graph = StateGraph(OverallState)
|
|
graph.add_node("generate_base_setting", generate_base_setting)
|
|
graph.add_node("generate_character_setting", generate_character_setting)
|
|
graph.add_node("generate_dark_lines", generate_dark_lines)
|
|
graph.add_node("finalize_novel_setting", finalize_novel_setting)
|
|
graph.add_node("generate_novel_directory", generate_novel_directory)
|
|
|
|
graph.add_edge(START, "generate_base_setting")
|
|
graph.add_edge("generate_base_setting", "generate_character_setting")
|
|
graph.add_edge("generate_character_setting", "generate_dark_lines")
|
|
graph.add_edge("generate_dark_lines", "finalize_novel_setting")
|
|
graph.add_edge("finalize_novel_setting", "generate_novel_directory")
|
|
graph.add_edge("generate_novel_directory", END)
|
|
|
|
app = graph.compile()
|
|
|
|
input_params = {
|
|
"topic": topic,
|
|
"genre": genre,
|
|
"number_of_chapters": number_of_chapters,
|
|
"word_number": word_number
|
|
}
|
|
result = app.invoke(input_params)
|
|
|
|
if not result:
|
|
print("⚠️ invoke() 结果为空，生成失败。")
|
|
return
|
|
|
|
final_novel_setting = result.get("final_novel_setting", "")
|
|
final_novel_directory = result.get("novel_directory", "")
|
|
if not final_novel_setting or not final_novel_directory:
|
|
print("⚠️ 生成失败：缺少 final_novel_setting 或 novel_directory。")
|
|
return
|
|
|
|
# 写入文件
|
|
filename_set = os.path.join(filepath, "Novel_setting.txt")
|
|
filename_novel_directory = os.path.join(filepath, "Novel_directory.txt")
|
|
|
|
final_novel_setting_cleaned = final_novel_setting.replace('#', '').replace('*', '')
|
|
final_novel_directory_cleaned = final_novel_directory.replace('#', '').replace('*', '')
|
|
|
|
append_text_to_file(final_novel_setting_cleaned, filename_set)
|
|
append_text_to_file(final_novel_directory_cleaned, filename_novel_directory)
|
|
|
|
|
|
# =============== 生成章节（含角色状态 & 全局摘要 & 向量检索） ===============
|
|
|
|
def init_vector_store(api_key: str, texts: list[str]) -> Chroma:
|
|
"""
|
|
初始化并返回一个Chroma向量库，将传入的文本进行嵌入。
|
|
若需要可对 texts 做分句或分块处理；这里只演示简单用法。
|
|
"""
|
|
embeddings = OpenAIEmbeddings(openai_api_key=api_key)
|
|
documents = [Document(page_content=t) for t in texts]
|
|
vectorstore = Chroma.from_documents(documents, embedding=embeddings, persist_directory=VECTOR_STORE_DIR)
|
|
vectorstore.persist()
|
|
return vectorstore
|
|
|
|
def load_vector_store(api_key: str) -> Chroma:
|
|
"""
|
|
读取已存在的向量库。若不存在则返回None或新建一个空的。
|
|
"""
|
|
if not os.path.exists(VECTOR_STORE_DIR):
|
|
return None
|
|
embeddings = OpenAIEmbeddings(openai_api_key=api_key)
|
|
return Chroma(persist_directory=VECTOR_STORE_DIR, embedding_function=embeddings)
|
|
|
|
|
|
def update_vector_store(api_key: str, new_chapter: str):
|
|
"""
|
|
将最新章节文本插入到向量库里，用于后续检索参考。
|
|
可根据实际需求做分块处理。此处仅作简单示范。
|
|
"""
|
|
store = load_vector_store(api_key)
|
|
if not store:
|
|
# 如果vector store不存在，先初始化
|
|
store = init_vector_store(api_key, [new_chapter])
|
|
return
|
|
|
|
embeddings = OpenAIEmbeddings(openai_api_key=api_key)
|
|
new_doc = Document(page_content=new_chapter)
|
|
store.add_documents([new_doc])
|
|
store.persist()
|
|
|
|
|
|
def get_relevant_context_from_vector_store(api_key: str, query: str, k: int=2) -> str:
|
|
"""
|
|
从向量库中检索与 query 最相关的 k 条文本，拼接后返回。
|
|
用于在生成大纲或写正文时，为大模型提供更多上下文。
|
|
"""
|
|
store = load_vector_store(api_key)
|
|
if not store:
|
|
return ""
|
|
docs = store.similarity_search(query, k=k)
|
|
# 简单拼接
|
|
combined = "\n".join([d.page_content for d in docs])
|
|
return combined
|
|
|
|
def generate_chapter_with_state(
|
|
novel_settings: str,
|
|
novel_novel_directory: str,
|
|
api_key: str,
|
|
base_url: str,
|
|
model_name: str,
|
|
novel_number: int,
|
|
filepath: str,
|
|
word_number: int,
|
|
lastchapter: str
|
|
) -> str:
|
|
"""
|
|
多步流程:
|
|
1) 更新/创建全局摘要
|
|
2) 更新/生成角色状态文档
|
|
3) 根据向量检索获取往期章节相关内容
|
|
4) 大纲 -> 正文
|
|
5) 更新向量库
|
|
最终写入 chapter.txt、lastchapter.txt、character_state.txt、global_summary.txt
|
|
"""
|
|
model = ChatOpenAI(
|
|
model=model_name,
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
temperature=0.9
|
|
)
|
|
|
|
# --- 文件名定义 ---
|
|
character_state_file = os.path.join(filepath, "character_state.txt")
|
|
global_summary_file = os.path.join(filepath, "global_summary.txt")
|
|
chapter_file = os.path.join(filepath, "chapter.txt")
|
|
lastchapter_file = os.path.join(filepath, "lastchapter.txt")
|
|
|
|
# --- 读取现有文档（可能为空） ---
|
|
old_char_state = read_file(character_state_file)
|
|
old_global_summary = read_file(global_summary_file)
|
|
|
|
# --- 1) 更新全局摘要 (若上一章文本不为空) ---
|
|
def update_global_summary(chapter_text: str, old_summary: str) -> str:
|
|
prompt = summary_prompt.format(chapter_text=chapter_text, global_summary=old_summary)
|
|
response = model.invoke(prompt)
|
|
if not response:
|
|
return old_summary
|
|
return response.content.strip()
|
|
|
|
if lastchapter.strip():
|
|
# 用上一章内容更新全局摘要
|
|
new_global_summary = update_global_summary(lastchapter, old_global_summary)
|
|
else:
|
|
new_global_summary = old_global_summary
|
|
|
|
# --- 2) 更新角色状态文档 ---
|
|
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)
|
|
response = model.invoke(prompt)
|
|
if not response:
|
|
return old_state
|
|
return response.content.strip()
|
|
|
|
if lastchapter.strip():
|
|
new_char_state = update_character_state(lastchapter, old_char_state)
|
|
else:
|
|
new_char_state = old_char_state
|
|
|
|
# --- 3) 从向量库检索相关上下文，用来帮助生成新的大纲 ---
|
|
# 例如，可以根据“角色状态”或“本章关键词”来查询。
|
|
# 简单示范：以 "回顾剧情" 作为检索Query
|
|
relevant_context = get_relevant_context_from_vector_store(api_key, "回顾剧情", k=2)
|
|
|
|
# --- 4) 大纲 -> 正文 ---
|
|
def outline_chapter(novel_setting: str, char_state: str, global_summary: str, chap_num: int, extra_context: str) -> str:
|
|
prompt = chapter_outline_prompt.format(
|
|
novel_setting=novel_setting,
|
|
character_state=char_state + "\n\n【历史上下文】\n" + extra_context,
|
|
global_summary=global_summary,
|
|
novel_number=chap_num
|
|
)
|
|
response = model.invoke(prompt)
|
|
if not response:
|
|
return ""
|
|
return response.content.strip()
|
|
|
|
chap_outline = outline_chapter(novel_settings, new_char_state, new_global_summary, novel_number, relevant_context)
|
|
|
|
def write_chapter(novel_setting: str, char_state: str, global_summary: str, outline: str, wnum: int, extra_context: str) -> str:
|
|
prompt = chapter_write_prompt.format(
|
|
novel_setting=novel_setting,
|
|
character_state=char_state + "\n\n【历史上下文】\n" + extra_context,
|
|
global_summary=global_summary,
|
|
chapter_outline=outline,
|
|
word_number=wnum
|
|
)
|
|
response = model.invoke(prompt)
|
|
if not response:
|
|
return ""
|
|
return response.content.strip()
|
|
|
|
chapter_content = write_chapter(novel_settings, new_char_state, new_global_summary, chap_outline, word_number, relevant_context)
|
|
|
|
if chapter_content:
|
|
# --- 写入 chapter.txt 与 lastchapter.txt ---
|
|
append_text_to_file(chapter_content, chapter_file)
|
|
|
|
clear_file_content(lastchapter_file)
|
|
save_string_to_txt(chapter_content, lastchapter_file)
|
|
|
|
# --- 更新全局摘要、角色状态到文件 ---
|
|
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)
|
|
|
|
# --- 5) 更新向量检索库 ---
|
|
update_vector_store(api_key, chapter_content)
|
|
|
|
return chapter_content
|