2025-01-31 13:50:07 +08:00
|
|
|
|
# novel_generator.py
|
|
|
|
|
|
# -*- coding: utf-8 -*-
|
2025-01-29 20:33:20 +08:00
|
|
|
|
import os
|
2025-01-29 21:59:36 +08:00
|
|
|
|
import logging
|
2025-01-31 13:50:07 +08:00
|
|
|
|
import re
|
2025-02-04 13:23:50 +08:00
|
|
|
|
import time
|
2025-02-02 19:22:54 +08:00
|
|
|
|
import traceback
|
2025-02-03 13:00:03 +08:00
|
|
|
|
from typing import List, Optional
|
2025-02-02 19:17:07 +08:00
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
# langchain 相关
|
2025-02-04 16:11:17 +08:00
|
|
|
|
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
|
2025-02-04 13:23:50 +08:00
|
|
|
|
from langchain_chroma import Chroma
|
2025-02-03 20:53:37 +08:00
|
|
|
|
from chromadb.config import Settings
|
2025-01-29 21:59:36 +08:00
|
|
|
|
from langchain.docstore.document import Document
|
2025-02-02 19:17:07 +08:00
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
# nltk、sentence_transformers 及文本处理相关
|
2025-01-31 13:50:07 +08:00
|
|
|
|
import nltk
|
|
|
|
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
|
|
from sklearn.metrics.pairwise import cosine_similarity
|
|
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
# 工具函数
|
2025-01-29 20:33:20 +08:00
|
|
|
|
from utils import (
|
2025-01-29 21:59:36 +08:00
|
|
|
|
read_file, append_text_to_file, clear_file_content,
|
|
|
|
|
|
save_string_to_txt
|
2025-01-29 20:33:20 +08:00
|
|
|
|
)
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
|
|
|
|
|
# prompt模板
|
2025-01-29 20:33:20 +08:00
|
|
|
|
from prompt_definitions import (
|
2025-02-05 21:55:48 +08:00
|
|
|
|
core_seed_prompt,
|
|
|
|
|
|
character_dynamics_prompt,
|
|
|
|
|
|
world_building_prompt,
|
|
|
|
|
|
plot_architecture_prompt,
|
|
|
|
|
|
chapter_blueprint_prompt,
|
|
|
|
|
|
summary_prompt,
|
|
|
|
|
|
update_character_state_prompt,
|
|
|
|
|
|
scene_dynamics_prompt
|
2025-01-29 20:33:20 +08:00
|
|
|
|
)
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
|
|
|
|
|
# Ollama嵌入 (如使用Ollama时需要)
|
2025-02-02 12:59:13 +08:00
|
|
|
|
from embedding_ollama import OllamaEmbeddings
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
|
|
|
|
|
# 用于目录解析章节标题/简介
|
2025-02-05 21:55:48 +08:00
|
|
|
|
from chapter_directory_parser import get_chapter_info_from_blueprint
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-02 14:39:30 +08:00
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
|
|
|
|
|
2025-02-02 15:50:19 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# ============ 基础工具 ============
|
2025-02-02 19:17:07 +08:00
|
|
|
|
def remove_think_tags(text: str) -> str:
|
2025-02-04 00:10:19 +08:00
|
|
|
|
"""移除 <think>...</think> 包裹的内容"""
|
2025-02-02 19:17:07 +08:00
|
|
|
|
return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
|
|
|
|
|
|
|
2025-02-04 16:11:17 +08:00
|
|
|
|
def debug_log(prompt: str, response_content: str):
|
|
|
|
|
|
logging.info(f"\n[######################################### Prompt #########################################]\n {prompt}\n")
|
|
|
|
|
|
logging.info(f"\n[######################################### Response #########################################]\n {response_content}\n")
|
|
|
|
|
|
|
2025-02-02 19:17:07 +08:00
|
|
|
|
def invoke_with_cleaning(model: ChatOpenAI, prompt: str) -> str:
|
2025-02-04 00:10:19 +08:00
|
|
|
|
"""通用封装:调用模型并移除 <think>...</think> 文本,记录日志后返回"""
|
2025-02-02 19:17:07 +08:00
|
|
|
|
response = model.invoke(prompt)
|
|
|
|
|
|
if not response:
|
|
|
|
|
|
logging.warning("No response from model.")
|
|
|
|
|
|
return ""
|
|
|
|
|
|
cleaned_text = remove_think_tags(response.content)
|
|
|
|
|
|
debug_log(prompt, cleaned_text)
|
|
|
|
|
|
return cleaned_text.strip()
|
|
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
def ensure_openai_base_url_has_v1(url: str) -> str:
|
|
|
|
|
|
"""
|
2025-02-04 00:10:19 +08:00
|
|
|
|
若用户输入的 url 不包含 '/v1',则在末尾追加 '/v1'。
|
2025-02-03 13:00:03 +08:00
|
|
|
|
"""
|
2025-02-04 00:10:19 +08:00
|
|
|
|
import re
|
2025-02-03 13:00:03 +08:00
|
|
|
|
url = url.strip()
|
|
|
|
|
|
if not url:
|
|
|
|
|
|
return url
|
|
|
|
|
|
if not re.search(r'/v\d+$', url):
|
|
|
|
|
|
if '/v1' not in url:
|
|
|
|
|
|
url = url.rstrip('/') + '/v1'
|
|
|
|
|
|
return url
|
|
|
|
|
|
|
2025-02-04 00:10:19 +08:00
|
|
|
|
def is_using_ollama_api(interface_format: str) -> bool:
|
|
|
|
|
|
return interface_format.lower() == "ollama"
|
|
|
|
|
|
|
|
|
|
|
|
def is_using_ml_studio_api(interface_format: str) -> bool:
|
|
|
|
|
|
return interface_format.lower() == "ml studio"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# ============ 获取 vectorstore 路径 ============
|
|
|
|
|
|
def get_vectorstore_dir(filepath: str) -> str:
|
|
|
|
|
|
"""
|
|
|
|
|
|
返回存储向量库的本地路径:
|
|
|
|
|
|
在用户指定的 `filepath` 下创建/使用 'vectorstore' 文件夹。
|
|
|
|
|
|
"""
|
|
|
|
|
|
return os.path.join(filepath, "vectorstore")
|
|
|
|
|
|
|
|
|
|
|
|
|
2025-02-02 18:25:35 +08:00
|
|
|
|
# ============ 创建 Embeddings 对象 ============
|
2025-02-02 12:59:13 +08:00
|
|
|
|
def create_embeddings_object(
|
|
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
interface_format: str,
|
|
|
|
|
|
embedding_model_name: str
|
|
|
|
|
|
):
|
|
|
|
|
|
"""
|
2025-02-04 00:10:19 +08:00
|
|
|
|
根据 embedding_interface_format,选择 Ollama 或 OpenAIEmbeddings 等不同后端。
|
2025-02-02 12:59:13 +08:00
|
|
|
|
"""
|
2025-02-04 00:10:19 +08:00
|
|
|
|
if is_using_ollama_api(interface_format):
|
|
|
|
|
|
fixed_url = base_url.rstrip("/")
|
2025-02-02 15:50:19 +08:00
|
|
|
|
return OllamaEmbeddings(
|
|
|
|
|
|
model_name=embedding_model_name,
|
|
|
|
|
|
base_url=fixed_url
|
|
|
|
|
|
)
|
2025-02-02 12:59:13 +08:00
|
|
|
|
else:
|
2025-02-04 16:11:17 +08:00
|
|
|
|
# OpenAI 或 ML Studio 均使用 OpenAIEmbeddings,注意 base_url 可能需要 ensure /v1
|
2025-02-04 00:10:19 +08:00
|
|
|
|
fixed_url = ensure_openai_base_url_has_v1(base_url)
|
2025-02-03 13:00:03 +08:00
|
|
|
|
return OpenAIEmbeddings(
|
|
|
|
|
|
openai_api_key=api_key,
|
|
|
|
|
|
openai_api_base=fixed_url,
|
|
|
|
|
|
model=embedding_model_name
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2025-02-02 12:59:13 +08:00
|
|
|
|
|
2025-02-04 00:10:19 +08:00
|
|
|
|
# ============ 向量库相关操作 ============
|
2025-02-04 16:11:17 +08:00
|
|
|
|
def clear_vector_store(filepath: str) -> bool:
|
2025-01-31 19:57:44 +08:00
|
|
|
|
"""
|
2025-02-04 16:11:17 +08:00
|
|
|
|
返回值表示是否成功清空向量库。
|
2025-01-31 19:57:44 +08:00
|
|
|
|
"""
|
2025-02-04 16:11:17 +08:00
|
|
|
|
import shutil
|
2025-02-04 13:23:50 +08:00
|
|
|
|
|
2025-02-04 00:10:19 +08:00
|
|
|
|
store_dir = get_vectorstore_dir(filepath)
|
2025-02-04 13:23:50 +08:00
|
|
|
|
if not os.path.exists(store_dir):
|
2025-01-31 19:57:44 +08:00
|
|
|
|
logging.info("No vector store found to clear.")
|
2025-02-04 16:11:17 +08:00
|
|
|
|
return False
|
2025-02-04 13:23:50 +08:00
|
|
|
|
|
|
|
|
|
|
try:
|
2025-02-04 16:11:17 +08:00
|
|
|
|
if os.path.exists(store_dir):
|
|
|
|
|
|
shutil.rmtree(store_dir)
|
|
|
|
|
|
logging.info(f"Vector store directory '{store_dir}' removed.")
|
|
|
|
|
|
return True
|
2025-02-04 13:23:50 +08:00
|
|
|
|
except Exception as e:
|
2025-02-04 16:11:17 +08:00
|
|
|
|
logging.error(f"程序正在运行,无法删除,请在程序关闭后手动前往 {store_dir} 删除目录。\n {str(e)}")
|
2025-02-04 13:23:50 +08:00
|
|
|
|
traceback.print_exc()
|
2025-02-04 16:11:17 +08:00
|
|
|
|
return False
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-02 12:59:13 +08:00
|
|
|
|
def init_vector_store(
|
2025-02-02 14:39:30 +08:00
|
|
|
|
api_key: str,
|
2025-02-03 13:00:03 +08:00
|
|
|
|
base_url: str,
|
2025-02-02 14:39:30 +08:00
|
|
|
|
interface_format: str,
|
|
|
|
|
|
embedding_model_name: str,
|
2025-02-03 13:00:03 +08:00
|
|
|
|
texts: List[str],
|
2025-02-04 00:10:19 +08:00
|
|
|
|
filepath: str
|
2025-02-02 14:39:30 +08:00
|
|
|
|
) -> Chroma:
|
2025-01-29 21:59:36 +08:00
|
|
|
|
"""
|
2025-02-04 00:10:19 +08:00
|
|
|
|
在 filepath 下创建/加载一个 Chroma 向量库并插入 texts。
|
2025-01-29 21:59:36 +08:00
|
|
|
|
"""
|
2025-02-04 00:10:19 +08:00
|
|
|
|
store_dir = get_vectorstore_dir(filepath)
|
|
|
|
|
|
os.makedirs(store_dir, exist_ok=True)
|
|
|
|
|
|
|
2025-02-02 14:39:30 +08:00
|
|
|
|
embeddings = create_embeddings_object(
|
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
|
base_url=base_url,
|
|
|
|
|
|
interface_format=interface_format,
|
|
|
|
|
|
embedding_model_name=embedding_model_name
|
|
|
|
|
|
)
|
2025-02-04 00:10:19 +08:00
|
|
|
|
documents = [Document(page_content=str(t)) for t in texts]
|
2025-01-29 21:59:36 +08:00
|
|
|
|
vectorstore = Chroma.from_documents(
|
|
|
|
|
|
documents,
|
|
|
|
|
|
embedding=embeddings,
|
2025-02-04 13:23:50 +08:00
|
|
|
|
persist_directory=store_dir,
|
|
|
|
|
|
client_settings=Settings(anonymized_telemetry=False),
|
|
|
|
|
|
collection_name="novel_collection"
|
2025-01-29 21:59:36 +08:00
|
|
|
|
)
|
|
|
|
|
|
return vectorstore
|
|
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-02 12:59:13 +08:00
|
|
|
|
def load_vector_store(
|
2025-02-02 14:39:30 +08:00
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
interface_format: str,
|
|
|
|
|
|
embedding_model_name: str,
|
2025-02-04 00:10:19 +08:00
|
|
|
|
filepath: str
|
2025-02-02 14:39:30 +08:00
|
|
|
|
) -> Optional[Chroma]:
|
2025-02-02 11:55:49 +08:00
|
|
|
|
"""
|
2025-02-04 00:10:19 +08:00
|
|
|
|
读取已存在的 Chroma 向量库。若不存在则返回 None。
|
2025-02-02 11:55:49 +08:00
|
|
|
|
"""
|
2025-02-04 00:10:19 +08:00
|
|
|
|
store_dir = get_vectorstore_dir(filepath)
|
|
|
|
|
|
if not os.path.exists(store_dir):
|
2025-02-03 13:00:03 +08:00
|
|
|
|
logging.info("Vector store not found. Will return None.")
|
2025-01-29 21:59:36 +08:00
|
|
|
|
return None
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-02 14:39:30 +08:00
|
|
|
|
embeddings = create_embeddings_object(
|
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
|
base_url=base_url,
|
|
|
|
|
|
interface_format=interface_format,
|
|
|
|
|
|
embedding_model_name=embedding_model_name
|
|
|
|
|
|
)
|
2025-02-04 00:10:19 +08:00
|
|
|
|
return Chroma(
|
|
|
|
|
|
persist_directory=store_dir,
|
2025-02-04 13:23:50 +08:00
|
|
|
|
embedding_function=embeddings,
|
|
|
|
|
|
client_settings=Settings(anonymized_telemetry=False),
|
|
|
|
|
|
collection_name="novel_collection"
|
2025-02-04 00:10:19 +08:00
|
|
|
|
)
|
2025-01-29 21:59:36 +08:00
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-04 16:11:17 +08:00
|
|
|
|
def split_by_length(text: str, max_length: int = 500) -> List[str]:
|
|
|
|
|
|
segments = []
|
|
|
|
|
|
start_idx = 0
|
|
|
|
|
|
while start_idx < len(text):
|
|
|
|
|
|
end_idx = min(start_idx + max_length, len(text))
|
|
|
|
|
|
segment = text[start_idx:end_idx]
|
|
|
|
|
|
segments.append(segment.strip())
|
|
|
|
|
|
start_idx = end_idx
|
|
|
|
|
|
return segments
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def split_text_for_vectorstore(chapter_text: str,
|
|
|
|
|
|
max_length: int = 500,
|
|
|
|
|
|
similarity_threshold: float = 0.7) -> List[str]:
|
|
|
|
|
|
"""
|
|
|
|
|
|
对新的章节文本进行分段后,再用于存入向量库。
|
|
|
|
|
|
"""
|
|
|
|
|
|
if not chapter_text.strip():
|
|
|
|
|
|
return []
|
|
|
|
|
|
|
|
|
|
|
|
nltk.download('punkt', quiet=True)
|
2025-02-04 16:37:29 +08:00
|
|
|
|
nltk.download('punkt_tab', quiet=True)
|
2025-02-04 16:11:17 +08:00
|
|
|
|
sentences = nltk.sent_tokenize(chapter_text)
|
|
|
|
|
|
if not sentences:
|
|
|
|
|
|
return []
|
|
|
|
|
|
|
|
|
|
|
|
# 先对相近句子进行合并
|
|
|
|
|
|
model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
|
|
|
|
|
|
embeddings = model.encode(sentences)
|
|
|
|
|
|
|
|
|
|
|
|
merged_paragraphs = []
|
|
|
|
|
|
current_sentences = [sentences[0]]
|
|
|
|
|
|
current_embedding = embeddings[0]
|
|
|
|
|
|
|
|
|
|
|
|
for i in range(1, len(sentences)):
|
|
|
|
|
|
sim = cosine_similarity([current_embedding], [embeddings[i]])[0][0]
|
|
|
|
|
|
if sim >= similarity_threshold:
|
|
|
|
|
|
current_sentences.append(sentences[i])
|
|
|
|
|
|
current_embedding = (current_embedding + embeddings[i]) / 2.0
|
|
|
|
|
|
else:
|
|
|
|
|
|
merged_paragraphs.append(" ".join(current_sentences))
|
|
|
|
|
|
current_sentences = [sentences[i]]
|
|
|
|
|
|
current_embedding = embeddings[i]
|
|
|
|
|
|
|
|
|
|
|
|
if current_sentences:
|
|
|
|
|
|
merged_paragraphs.append(" ".join(current_sentences))
|
|
|
|
|
|
|
|
|
|
|
|
# 再对合并好的段落做 max_length 切分
|
|
|
|
|
|
final_segments = []
|
|
|
|
|
|
for para in merged_paragraphs:
|
|
|
|
|
|
if len(para) > max_length:
|
|
|
|
|
|
sub_segments = split_by_length(para, max_length=max_length)
|
|
|
|
|
|
final_segments.extend(sub_segments)
|
|
|
|
|
|
else:
|
|
|
|
|
|
final_segments.append(para)
|
|
|
|
|
|
|
|
|
|
|
|
return final_segments
|
|
|
|
|
|
|
|
|
|
|
|
|
2025-02-02 12:59:13 +08:00
|
|
|
|
def update_vector_store(
|
2025-02-03 13:00:03 +08:00
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
new_chapter: str,
|
2025-02-02 15:50:19 +08:00
|
|
|
|
interface_format: str,
|
|
|
|
|
|
embedding_model_name: str,
|
2025-02-04 00:10:19 +08:00
|
|
|
|
filepath: str
|
|
|
|
|
|
):
|
2025-02-02 14:39:30 +08:00
|
|
|
|
"""
|
2025-02-04 00:10:19 +08:00
|
|
|
|
将最新章节文本插入到向量库中。若库不存在则初始化。
|
2025-02-02 14:39:30 +08:00
|
|
|
|
"""
|
2025-02-04 16:11:17 +08:00
|
|
|
|
splitted_texts = split_text_for_vectorstore(new_chapter)
|
|
|
|
|
|
if not splitted_texts:
|
|
|
|
|
|
logging.warning("No valid text to insert into vector store. Skipping.")
|
|
|
|
|
|
return
|
|
|
|
|
|
|
2025-02-02 14:39:30 +08:00
|
|
|
|
store = load_vector_store(
|
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
|
base_url=base_url,
|
|
|
|
|
|
interface_format=interface_format,
|
|
|
|
|
|
embedding_model_name=embedding_model_name,
|
2025-02-04 00:10:19 +08:00
|
|
|
|
filepath=filepath
|
2025-02-02 14:39:30 +08:00
|
|
|
|
)
|
2025-01-29 21:59:36 +08:00
|
|
|
|
if not store:
|
2025-02-02 18:25:35 +08:00
|
|
|
|
logging.info("Vector store does not exist. Initializing a new one for new chapter...")
|
2025-02-02 14:39:30 +08:00
|
|
|
|
init_vector_store(
|
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
|
base_url=base_url,
|
|
|
|
|
|
interface_format=interface_format,
|
|
|
|
|
|
embedding_model_name=embedding_model_name,
|
2025-02-04 16:11:17 +08:00
|
|
|
|
texts=splitted_texts,
|
2025-02-04 00:10:19 +08:00
|
|
|
|
filepath=filepath
|
2025-02-02 14:39:30 +08:00
|
|
|
|
)
|
2025-01-29 21:59:36 +08:00
|
|
|
|
return
|
|
|
|
|
|
|
2025-02-04 16:11:17 +08:00
|
|
|
|
docs = [Document(page_content=str(t)) for t in splitted_texts]
|
|
|
|
|
|
store.add_documents(docs)
|
|
|
|
|
|
logging.info("Vector store updated with the new chapter splitted segments.")
|
2025-01-29 21:59:36 +08:00
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-02 12:59:13 +08:00
|
|
|
|
def get_relevant_context_from_vector_store(
|
2025-02-02 14:39:30 +08:00
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
query: str,
|
2025-02-02 15:50:19 +08:00
|
|
|
|
interface_format: str,
|
|
|
|
|
|
embedding_model_name: str,
|
2025-02-04 00:10:19 +08:00
|
|
|
|
filepath: str,
|
2025-02-02 14:39:30 +08:00
|
|
|
|
k: int = 2
|
|
|
|
|
|
) -> str:
|
2025-01-29 21:59:36 +08:00
|
|
|
|
"""
|
|
|
|
|
|
从向量库中检索与 query 最相关的 k 条文本,拼接后返回。
|
|
|
|
|
|
"""
|
2025-02-02 14:39:30 +08:00
|
|
|
|
store = load_vector_store(
|
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
|
base_url=base_url,
|
|
|
|
|
|
interface_format=interface_format,
|
|
|
|
|
|
embedding_model_name=embedding_model_name,
|
2025-02-04 00:10:19 +08:00
|
|
|
|
filepath=filepath
|
2025-02-02 14:39:30 +08:00
|
|
|
|
)
|
2025-01-29 21:59:36 +08:00
|
|
|
|
if not store:
|
2025-02-02 18:25:35 +08:00
|
|
|
|
logging.info("No vector store found. Returning empty context.")
|
2025-01-29 21:59:36 +08:00
|
|
|
|
return ""
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-01-29 21:59:36 +08:00
|
|
|
|
docs = store.similarity_search(query, k=k)
|
2025-02-02 18:25:35 +08:00
|
|
|
|
if not docs:
|
|
|
|
|
|
logging.info(f"No relevant documents found for query '{query}'. Returning empty context.")
|
|
|
|
|
|
return ""
|
|
|
|
|
|
|
2025-01-29 21:59:36 +08:00
|
|
|
|
combined = "\n".join([d.page_content for d in docs])
|
|
|
|
|
|
return combined
|
|
|
|
|
|
|
2025-01-29 20:33:20 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# ========== 1) 生成总体架构 (Novel_architecture.txt) ==========
|
|
|
|
|
|
def Novel_architecture_generate(
|
2025-01-29 20:33:20 +08:00
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
llm_model: str,
|
|
|
|
|
|
topic: str,
|
|
|
|
|
|
genre: str,
|
|
|
|
|
|
number_of_chapters: int,
|
|
|
|
|
|
word_number: int,
|
2025-01-31 13:50:07 +08:00
|
|
|
|
filepath: str,
|
|
|
|
|
|
temperature: float = 0.7
|
2025-01-29 21:59:36 +08:00
|
|
|
|
) -> None:
|
2025-02-05 21:55:48 +08:00
|
|
|
|
"""
|
|
|
|
|
|
依次调用:
|
|
|
|
|
|
1. core_seed_prompt
|
|
|
|
|
|
2. character_dynamics_prompt
|
|
|
|
|
|
3. world_building_prompt
|
|
|
|
|
|
4. plot_architecture_prompt
|
|
|
|
|
|
将结果整合为“Novel_architecture.txt”。
|
|
|
|
|
|
"""
|
2025-01-29 21:59:36 +08:00
|
|
|
|
os.makedirs(filepath, exist_ok=True)
|
2025-01-29 20:33:20 +08:00
|
|
|
|
model = ChatOpenAI(
|
|
|
|
|
|
model=llm_model,
|
|
|
|
|
|
api_key=api_key,
|
2025-02-04 00:10:19 +08:00
|
|
|
|
base_url=ensure_openai_base_url_has_v1(base_url),
|
2025-01-31 13:50:07 +08:00
|
|
|
|
temperature=temperature
|
2025-01-29 20:33:20 +08:00
|
|
|
|
)
|
|
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# 1) 核心种子
|
|
|
|
|
|
prompt_core = core_seed_prompt.format(
|
2025-02-03 13:00:03 +08:00
|
|
|
|
topic=topic,
|
|
|
|
|
|
genre=genre,
|
|
|
|
|
|
number_of_chapters=number_of_chapters,
|
|
|
|
|
|
word_number=word_number
|
|
|
|
|
|
)
|
2025-02-05 21:55:48 +08:00
|
|
|
|
core_seed_result = invoke_with_cleaning(model, prompt_core)
|
|
|
|
|
|
core_seed_text = core_seed_result.strip()
|
2025-01-29 20:33:20 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# 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
|
2025-02-03 13:00:03 +08:00
|
|
|
|
)
|
2025-02-05 21:55:48 +08:00
|
|
|
|
plot_arch_result = invoke_with_cleaning(model, prompt_plot)
|
|
|
|
|
|
plot_arch_text = plot_arch_result.strip()
|
2025-01-29 20:33:20 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# 整合并写入 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"
|
2025-02-03 13:00:03 +08:00
|
|
|
|
)
|
2025-01-29 20:33:20 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
arch_file = os.path.join(filepath, "Novel_architecture.txt")
|
|
|
|
|
|
clear_file_content(arch_file)
|
|
|
|
|
|
save_string_to_txt(final_content, arch_file)
|
2025-01-29 20:33:20 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
logging.info("Novel_architecture.txt has been generated successfully.")
|
2025-02-02 18:35:10 +08:00
|
|
|
|
|
2025-02-02 14:39:30 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# ========== 2) 生成章节蓝图 (Novel_directory.txt) ==========
|
|
|
|
|
|
def Chapter_blueprint_generate(
|
2025-02-03 13:00:03 +08:00
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
llm_model: str,
|
|
|
|
|
|
filepath: str,
|
|
|
|
|
|
temperature: float = 0.7
|
|
|
|
|
|
) -> None:
|
2025-02-05 21:55:48 +08:00
|
|
|
|
"""
|
|
|
|
|
|
基于“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.")
|
2025-02-03 13:00:03 +08:00
|
|
|
|
return
|
|
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
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()
|
|
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
model = ChatOpenAI(
|
|
|
|
|
|
model=llm_model,
|
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
|
base_url=ensure_openai_base_url_has_v1(base_url),
|
|
|
|
|
|
temperature=temperature
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
prompt = chapter_blueprint_prompt.format(
|
|
|
|
|
|
plot_architecture=plot_arch_text,
|
2025-02-03 13:00:03 +08:00
|
|
|
|
number_of_chapters=number_of_chapters
|
|
|
|
|
|
)
|
2025-02-05 21:55:48 +08:00
|
|
|
|
blueprint_text = invoke_with_cleaning(model, prompt)
|
|
|
|
|
|
if not blueprint_text.strip():
|
|
|
|
|
|
logging.warning("Chapter blueprint generation result is empty.")
|
2025-02-03 13:00:03 +08:00
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
|
|
filename_dir = os.path.join(filepath, "Novel_directory.txt")
|
|
|
|
|
|
clear_file_content(filename_dir)
|
2025-02-05 21:55:48 +08:00
|
|
|
|
save_string_to_txt(blueprint_text, filename_dir)
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
logging.info("Novel_directory.txt (chapter blueprint) has been generated successfully.")
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# ============ 获取最近 N 章内容,生成短期摘要 ============
|
2025-01-31 19:57:44 +08:00
|
|
|
|
def get_last_n_chapters_text(chapters_dir: str, current_chapter_num: int, n: int = 3) -> List[str]:
|
|
|
|
|
|
texts = []
|
|
|
|
|
|
start_chap = max(1, current_chapter_num - n)
|
|
|
|
|
|
for c in range(start_chap, current_chapter_num):
|
|
|
|
|
|
chap_file = os.path.join(chapters_dir, f"chapter_{c}.txt")
|
|
|
|
|
|
if os.path.exists(chap_file):
|
|
|
|
|
|
text = read_file(chap_file).strip()
|
|
|
|
|
|
if text:
|
|
|
|
|
|
texts.append(text)
|
2025-02-02 18:25:35 +08:00
|
|
|
|
if len(texts) < n:
|
|
|
|
|
|
texts = [''] * (n - len(texts)) + texts
|
2025-01-31 19:57:44 +08:00
|
|
|
|
return texts
|
2025-01-31 13:50:07 +08:00
|
|
|
|
|
2025-02-02 15:50:19 +08:00
|
|
|
|
def summarize_recent_chapters(
|
2025-02-02 18:25:35 +08:00
|
|
|
|
llm_model: str,
|
|
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
temperature: float,
|
|
|
|
|
|
chapters_text_list: List[str]
|
|
|
|
|
|
) -> str:
|
|
|
|
|
|
if not chapters_text_list:
|
|
|
|
|
|
return ""
|
2025-02-02 18:35:10 +08:00
|
|
|
|
if all(not txt.strip() for txt in chapters_text_list):
|
2025-02-02 18:25:35 +08:00
|
|
|
|
return "暂无摘要。"
|
2025-02-02 18:35:10 +08:00
|
|
|
|
|
2025-02-02 15:50:19 +08:00
|
|
|
|
model = ChatOpenAI(
|
|
|
|
|
|
model=llm_model,
|
|
|
|
|
|
api_key=api_key,
|
2025-02-03 13:00:03 +08:00
|
|
|
|
base_url=ensure_openai_base_url_has_v1(base_url),
|
2025-02-02 15:50:19 +08:00
|
|
|
|
temperature=temperature
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2025-01-31 19:57:44 +08:00
|
|
|
|
combined_text = "\n".join(chapters_text_list)
|
2025-02-02 18:25:35 +08:00
|
|
|
|
prompt = f"""你是一名资深长篇小说写作辅助AI,下面是最近几章的合并文本:
|
2025-02-02 14:39:30 +08:00
|
|
|
|
{combined_text}
|
|
|
|
|
|
|
2025-02-02 18:25:35 +08:00
|
|
|
|
请用中文输出不超过500字的摘要,只包含主要剧情进展、角色变化、冲突焦点等要点:"""
|
|
|
|
|
|
|
2025-02-02 19:17:07 +08:00
|
|
|
|
summary_text = invoke_with_cleaning(model, prompt)
|
|
|
|
|
|
if not summary_text:
|
2025-02-03 13:00:03 +08:00
|
|
|
|
return (combined_text[:800] + "...") if len(combined_text) > 800 else combined_text
|
2025-02-02 19:17:07 +08:00
|
|
|
|
return summary_text
|
2025-02-02 14:39:30 +08:00
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-04 00:10:19 +08:00
|
|
|
|
# ============ 剧情要点/冲突 ============
|
2025-02-01 18:49:52 +08:00
|
|
|
|
PLOT_ARCS_PROMPT = """\
|
|
|
|
|
|
下面是新生成的章节内容:
|
|
|
|
|
|
{chapter_text}
|
|
|
|
|
|
|
|
|
|
|
|
这里是已记录的剧情要点/未解决冲突(可能为空):
|
|
|
|
|
|
{old_plot_arcs}
|
|
|
|
|
|
|
2025-02-02 18:25:35 +08:00
|
|
|
|
请基于新的章节内容,提炼本章引入或延续的悬念、冲突、角色暗线等,将其合并到旧的剧情要点中。
|
2025-02-01 18:49:52 +08:00
|
|
|
|
若有新的冲突则添加,若有已解决/不再重要的冲突可标注或移除。
|
2025-02-02 18:25:35 +08:00
|
|
|
|
最终输出更新后的剧情要点列表,以帮助后续保持故事整体的一致性和悬念延续。
|
2025-02-01 18:49:52 +08:00
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
def update_plot_arcs(
|
|
|
|
|
|
chapter_text: str,
|
|
|
|
|
|
old_plot_arcs: str,
|
|
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
model_name: str,
|
|
|
|
|
|
temperature: float
|
|
|
|
|
|
) -> str:
|
|
|
|
|
|
model = ChatOpenAI(
|
|
|
|
|
|
model=model_name,
|
|
|
|
|
|
api_key=api_key,
|
2025-02-03 13:00:03 +08:00
|
|
|
|
base_url=ensure_openai_base_url_has_v1(base_url),
|
2025-02-01 18:49:52 +08:00
|
|
|
|
temperature=temperature
|
|
|
|
|
|
)
|
|
|
|
|
|
prompt = PLOT_ARCS_PROMPT.format(
|
|
|
|
|
|
chapter_text=chapter_text,
|
|
|
|
|
|
old_plot_arcs=old_plot_arcs
|
|
|
|
|
|
)
|
2025-02-02 19:17:07 +08:00
|
|
|
|
arcs_text = invoke_with_cleaning(model, prompt)
|
|
|
|
|
|
if not arcs_text:
|
|
|
|
|
|
logging.warning("update_plot_arcs: No response or empty result.")
|
2025-02-01 18:49:52 +08:00
|
|
|
|
return old_plot_arcs
|
2025-02-02 19:17:07 +08:00
|
|
|
|
return arcs_text
|
2025-02-01 18:49:52 +08:00
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# ========== 3) 生成章节草稿 ==========
|
|
|
|
|
|
|
2025-01-31 19:57:44 +08:00
|
|
|
|
def generate_chapter_draft(
|
2025-01-29 20:33:20 +08:00
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
model_name: str,
|
2025-02-05 21:55:48 +08:00
|
|
|
|
filepath: str,
|
2025-01-29 20:33:20 +08:00
|
|
|
|
novel_number: int,
|
|
|
|
|
|
word_number: int,
|
2025-01-31 19:57:44 +08:00
|
|
|
|
temperature: float,
|
2025-02-05 21:55:48 +08:00
|
|
|
|
user_guidance: str,
|
|
|
|
|
|
characters_involved: str,
|
|
|
|
|
|
key_items: str,
|
|
|
|
|
|
scene_location: str,
|
|
|
|
|
|
time_constraint: str,
|
|
|
|
|
|
embedding_retrieval_k: int = 2
|
2025-01-29 20:33:20 +08:00
|
|
|
|
) -> str:
|
2025-02-05 21:55:48 +08:00
|
|
|
|
"""
|
|
|
|
|
|
根据 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)
|
2025-01-31 20:39:05 +08:00
|
|
|
|
chapter_title = chapter_info["chapter_title"]
|
2025-02-05 21:55:48 +08:00
|
|
|
|
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"]
|
2025-01-31 20:39:05 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# 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
|
2025-02-02 18:25:35 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# 4) 检索向量库上下文
|
2025-02-04 16:11:17 +08:00
|
|
|
|
relevant_context = get_relevant_context_from_vector_store(
|
|
|
|
|
|
api_key=api_key,
|
2025-02-05 21:55:48 +08:00
|
|
|
|
base_url=base_url,
|
2025-02-04 16:11:17 +08:00
|
|
|
|
query=merged_query_str,
|
2025-02-05 21:55:48 +08:00
|
|
|
|
embedding_model_name=model_name,
|
2025-02-04 16:11:17 +08:00
|
|
|
|
filepath=filepath,
|
|
|
|
|
|
k=embedding_retrieval_k
|
|
|
|
|
|
)
|
2025-01-31 19:57:44 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
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"
|
|
|
|
|
|
|
2025-01-29 20:33:20 +08:00
|
|
|
|
model = ChatOpenAI(
|
|
|
|
|
|
model=model_name,
|
|
|
|
|
|
api_key=api_key,
|
2025-02-03 13:00:03 +08:00
|
|
|
|
base_url=ensure_openai_base_url_has_v1(base_url),
|
2025-01-31 13:50:07 +08:00
|
|
|
|
temperature=temperature
|
2025-01-29 20:33:20 +08:00
|
|
|
|
)
|
|
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
chapter_content = invoke_with_cleaning(model, prompt_text)
|
|
|
|
|
|
if not chapter_content.strip():
|
|
|
|
|
|
logging.warning("Generated chapter draft is empty.")
|
2025-01-31 19:57:44 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# 6) 写入 chapters 目录
|
2025-01-29 21:59:36 +08:00
|
|
|
|
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")
|
2025-02-05 21:55:48 +08:00
|
|
|
|
|
2025-01-31 19:57:44 +08:00
|
|
|
|
clear_file_content(chapter_file)
|
|
|
|
|
|
save_string_to_txt(chapter_content, chapter_file)
|
|
|
|
|
|
|
|
|
|
|
|
logging.info(f"[Draft] Chapter {novel_number} generated as a draft.")
|
|
|
|
|
|
return chapter_content
|
|
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# ========== 4) 定稿章节 ==========
|
2025-01-31 19:57:44 +08:00
|
|
|
|
def finalize_chapter(
|
|
|
|
|
|
novel_number: int,
|
|
|
|
|
|
word_number: int,
|
|
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
model_name: str,
|
|
|
|
|
|
temperature: float,
|
2025-02-04 00:10:19 +08:00
|
|
|
|
filepath: str,
|
2025-02-05 21:55:48 +08:00
|
|
|
|
embedding_model_name: str
|
2025-01-31 19:57:44 +08:00
|
|
|
|
):
|
2025-02-05 21:55:48 +08:00
|
|
|
|
"""
|
|
|
|
|
|
定稿:更新全局摘要、角色状态,并将本章文本插入向量库。
|
|
|
|
|
|
"""
|
2025-01-31 19:57:44 +08:00
|
|
|
|
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()
|
|
|
|
|
|
if not chapter_text:
|
|
|
|
|
|
logging.warning(f"Chapter {novel_number} is empty, cannot finalize.")
|
|
|
|
|
|
return
|
|
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# 如果长度比目标少很多,可考虑在此扩写
|
|
|
|
|
|
if len(chapter_text) < 0.6 * word_number:
|
|
|
|
|
|
chapter_text = enrich_chapter_text(chapter_text, word_number, api_key, base_url, model_name, temperature)
|
2025-01-31 19:57:44 +08:00
|
|
|
|
clear_file_content(chapter_file)
|
|
|
|
|
|
save_string_to_txt(chapter_text, chapter_file)
|
|
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# 读取全局摘要、角色状态
|
|
|
|
|
|
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) 更新全局摘要
|
2025-01-31 19:57:44 +08:00
|
|
|
|
model = ChatOpenAI(
|
|
|
|
|
|
model=model_name,
|
|
|
|
|
|
api_key=api_key,
|
2025-02-03 13:00:03 +08:00
|
|
|
|
base_url=ensure_openai_base_url_has_v1(base_url),
|
2025-01-31 19:57:44 +08:00
|
|
|
|
temperature=temperature
|
|
|
|
|
|
)
|
2025-02-05 21:55:48 +08:00
|
|
|
|
prompt_summary = summary_prompt.format(
|
2025-02-01 18:49:52 +08:00
|
|
|
|
chapter_text=chapter_text,
|
2025-02-05 21:55:48 +08:00
|
|
|
|
global_summary=old_global_summary
|
2025-02-01 18:49:52 +08:00
|
|
|
|
)
|
2025-02-05 21:55:48 +08:00
|
|
|
|
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
|
2025-02-01 18:49:52 +08:00
|
|
|
|
|
2025-02-02 18:25:35 +08:00
|
|
|
|
# 写回文件
|
2025-01-31 19:57:44 +08:00
|
|
|
|
clear_file_content(global_summary_file)
|
|
|
|
|
|
save_string_to_txt(new_global_summary, global_summary_file)
|
|
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
clear_file_content(character_state_file)
|
|
|
|
|
|
save_string_to_txt(new_char_state, character_state_file)
|
2025-02-01 18:49:52 +08:00
|
|
|
|
|
2025-02-05 21:55:48 +08:00
|
|
|
|
# 3) 更新向量库
|
2025-02-02 14:39:30 +08:00
|
|
|
|
update_vector_store(
|
2025-02-05 21:55:48 +08:00
|
|
|
|
api_key=api_key,
|
|
|
|
|
|
base_url=base_url,
|
2025-02-02 14:39:30 +08:00
|
|
|
|
new_chapter=chapter_text,
|
2025-02-05 21:55:48 +08:00
|
|
|
|
model_name=embedding_model_name, # 用于embedding
|
2025-02-04 00:10:19 +08:00
|
|
|
|
filepath=filepath
|
2025-02-02 14:39:30 +08:00
|
|
|
|
)
|
2025-01-31 19:57:44 +08:00
|
|
|
|
|
2025-02-02 14:39:30 +08:00
|
|
|
|
logging.info(f"Chapter {novel_number} has been finalized.")
|
2025-01-31 19:57:44 +08:00
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-01-31 19:57:44 +08:00
|
|
|
|
def enrich_chapter_text(
|
|
|
|
|
|
chapter_text: str,
|
|
|
|
|
|
word_number: int,
|
|
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
model_name: str,
|
|
|
|
|
|
temperature: float
|
|
|
|
|
|
) -> str:
|
|
|
|
|
|
model = ChatOpenAI(
|
|
|
|
|
|
model=model_name,
|
|
|
|
|
|
api_key=api_key,
|
2025-02-03 13:00:03 +08:00
|
|
|
|
base_url=ensure_openai_base_url_has_v1(base_url),
|
2025-01-31 19:57:44 +08:00
|
|
|
|
temperature=temperature
|
2025-01-29 21:59:36 +08:00
|
|
|
|
)
|
2025-02-02 14:39:30 +08:00
|
|
|
|
prompt = f"""以下是当前章节文本,可能篇幅较短,请在保持剧情连贯的前提下进行扩写,使其更充实、生动,并尽量靠近目标 {word_number} 字数。
|
2025-01-29 20:33:20 +08:00
|
|
|
|
|
2025-01-31 19:57:44 +08:00
|
|
|
|
原章节内容:
|
2025-02-02 14:39:30 +08:00
|
|
|
|
{chapter_text}"""
|
2025-02-02 19:17:07 +08:00
|
|
|
|
enriched_text = invoke_with_cleaning(model, prompt)
|
|
|
|
|
|
return enriched_text if enriched_text else chapter_text
|
2025-01-31 13:50:07 +08:00
|
|
|
|
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-04 00:10:19 +08:00
|
|
|
|
# ============ 导入外部知识文本到向量库 ============
|
2025-01-31 13:50:07 +08:00
|
|
|
|
def advanced_split_content(content: str,
|
|
|
|
|
|
similarity_threshold: float = 0.7,
|
|
|
|
|
|
max_length: int = 500) -> List[str]:
|
|
|
|
|
|
"""
|
2025-02-03 13:00:03 +08:00
|
|
|
|
将文本先按句子切分,然后根据语义相似度进行合并,最后按 max_length 二次切分。
|
2025-01-31 13:50:07 +08:00
|
|
|
|
"""
|
2025-02-04 16:11:17 +08:00
|
|
|
|
nltk.download('punkt', quiet=True)
|
2025-02-02 18:25:35 +08:00
|
|
|
|
sentences = nltk.sent_tokenize(content)
|
2025-01-31 13:50:07 +08:00
|
|
|
|
if not sentences:
|
|
|
|
|
|
return []
|
|
|
|
|
|
|
|
|
|
|
|
model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
|
|
|
|
|
|
embeddings = model.encode(sentences)
|
|
|
|
|
|
|
|
|
|
|
|
merged_paragraphs = []
|
|
|
|
|
|
current_sentences = [sentences[0]]
|
|
|
|
|
|
current_embedding = embeddings[0]
|
|
|
|
|
|
|
|
|
|
|
|
for i in range(1, len(sentences)):
|
|
|
|
|
|
sim = cosine_similarity([current_embedding], [embeddings[i]])[0][0]
|
|
|
|
|
|
if sim >= similarity_threshold:
|
|
|
|
|
|
current_sentences.append(sentences[i])
|
|
|
|
|
|
current_embedding = (current_embedding + embeddings[i]) / 2.0
|
|
|
|
|
|
else:
|
|
|
|
|
|
merged_paragraphs.append(" ".join(current_sentences))
|
|
|
|
|
|
current_sentences = [sentences[i]]
|
|
|
|
|
|
current_embedding = embeddings[i]
|
|
|
|
|
|
|
|
|
|
|
|
if current_sentences:
|
|
|
|
|
|
merged_paragraphs.append(" ".join(current_sentences))
|
|
|
|
|
|
|
|
|
|
|
|
final_segments = []
|
|
|
|
|
|
for para in merged_paragraphs:
|
|
|
|
|
|
if len(para) > max_length:
|
|
|
|
|
|
sub_segments = split_by_length(para, max_length=max_length)
|
|
|
|
|
|
final_segments.extend(sub_segments)
|
|
|
|
|
|
else:
|
|
|
|
|
|
final_segments.append(para)
|
|
|
|
|
|
|
|
|
|
|
|
return final_segments
|
|
|
|
|
|
|
2025-02-04 16:11:17 +08:00
|
|
|
|
def import_knowledge_file(
|
|
|
|
|
|
api_key: str,
|
|
|
|
|
|
base_url: str,
|
|
|
|
|
|
interface_format: str,
|
|
|
|
|
|
embedding_model_name: str,
|
|
|
|
|
|
file_path: str,
|
|
|
|
|
|
embedding_base_url: str,
|
|
|
|
|
|
filepath: str
|
|
|
|
|
|
):
|
|
|
|
|
|
logging.info(f"开始导入知识库文件: {file_path}, 接口格式: {interface_format}, 模型: {embedding_model_name}")
|
|
|
|
|
|
if not os.path.exists(file_path):
|
|
|
|
|
|
logging.warning(f"知识库文件不存在: {file_path}")
|
|
|
|
|
|
return
|
2025-02-03 13:00:03 +08:00
|
|
|
|
|
2025-02-04 16:11:17 +08:00
|
|
|
|
content = read_file(file_path)
|
|
|
|
|
|
if not content.strip():
|
|
|
|
|
|
logging.warning("知识库文件内容为空。")
|
|
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
|
|
paragraphs = advanced_split_content(content)
|
|
|
|
|
|
|
|
|
|
|
|
# 若向量库不存在则初始化,否则追加
|
|
|
|
|
|
store = load_vector_store(
|
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
|
base_url=base_url if base_url else "http://localhost:11434/v1",
|
|
|
|
|
|
interface_format=interface_format,
|
|
|
|
|
|
embedding_model_name=embedding_model_name,
|
|
|
|
|
|
filepath=filepath
|
|
|
|
|
|
)
|
|
|
|
|
|
if not store:
|
|
|
|
|
|
logging.info("Vector store does not exist. Initializing a new one for knowledge import...")
|
|
|
|
|
|
init_vector_store(
|
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
|
base_url=base_url if base_url else "http://localhost:11434/v1",
|
|
|
|
|
|
interface_format=interface_format,
|
|
|
|
|
|
embedding_model_name=embedding_model_name,
|
|
|
|
|
|
texts=paragraphs,
|
|
|
|
|
|
filepath=filepath
|
|
|
|
|
|
)
|
|
|
|
|
|
else:
|
|
|
|
|
|
docs = [Document(page_content=str(p)) for p in paragraphs]
|
|
|
|
|
|
store.add_documents(docs)
|
|
|
|
|
|
logging.info("知识库文件已成功导入至向量库。")
|