@@ -0,0 +1,55 @@
|
|||||||
|
# embedding_ollama.py
|
||||||
|
import requests
|
||||||
|
import traceback
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
class OllamaEmbeddings:
|
||||||
|
"""
|
||||||
|
Ollama 本地服务提供的 Embedding 接口,
|
||||||
|
最终拼出形如: http://localhost:11434/api/embed
|
||||||
|
即 base_url + "/embed"
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, model_name: str, base_url: str):
|
||||||
|
self.model_name = model_name
|
||||||
|
self.base_url = base_url
|
||||||
|
|
||||||
|
def embed(self, texts: List[str]) -> List[List[float]]:
|
||||||
|
"""
|
||||||
|
批量将多段文本转换为embedding向量
|
||||||
|
"""
|
||||||
|
embeddings = []
|
||||||
|
for text in texts:
|
||||||
|
embeddings.append(self.embed_single_document(text))
|
||||||
|
return embeddings
|
||||||
|
|
||||||
|
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
||||||
|
"""
|
||||||
|
兼容langchain的接口写法
|
||||||
|
"""
|
||||||
|
return self.embed(texts)
|
||||||
|
|
||||||
|
def embed_query(self, query: str) -> List[float]:
|
||||||
|
"""
|
||||||
|
将单条 query 转换为 embedding 向量
|
||||||
|
"""
|
||||||
|
return self.embed_single_document(query)
|
||||||
|
|
||||||
|
def embed_single_document(self, text: str) -> List[float]:
|
||||||
|
"""
|
||||||
|
调用 Ollama 本地服务接口,获取文本的 embedding。
|
||||||
|
"""
|
||||||
|
url = f"{self.base_url}/embed"
|
||||||
|
data = {
|
||||||
|
"model": self.model_name,
|
||||||
|
"prompt": text
|
||||||
|
}
|
||||||
|
try:
|
||||||
|
response = requests.post(url, json=data)
|
||||||
|
response.raise_for_status()
|
||||||
|
result = response.json()
|
||||||
|
if "embedding" not in result:
|
||||||
|
raise ValueError("No 'embedding' field in Ollama response.")
|
||||||
|
return result["embedding"]
|
||||||
|
except requests.exceptions.RequestException as e:
|
||||||
|
raise Exception(f"Ollama embeddings request error: {e}\n{traceback.format_exc()}")
|
||||||
@@ -40,7 +40,7 @@ exe = EXE(
|
|||||||
a.scripts,
|
a.scripts,
|
||||||
[],
|
[],
|
||||||
exclude_binaries=True,
|
exclude_binaries=True,
|
||||||
name='AI_NovelGenerator_V1.2.3',
|
name='AI_NovelGenerator_V1.2.4',
|
||||||
debug=True,
|
debug=True,
|
||||||
bootloader_ignore_signals=False,
|
bootloader_ignore_signals=False,
|
||||||
strip=False,
|
strip=False,
|
||||||
@@ -60,5 +60,5 @@ coll = COLLECT(
|
|||||||
strip=False,
|
strip=False,
|
||||||
upx=True,
|
upx=True,
|
||||||
upx_exclude=[],
|
upx_exclude=[],
|
||||||
name='AI_NovelGenerator_V1.2.3'
|
name='AI_NovelGenerator_V1.2.4'
|
||||||
)
|
)
|
||||||
|
|||||||
+298
-179
@@ -3,11 +3,9 @@
|
|||||||
import os
|
import os
|
||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
|
import traceback
|
||||||
from typing import Dict, List, Optional
|
from typing import Dict, List, Optional
|
||||||
try:
|
from typing import TypedDict
|
||||||
from typing import TypedDict # Python 3.8+ 直接可用;若是3.7可改用 typing_extensions
|
|
||||||
except ImportError:
|
|
||||||
from typing_extensions import TypedDict
|
|
||||||
|
|
||||||
from langchain_openai import ChatOpenAI
|
from langchain_openai import ChatOpenAI
|
||||||
from langgraph.graph import StateGraph, START, END
|
from langgraph.graph import StateGraph, START, END
|
||||||
@@ -30,20 +28,79 @@ from prompt_definitions import (
|
|||||||
summary_prompt, update_character_state_prompt,
|
summary_prompt, update_character_state_prompt,
|
||||||
chapter_outline_prompt, chapter_write_prompt
|
chapter_outline_prompt, chapter_write_prompt
|
||||||
)
|
)
|
||||||
|
from embedding_ollama import OllamaEmbeddings
|
||||||
# ============ 新增:导入 chapter_directory_parser ============
|
|
||||||
from chapter_directory_parser import get_chapter_info_from_directory
|
from chapter_directory_parser import get_chapter_info_from_directory
|
||||||
|
|
||||||
# ============ 日志配置 ============
|
# ============ 日志配置 ============
|
||||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
||||||
|
|
||||||
def debug_log(prompt: str, response_content: str):
|
def debug_log(prompt: str, response_content: str):
|
||||||
"""在控制台打印或记录下每次Prompt与Response,[调试]"""
|
"""
|
||||||
|
打印prompt和response的辅助函数
|
||||||
|
"""
|
||||||
logging.info(f"\n[Prompt >>>] {prompt}\n")
|
logging.info(f"\n[Prompt >>>] {prompt}\n")
|
||||||
logging.info(f"[Response >>>] {response_content}\n")
|
logging.info(f"[Response >>>] {response_content}\n")
|
||||||
|
|
||||||
# ============ 向量检索相关 ============
|
def remove_think_tags(text: str) -> str:
|
||||||
|
"""
|
||||||
|
移除 <think>...</think> 包裹的内容
|
||||||
|
"""
|
||||||
|
return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
|
||||||
|
|
||||||
|
def invoke_with_cleaning(model: ChatOpenAI, prompt: str) -> str:
|
||||||
|
"""
|
||||||
|
通用封装:调用模型并移除 <think>...</think> 文本,记录日志后返回
|
||||||
|
"""
|
||||||
|
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()
|
||||||
|
|
||||||
|
# ============ 判断接口格式相关 ============
|
||||||
|
def is_using_ollama_api(interface_format: str, base_url: str) -> bool:
|
||||||
|
"""
|
||||||
|
当 interface_format == "Ollama" 时返回 True
|
||||||
|
"""
|
||||||
|
return interface_format.lower() == "ollama"
|
||||||
|
|
||||||
|
def is_using_ml_studio_api(interface_format: str, base_url: str) -> bool:
|
||||||
|
"""
|
||||||
|
如果用户在下拉里选择了 ML Studio
|
||||||
|
"""
|
||||||
|
return interface_format.lower() == "ml studio"
|
||||||
|
|
||||||
|
# ============ 创建 Embeddings 对象 ============
|
||||||
|
def create_embeddings_object(
|
||||||
|
api_key: str,
|
||||||
|
base_url: str,
|
||||||
|
embed_url: str,
|
||||||
|
interface_format: str,
|
||||||
|
embedding_model_name: str
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
根据用户在UI中配置的参数,返回对应的 embeddings 对象。
|
||||||
|
- 当 interface_format = "Ollama" => OllamaEmbeddings(...)
|
||||||
|
- 当 interface_format = "OpenAI"/"ML Studio" => OpenAIEmbeddings(...)
|
||||||
|
- 其它情况可扩展
|
||||||
|
"""
|
||||||
|
if is_using_ollama_api(interface_format, embed_url):
|
||||||
|
fixed_url = embed_url.rstrip("/")
|
||||||
|
# Ollama embedding接口通常是 /api/embed
|
||||||
|
fixed_url = fixed_url.replace("/v1", "/api")
|
||||||
|
return OllamaEmbeddings(
|
||||||
|
model_name=embedding_model_name,
|
||||||
|
base_url=fixed_url
|
||||||
|
)
|
||||||
|
elif is_using_ml_studio_api(interface_format, base_url):
|
||||||
|
return OpenAIEmbeddings(openai_api_key=api_key, openai_api_base=base_url)
|
||||||
|
else:
|
||||||
|
# 默认使用 OpenAIEmbeddings
|
||||||
|
return OpenAIEmbeddings(openai_api_key=api_key, openai_api_base=base_url)
|
||||||
|
|
||||||
|
# ============ 向量库相关 ============
|
||||||
VECTOR_STORE_DIR = os.path.join(os.getcwd(), "vectorstore")
|
VECTOR_STORE_DIR = os.path.join(os.getcwd(), "vectorstore")
|
||||||
if not os.path.exists(VECTOR_STORE_DIR):
|
if not os.path.exists(VECTOR_STORE_DIR):
|
||||||
os.makedirs(VECTOR_STORE_DIR)
|
os.makedirs(VECTOR_STORE_DIR)
|
||||||
@@ -51,11 +108,10 @@ if not os.path.exists(VECTOR_STORE_DIR):
|
|||||||
def clear_vector_store():
|
def clear_vector_store():
|
||||||
"""
|
"""
|
||||||
清空本地向量库(删除 vectorstore 文件夹内的内容)。
|
清空本地向量库(删除 vectorstore 文件夹内的内容)。
|
||||||
需要在UI中加一个二次确认弹窗,防止误删。
|
|
||||||
"""
|
"""
|
||||||
if os.path.exists(VECTOR_STORE_DIR):
|
if os.path.exists(VECTOR_STORE_DIR):
|
||||||
try:
|
|
||||||
import shutil
|
import shutil
|
||||||
|
try:
|
||||||
for filename in os.listdir(VECTOR_STORE_DIR):
|
for filename in os.listdir(VECTOR_STORE_DIR):
|
||||||
file_path = os.path.join(VECTOR_STORE_DIR, filename)
|
file_path = os.path.join(VECTOR_STORE_DIR, filename)
|
||||||
if os.path.isfile(file_path) or os.path.islink(file_path):
|
if os.path.isfile(file_path) or os.path.islink(file_path):
|
||||||
@@ -63,19 +119,29 @@ def clear_vector_store():
|
|||||||
elif os.path.isdir(file_path):
|
elif os.path.isdir(file_path):
|
||||||
shutil.rmtree(file_path)
|
shutil.rmtree(file_path)
|
||||||
logging.info("Local vector store has been cleared.")
|
logging.info("Local vector store has been cleared.")
|
||||||
except Exception as e:
|
except Exception:
|
||||||
logging.warning(f"Failed to clear vector store: {e}")
|
logging.warning(f"Failed to clear vector store:\n{traceback.format_exc()}")
|
||||||
else:
|
else:
|
||||||
logging.info("No vector store found to clear.")
|
logging.info("No vector store found to clear.")
|
||||||
|
|
||||||
def init_vector_store(api_key: str, base_url: str, texts: List[str]) -> Chroma:
|
def init_vector_store(
|
||||||
|
api_key: str,
|
||||||
|
base_url: str,
|
||||||
|
interface_format: str,
|
||||||
|
embedding_model_name: str,
|
||||||
|
texts: List[str],
|
||||||
|
embedding_base_url: str = ""
|
||||||
|
) -> Chroma:
|
||||||
"""
|
"""
|
||||||
初始化并返回一个Chroma向量库,将传入的文本进行嵌入并保存到本地目录。
|
初始化并返回一个Chroma向量库,将传入的文本进行嵌入并保存到本地目录。
|
||||||
如果不存在该目录,会自动创建。
|
|
||||||
"""
|
"""
|
||||||
embeddings = OpenAIEmbeddings(
|
embed_url = embedding_base_url if embedding_base_url else base_url
|
||||||
openai_api_key=api_key,
|
embeddings = create_embeddings_object(
|
||||||
openai_api_base=base_url
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
embed_url=embed_url,
|
||||||
|
interface_format=interface_format,
|
||||||
|
embedding_model_name=embedding_model_name
|
||||||
)
|
)
|
||||||
documents = [Document(page_content=t) for t in texts]
|
documents = [Document(page_content=t) for t in texts]
|
||||||
vectorstore = Chroma.from_documents(
|
vectorstore = Chroma.from_documents(
|
||||||
@@ -86,43 +152,101 @@ def init_vector_store(api_key: str, base_url: str, texts: List[str]) -> Chroma:
|
|||||||
vectorstore.persist()
|
vectorstore.persist()
|
||||||
return vectorstore
|
return vectorstore
|
||||||
|
|
||||||
def load_vector_store(api_key: str, base_url: str) -> Optional[Chroma]:
|
def load_vector_store(
|
||||||
"""读取已存在的向量库。若不存在则返回 None。"""
|
api_key: str,
|
||||||
|
base_url: str,
|
||||||
|
interface_format: str,
|
||||||
|
embedding_model_name: str,
|
||||||
|
embedding_base_url: str = ""
|
||||||
|
) -> Optional[Chroma]:
|
||||||
|
"""
|
||||||
|
读取已存在的向量库。若不存在则返回 None。
|
||||||
|
"""
|
||||||
if not os.path.exists(VECTOR_STORE_DIR):
|
if not os.path.exists(VECTOR_STORE_DIR):
|
||||||
|
logging.info("Vector store not found. Initializing a new one...")
|
||||||
return None
|
return None
|
||||||
embeddings = OpenAIEmbeddings(
|
embed_url = embedding_base_url if embedding_base_url else base_url
|
||||||
openai_api_key=api_key,
|
embeddings = create_embeddings_object(
|
||||||
openai_api_base=base_url
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
embed_url=embed_url,
|
||||||
|
interface_format=interface_format,
|
||||||
|
embedding_model_name=embedding_model_name
|
||||||
)
|
)
|
||||||
return Chroma(persist_directory=VECTOR_STORE_DIR, embedding_function=embeddings)
|
return Chroma(persist_directory=VECTOR_STORE_DIR, embedding_function=embeddings)
|
||||||
|
|
||||||
def update_vector_store(api_key: str, base_url: str, new_chapter: str) -> None:
|
def update_vector_store(
|
||||||
"""将最新章节文本插入到向量库里,用于后续检索参考。若库不存在则初始化。"""
|
api_key: str,
|
||||||
store = load_vector_store(api_key, base_url)
|
base_url: str,
|
||||||
|
new_chapter: str,
|
||||||
|
interface_format: str,
|
||||||
|
embedding_model_name: str,
|
||||||
|
embedding_base_url: str = ""
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
将最新章节文本插入到向量库里,用于后续检索参考。若库不存在则初始化。
|
||||||
|
"""
|
||||||
|
store = load_vector_store(
|
||||||
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
interface_format=interface_format,
|
||||||
|
embedding_model_name=embedding_model_name,
|
||||||
|
embedding_base_url=embedding_base_url
|
||||||
|
)
|
||||||
|
|
||||||
|
# 如果向量库不存在,初始化它
|
||||||
if not store:
|
if not store:
|
||||||
logging.info("Vector store does not exist. Initializing a new one...")
|
logging.info("Vector store does not exist. Initializing a new one for new chapter...")
|
||||||
init_vector_store(api_key, base_url, [new_chapter])
|
init_vector_store(
|
||||||
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
interface_format=interface_format,
|
||||||
|
embedding_model_name=embedding_model_name,
|
||||||
|
texts=[new_chapter],
|
||||||
|
embedding_base_url=embedding_base_url
|
||||||
|
)
|
||||||
return
|
return
|
||||||
|
|
||||||
new_doc = Document(page_content=new_chapter)
|
new_doc = Document(page_content=new_chapter)
|
||||||
store.add_documents([new_doc])
|
store.add_documents([new_doc])
|
||||||
store.persist()
|
store.persist()
|
||||||
|
logging.info("Vector store updated with the new chapter.")
|
||||||
|
|
||||||
def get_relevant_context_from_vector_store(api_key: str, base_url: str, query: str, k: int = 2) -> str:
|
def get_relevant_context_from_vector_store(
|
||||||
|
api_key: str,
|
||||||
|
base_url: str,
|
||||||
|
query: str,
|
||||||
|
interface_format: str,
|
||||||
|
embedding_model_name: str,
|
||||||
|
embedding_base_url: str = "",
|
||||||
|
k: int = 2
|
||||||
|
) -> str:
|
||||||
"""
|
"""
|
||||||
从向量库中检索与 query 最相关的 k 条文本,拼接后返回。
|
从向量库中检索与 query 最相关的 k 条文本,拼接后返回。
|
||||||
若向量库不存在则返回空字符串。
|
若向量库不存在或没有足够的内容,则返回空字符串。
|
||||||
"""
|
"""
|
||||||
store = load_vector_store(api_key, base_url)
|
store = load_vector_store(
|
||||||
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
interface_format=interface_format,
|
||||||
|
embedding_model_name=embedding_model_name,
|
||||||
|
embedding_base_url=embedding_base_url
|
||||||
|
)
|
||||||
|
|
||||||
|
# 如果向量库为空,直接返回空字符串
|
||||||
if not store:
|
if not store:
|
||||||
logging.warning("Vector store not found. Returning empty context.")
|
logging.info("No vector store found. Returning empty context.")
|
||||||
return ""
|
return ""
|
||||||
|
|
||||||
docs = store.similarity_search(query, k=k)
|
docs = store.similarity_search(query, k=k)
|
||||||
|
if not docs:
|
||||||
|
logging.info(f"No relevant documents found for query '{query}'. Returning empty context.")
|
||||||
|
return ""
|
||||||
|
|
||||||
combined = "\n".join([d.page_content for d in docs])
|
combined = "\n".join([d.page_content for d in docs])
|
||||||
return combined
|
return combined
|
||||||
|
|
||||||
# ============ 多步生成:设置 & 目录 ============
|
# ============ 多步生成:设置 & 目录 ============
|
||||||
|
|
||||||
class OverallState(TypedDict):
|
class OverallState(TypedDict):
|
||||||
topic: str
|
topic: str
|
||||||
genre: str
|
genre: str
|
||||||
@@ -148,7 +272,6 @@ def Novel_novel_directory_generate(
|
|||||||
"""
|
"""
|
||||||
使用多步流程,生成 Novel_setting.txt 与 Novel_directory.txt 并保存到 filepath。
|
使用多步流程,生成 Novel_setting.txt 与 Novel_directory.txt 并保存到 filepath。
|
||||||
"""
|
"""
|
||||||
# 确保文件夹存在
|
|
||||||
os.makedirs(filepath, exist_ok=True)
|
os.makedirs(filepath, exist_ok=True)
|
||||||
|
|
||||||
model = ChatOpenAI(
|
model = ChatOpenAI(
|
||||||
@@ -165,67 +288,46 @@ def Novel_novel_directory_generate(
|
|||||||
number_of_chapters=state["number_of_chapters"],
|
number_of_chapters=state["number_of_chapters"],
|
||||||
word_number=state["word_number"]
|
word_number=state["word_number"]
|
||||||
)
|
)
|
||||||
response = model.invoke(prompt)
|
result_text = invoke_with_cleaning(model, prompt)
|
||||||
if not response:
|
return {"novel_setting_base": result_text}
|
||||||
logging.warning("generate_base_setting: No response.")
|
|
||||||
return {"novel_setting_base": ""}
|
|
||||||
debug_log(prompt, response.content)
|
|
||||||
return {"novel_setting_base": response.content.strip()}
|
|
||||||
|
|
||||||
def generate_character_setting(state: OverallState) -> Dict[str, str]:
|
def generate_character_setting(state: OverallState) -> Dict[str, str]:
|
||||||
prompt = character_prompt.format(
|
prompt = character_prompt.format(
|
||||||
novel_setting=state["novel_setting_base"]
|
novel_setting=state["novel_setting_base"]
|
||||||
)
|
)
|
||||||
response = model.invoke(prompt)
|
result_text = invoke_with_cleaning(model, prompt)
|
||||||
if not response:
|
return {"character_setting": result_text}
|
||||||
logging.warning("generate_character_setting: No response.")
|
|
||||||
return {"character_setting": ""}
|
|
||||||
debug_log(prompt, response.content)
|
|
||||||
return {"character_setting": response.content.strip()}
|
|
||||||
|
|
||||||
def generate_dark_lines(state: OverallState) -> Dict[str, str]:
|
def generate_dark_lines(state: OverallState) -> Dict[str, str]:
|
||||||
prompt = dark_lines_prompt.format(
|
prompt = dark_lines_prompt.format(
|
||||||
character_info=state["character_setting"]
|
character_info=state["character_setting"]
|
||||||
)
|
)
|
||||||
response = model.invoke(prompt)
|
result_text = invoke_with_cleaning(model, prompt)
|
||||||
if not response:
|
return {"dark_lines": result_text}
|
||||||
logging.warning("generate_dark_lines: No response.")
|
|
||||||
return {"dark_lines": ""}
|
|
||||||
debug_log(prompt, response.content)
|
|
||||||
return {"dark_lines": response.content.strip()}
|
|
||||||
|
|
||||||
def finalize_novel_setting(state: OverallState) -> Dict[str, str]:
|
def finalize_novel_setting_func(state: OverallState) -> Dict[str, str]:
|
||||||
prompt = finalize_setting_prompt.format(
|
prompt = finalize_setting_prompt.format(
|
||||||
novel_setting_base=state["novel_setting_base"],
|
novel_setting_base=state["novel_setting_base"],
|
||||||
character_setting=state["character_setting"],
|
character_setting=state["character_setting"],
|
||||||
dark_lines=state["dark_lines"]
|
dark_lines=state["dark_lines"]
|
||||||
)
|
)
|
||||||
response = model.invoke(prompt)
|
result_text = invoke_with_cleaning(model, prompt)
|
||||||
if not response:
|
return {"final_novel_setting": result_text}
|
||||||
logging.warning("finalize_novel_setting: No response.")
|
|
||||||
return {"final_novel_setting": ""}
|
|
||||||
debug_log(prompt, response.content)
|
|
||||||
return {"final_novel_setting": response.content.strip()}
|
|
||||||
|
|
||||||
def generate_novel_directory(state: OverallState) -> Dict[str, str]:
|
def generate_novel_directory_func(state: OverallState) -> Dict[str, str]:
|
||||||
prompt = novel_directory_prompt.format(
|
prompt = novel_directory_prompt.format(
|
||||||
final_novel_setting=state["final_novel_setting"],
|
final_novel_setting=state["final_novel_setting"],
|
||||||
number_of_chapters=state["number_of_chapters"]
|
number_of_chapters=state["number_of_chapters"]
|
||||||
)
|
)
|
||||||
response = model.invoke(prompt)
|
result_text = invoke_with_cleaning(model, prompt)
|
||||||
if not response:
|
return {"novel_directory": result_text}
|
||||||
logging.warning("generate_novel_directory: No response.")
|
|
||||||
return {"novel_directory": ""}
|
|
||||||
debug_log(prompt, response.content)
|
|
||||||
return {"novel_directory": response.content.strip()}
|
|
||||||
|
|
||||||
# 构建状态图
|
|
||||||
graph = StateGraph(OverallState)
|
graph = StateGraph(OverallState)
|
||||||
graph.add_node("generate_base_setting", generate_base_setting)
|
graph.add_node("generate_base_setting", generate_base_setting)
|
||||||
graph.add_node("generate_character_setting", generate_character_setting)
|
graph.add_node("generate_character_setting", generate_character_setting)
|
||||||
graph.add_node("generate_dark_lines", generate_dark_lines)
|
graph.add_node("generate_dark_lines", generate_dark_lines)
|
||||||
graph.add_node("finalize_novel_setting", finalize_novel_setting)
|
graph.add_node("finalize_novel_setting", finalize_novel_setting_func)
|
||||||
graph.add_node("generate_novel_directory", generate_novel_directory)
|
graph.add_node("generate_novel_directory", generate_novel_directory_func)
|
||||||
|
|
||||||
graph.add_edge(START, "generate_base_setting")
|
graph.add_edge(START, "generate_base_setting")
|
||||||
graph.add_edge("generate_base_setting", "generate_character_setting")
|
graph.add_edge("generate_base_setting", "generate_character_setting")
|
||||||
@@ -255,24 +357,25 @@ def Novel_novel_directory_generate(
|
|||||||
logging.warning("生成失败:缺少 final_novel_setting 或 novel_directory。")
|
logging.warning("生成失败:缺少 final_novel_setting 或 novel_directory。")
|
||||||
return
|
return
|
||||||
|
|
||||||
# 写入文件
|
|
||||||
filename_set = os.path.join(filepath, "Novel_setting.txt")
|
filename_set = os.path.join(filepath, "Novel_setting.txt")
|
||||||
filename_novel_directory = os.path.join(filepath, "Novel_directory.txt")
|
filename_novel_directory = os.path.join(filepath, "Novel_directory.txt")
|
||||||
|
|
||||||
# 清理文本(可根据需要去除多余字符)
|
|
||||||
def clean_text(txt: str) -> str:
|
def clean_text(txt: str) -> str:
|
||||||
return txt.replace('#', '').replace('*', '')
|
return txt.replace('#', '').replace('*', '')
|
||||||
|
|
||||||
final_novel_setting_cleaned = clean_text(final_novel_setting)
|
final_novel_setting_cleaned = clean_text(final_novel_setting)
|
||||||
final_novel_directory_cleaned = clean_text(final_novel_directory)
|
final_novel_directory_cleaned = clean_text(final_novel_directory)
|
||||||
|
|
||||||
append_text_to_file(final_novel_setting_cleaned, filename_set)
|
# 改进:写文件时先清空再写入
|
||||||
append_text_to_file(final_novel_directory_cleaned, filename_novel_directory)
|
clear_file_content(filename_set)
|
||||||
|
save_string_to_txt(final_novel_setting_cleaned, filename_set)
|
||||||
|
|
||||||
|
clear_file_content(filename_novel_directory)
|
||||||
|
save_string_to_txt(final_novel_directory_cleaned, filename_novel_directory)
|
||||||
|
|
||||||
logging.info("Novel settings and directory generated successfully.")
|
logging.info("Novel settings and directory generated successfully.")
|
||||||
|
|
||||||
# ============ 获取最近N章内容,生成短期摘要 ============
|
# ============ 获取最近N章内容,生成短期摘要 ============
|
||||||
|
|
||||||
def get_last_n_chapters_text(chapters_dir: str, current_chapter_num: int, n: int = 3) -> List[str]:
|
def get_last_n_chapters_text(chapters_dir: str, current_chapter_num: int, n: int = 3) -> List[str]:
|
||||||
"""
|
"""
|
||||||
从指定文件夹中,读取最近 n 章的内容(如果存在),并按从旧到新的顺序返回文本列表。
|
从指定文件夹中,读取最近 n 章的内容(如果存在),并按从旧到新的顺序返回文本列表。
|
||||||
@@ -286,30 +389,45 @@ def get_last_n_chapters_text(chapters_dir: str, current_chapter_num: int, n: int
|
|||||||
text = read_file(chap_file).strip()
|
text = read_file(chap_file).strip()
|
||||||
if text:
|
if text:
|
||||||
texts.append(text)
|
texts.append(text)
|
||||||
|
if len(texts) < n:
|
||||||
|
texts = [''] * (n - len(texts)) + texts
|
||||||
return texts
|
return texts
|
||||||
|
|
||||||
def summarize_recent_chapters(model: ChatOpenAI, chapters_text_list: List[str]) -> str:
|
def summarize_recent_chapters(
|
||||||
|
llm_model: str,
|
||||||
|
api_key: str,
|
||||||
|
base_url: str,
|
||||||
|
temperature: float,
|
||||||
|
chapters_text_list: List[str]
|
||||||
|
) -> str:
|
||||||
"""
|
"""
|
||||||
将最近几章的文本拼接后,通过模型生成一个相对详细的“短期内容摘要”。
|
将最近几章文本拼接,通过模型生成相对简要的“短期内容摘要”。
|
||||||
"""
|
"""
|
||||||
if not chapters_text_list:
|
if not chapters_text_list:
|
||||||
return ""
|
return ""
|
||||||
|
if all(not txt.strip() for txt in chapters_text_list):
|
||||||
|
return "暂无摘要。"
|
||||||
|
|
||||||
|
model = ChatOpenAI(
|
||||||
|
model=llm_model,
|
||||||
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
temperature=temperature
|
||||||
|
)
|
||||||
|
|
||||||
combined_text = "\n".join(chapters_text_list)
|
combined_text = "\n".join(chapters_text_list)
|
||||||
prompt = f"""\
|
prompt = f"""你是一名资深长篇小说写作辅助AI,下面是最近几章的合并文本:
|
||||||
这是最近几章的故事内容,请生成一份详细的短期内容摘要(不少于一章篇幅的细节),用于帮助后续创作时回顾细节。
|
|
||||||
请着重强调发生的事件、角色的心理和关系变化、冲突或悬念等。
|
|
||||||
|
|
||||||
{combined_text}
|
{combined_text}
|
||||||
"""
|
|
||||||
response = model.invoke(prompt)
|
|
||||||
if not response:
|
|
||||||
return ""
|
|
||||||
debug_log(prompt, response.content)
|
|
||||||
return response.content.strip()
|
|
||||||
|
|
||||||
# ============ 新增1:记录剧情要点/未解决冲突 ============
|
请用中文输出不超过500字的摘要,只包含主要剧情进展、角色变化、冲突焦点等要点:"""
|
||||||
|
|
||||||
|
summary_text = invoke_with_cleaning(model, prompt)
|
||||||
|
if not summary_text:
|
||||||
|
# 若模型无响应,就截取一段作为“备选”
|
||||||
|
return combined_text[:800] + "..." if len(combined_text) > 800 else combined_text
|
||||||
|
return summary_text
|
||||||
|
|
||||||
|
# ============ 新增:剧情要点/未解决冲突 ============
|
||||||
PLOT_ARCS_PROMPT = """\
|
PLOT_ARCS_PROMPT = """\
|
||||||
下面是新生成的章节内容:
|
下面是新生成的章节内容:
|
||||||
{chapter_text}
|
{chapter_text}
|
||||||
@@ -317,9 +435,9 @@ PLOT_ARCS_PROMPT = """\
|
|||||||
这里是已记录的剧情要点/未解决冲突(可能为空):
|
这里是已记录的剧情要点/未解决冲突(可能为空):
|
||||||
{old_plot_arcs}
|
{old_plot_arcs}
|
||||||
|
|
||||||
请基于新的章节内容,提炼出本章引入或延续的悬念、冲突、角色暗线等,将其合并到旧的剧情要点中。
|
请基于新的章节内容,提炼本章引入或延续的悬念、冲突、角色暗线等,将其合并到旧的剧情要点中。
|
||||||
若有新的冲突则添加,若有已解决/不再重要的冲突可标注或移除。
|
若有新的冲突则添加,若有已解决/不再重要的冲突可标注或移除。
|
||||||
最终输出一份更新后的剧情要点列表,以帮助后续保持故事的整体一致性和悬念延续。
|
最终输出更新后的剧情要点列表,以帮助后续保持故事整体的一致性和悬念延续。
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def update_plot_arcs(
|
def update_plot_arcs(
|
||||||
@@ -330,10 +448,6 @@ def update_plot_arcs(
|
|||||||
model_name: str,
|
model_name: str,
|
||||||
temperature: float
|
temperature: float
|
||||||
) -> str:
|
) -> str:
|
||||||
"""
|
|
||||||
利用模型分析最新章节文本,提炼或更新“未解决冲突或剧情要点”。
|
|
||||||
并返回更新后的字符串。
|
|
||||||
"""
|
|
||||||
model = ChatOpenAI(
|
model = ChatOpenAI(
|
||||||
model=model_name,
|
model=model_name,
|
||||||
api_key=api_key,
|
api_key=api_key,
|
||||||
@@ -344,15 +458,13 @@ def update_plot_arcs(
|
|||||||
chapter_text=chapter_text,
|
chapter_text=chapter_text,
|
||||||
old_plot_arcs=old_plot_arcs
|
old_plot_arcs=old_plot_arcs
|
||||||
)
|
)
|
||||||
response = model.invoke(prompt)
|
arcs_text = invoke_with_cleaning(model, prompt)
|
||||||
if not response:
|
if not arcs_text:
|
||||||
logging.warning("update_plot_arcs: No response.")
|
logging.warning("update_plot_arcs: No response or empty result.")
|
||||||
return old_plot_arcs
|
return old_plot_arcs
|
||||||
debug_log(prompt, response.content)
|
return arcs_text
|
||||||
return response.content.strip()
|
|
||||||
|
|
||||||
# ============ 生成章节草稿 & 定稿 ============
|
# ============ 生成章节草稿 & 定稿 ============
|
||||||
|
|
||||||
def generate_chapter_draft(
|
def generate_chapter_draft(
|
||||||
novel_settings: str,
|
novel_settings: str,
|
||||||
global_summary: str,
|
global_summary: str,
|
||||||
@@ -366,24 +478,43 @@ def generate_chapter_draft(
|
|||||||
word_number: int,
|
word_number: int,
|
||||||
temperature: float,
|
temperature: float,
|
||||||
novel_novel_directory: str,
|
novel_novel_directory: str,
|
||||||
filepath: str
|
filepath: str,
|
||||||
|
interface_format: str,
|
||||||
|
embedding_model_name: str,
|
||||||
|
embedding_base_url: str
|
||||||
) -> str:
|
) -> str:
|
||||||
"""
|
"""
|
||||||
仅生成当前章节的草稿,不更新全局摘要/角色状态/向量库。
|
生成当前章节的草稿,不更新全局摘要/角色状态/向量库。
|
||||||
并将生成的内容写到 "chapter_{novel_number}.txt" 覆盖写入。
|
|
||||||
同时生成 "outline_{novel_number}.txt" 存储大纲内容。
|
|
||||||
"""
|
"""
|
||||||
# 0) 根据 novel_number 从 novel_novel_directory 中获取本章标题及简述
|
# 根据目录信息获取本章标题、简介
|
||||||
chapter_info = get_chapter_info_from_directory(novel_novel_directory, novel_number)
|
chapter_info = get_chapter_info_from_directory(novel_novel_directory, novel_number)
|
||||||
chapter_title = chapter_info["chapter_title"]
|
chapter_title = chapter_info["chapter_title"]
|
||||||
chapter_brief = chapter_info["chapter_brief"]
|
chapter_brief = chapter_info["chapter_brief"]
|
||||||
|
|
||||||
# 1) 从向量库检索往期上下文
|
# 从向量库检索上下文
|
||||||
relevant_context = get_relevant_context_from_vector_store(
|
queries = []
|
||||||
api_key, base_url, "回顾剧情", k=2
|
if user_guidance.strip():
|
||||||
|
queries.append(user_guidance)
|
||||||
|
if chapter_brief.strip():
|
||||||
|
queries.append(chapter_brief)
|
||||||
|
queries.append("回顾剧情")
|
||||||
|
|
||||||
|
relevant_context = ""
|
||||||
|
for q in queries:
|
||||||
|
partial_context = get_relevant_context_from_vector_store(
|
||||||
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
query=q,
|
||||||
|
interface_format=interface_format,
|
||||||
|
embedding_model_name=embedding_model_name,
|
||||||
|
embedding_base_url=embedding_base_url,
|
||||||
|
k=2
|
||||||
)
|
)
|
||||||
|
if partial_context.strip():
|
||||||
|
relevant_context += "\n" + partial_context
|
||||||
|
if not relevant_context:
|
||||||
|
relevant_context = "暂无相关内容。"
|
||||||
|
|
||||||
# 2) 生成大纲
|
|
||||||
model = ChatOpenAI(
|
model = ChatOpenAI(
|
||||||
model=model_name,
|
model=model_name,
|
||||||
api_key=api_key,
|
api_key=api_key,
|
||||||
@@ -391,26 +522,19 @@ def generate_chapter_draft(
|
|||||||
temperature=temperature
|
temperature=temperature
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# 1) 生成本章大纲
|
||||||
outline_prompt_text = chapter_outline_prompt.format(
|
outline_prompt_text = chapter_outline_prompt.format(
|
||||||
novel_setting=novel_settings,
|
novel_setting=novel_settings,
|
||||||
character_state=character_state + "\n\n【历史上下文】\n" + relevant_context,
|
character_state=character_state + "\n\n【检索到的上下文】\n" + relevant_context,
|
||||||
global_summary=global_summary,
|
global_summary=global_summary,
|
||||||
novel_number=novel_number,
|
novel_number=novel_number,
|
||||||
chapter_title=chapter_title,
|
chapter_title=chapter_title,
|
||||||
chapter_brief=chapter_brief
|
chapter_brief=chapter_brief
|
||||||
)
|
)
|
||||||
|
outline_prompt_text += f"\n\n【最近几章摘要】\n{recent_chapters_summary}"
|
||||||
outline_prompt_text += f"\n\n【本章目录标题与简述】\n标题:{chapter_title}\n简述:{chapter_brief}\n"
|
|
||||||
outline_prompt_text += f"\n【最近几章摘要】\n{recent_chapters_summary}"
|
|
||||||
outline_prompt_text += f"\n\n【用户指导】\n{user_guidance if user_guidance else '(无)'}"
|
outline_prompt_text += f"\n\n【用户指导】\n{user_guidance if user_guidance else '(无)'}"
|
||||||
|
|
||||||
response_outline = model.invoke(outline_prompt_text)
|
chapter_outline = invoke_with_cleaning(model, outline_prompt_text)
|
||||||
if not response_outline:
|
|
||||||
logging.warning("generate_chapter_draft: outline no response.")
|
|
||||||
chapter_outline = ""
|
|
||||||
else:
|
|
||||||
debug_log(outline_prompt_text, response_outline.content)
|
|
||||||
chapter_outline = response_outline.content.strip()
|
|
||||||
|
|
||||||
outlines_dir = os.path.join(filepath, "outlines")
|
outlines_dir = os.path.join(filepath, "outlines")
|
||||||
os.makedirs(outlines_dir, exist_ok=True)
|
os.makedirs(outlines_dir, exist_ok=True)
|
||||||
@@ -418,28 +542,20 @@ def generate_chapter_draft(
|
|||||||
clear_file_content(outline_file)
|
clear_file_content(outline_file)
|
||||||
save_string_to_txt(chapter_outline, outline_file)
|
save_string_to_txt(chapter_outline, outline_file)
|
||||||
|
|
||||||
# 3) 生成正文草稿
|
# 2) 生成正文草稿
|
||||||
writing_prompt_text = chapter_write_prompt.format(
|
writing_prompt_text = chapter_write_prompt.format(
|
||||||
novel_setting=novel_settings,
|
novel_setting=novel_settings,
|
||||||
character_state=character_state + "\n\n【历史上下文】\n" + relevant_context,
|
character_state=character_state + "\n\n【检索到的上下文】\n" + relevant_context,
|
||||||
global_summary=global_summary,
|
global_summary=global_summary,
|
||||||
chapter_outline=chapter_outline,
|
chapter_outline=chapter_outline,
|
||||||
word_number=word_number,
|
word_number=word_number,
|
||||||
chapter_title=chapter_title,
|
chapter_title=chapter_title,
|
||||||
chapter_brief=chapter_brief
|
chapter_brief=chapter_brief
|
||||||
)
|
)
|
||||||
|
writing_prompt_text += f"\n\n【最近几章摘要】\n{recent_chapters_summary}"
|
||||||
writing_prompt_text += f"\n\n【本章目录标题与简述】\n标题:{chapter_title}\n简述:{chapter_brief}\n"
|
|
||||||
writing_prompt_text += f"\n【最近几章摘要】\n{recent_chapters_summary}"
|
|
||||||
writing_prompt_text += f"\n\n【用户指导】\n{user_guidance if user_guidance else '(无)'}"
|
writing_prompt_text += f"\n\n【用户指导】\n{user_guidance if user_guidance else '(无)'}"
|
||||||
|
|
||||||
response_chapter = model.invoke(writing_prompt_text)
|
chapter_content = invoke_with_cleaning(model, writing_prompt_text)
|
||||||
if not response_chapter:
|
|
||||||
logging.warning("generate_chapter_draft: writing no response.")
|
|
||||||
chapter_content = ""
|
|
||||||
else:
|
|
||||||
debug_log(writing_prompt_text, response_chapter.content)
|
|
||||||
chapter_content = response_chapter.content.strip()
|
|
||||||
|
|
||||||
chapters_dir = os.path.join(filepath, "chapters")
|
chapters_dir = os.path.join(filepath, "chapters")
|
||||||
os.makedirs(chapters_dir, exist_ok=True)
|
os.makedirs(chapters_dir, exist_ok=True)
|
||||||
@@ -455,19 +571,20 @@ def finalize_chapter(
|
|||||||
word_number: int,
|
word_number: int,
|
||||||
api_key: str,
|
api_key: str,
|
||||||
base_url: str,
|
base_url: str,
|
||||||
|
interface_format: str,
|
||||||
|
embedding_model_name: str,
|
||||||
model_name: str,
|
model_name: str,
|
||||||
temperature: float,
|
temperature: float,
|
||||||
filepath: str
|
filepath: str
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
对当前章节进行定稿:
|
对当前章节进行定稿:
|
||||||
1. 读取 chapter_{novel_number}.txt 的最终内容;
|
1. 读取草稿文本
|
||||||
2. 更新全局摘要、角色状态文件;
|
2. 若字数太短则再次扩写
|
||||||
3. 如果字数明显少于 word_number 的 80%,则自动调用 enrich_chapter_text 再次扩写;
|
3. 更新全局摘要、角色状态
|
||||||
4. 更新向量库;
|
4. 更新剧情要点
|
||||||
5. 新增:更新剧情要点/未解决冲突 -> plot_arcs.txt
|
5. 更新向量库
|
||||||
"""
|
"""
|
||||||
# 读取当前章节内容
|
|
||||||
chapters_dir = os.path.join(filepath, "chapters")
|
chapters_dir = os.path.join(filepath, "chapters")
|
||||||
chapter_file = os.path.join(chapters_dir, f"chapter_{novel_number}.txt")
|
chapter_file = os.path.join(chapters_dir, f"chapter_{novel_number}.txt")
|
||||||
chapter_text = read_file(chapter_file).strip()
|
chapter_text = read_file(chapter_file).strip()
|
||||||
@@ -475,18 +592,17 @@ def finalize_chapter(
|
|||||||
logging.warning(f"Chapter {novel_number} is empty, cannot finalize.")
|
logging.warning(f"Chapter {novel_number} is empty, cannot finalize.")
|
||||||
return
|
return
|
||||||
|
|
||||||
# 读取角色状态 & 全局摘要 & 剧情要点
|
|
||||||
character_state_file = os.path.join(filepath, "character_state.txt")
|
character_state_file = os.path.join(filepath, "character_state.txt")
|
||||||
global_summary_file = os.path.join(filepath, "global_summary.txt")
|
global_summary_file = os.path.join(filepath, "global_summary.txt")
|
||||||
plot_arcs_file = os.path.join(filepath, "plot_arcs.txt") # 新增文件
|
plot_arcs_file = os.path.join(filepath, "plot_arcs.txt")
|
||||||
|
|
||||||
old_char_state = read_file(character_state_file)
|
old_char_state = read_file(character_state_file)
|
||||||
old_global_summary = read_file(global_summary_file)
|
old_global_summary = read_file(global_summary_file)
|
||||||
old_plot_arcs = read_file(plot_arcs_file)
|
old_plot_arcs = read_file(plot_arcs_file)
|
||||||
|
|
||||||
# 1) 先检查字数是否过少,若少于 80% 则调用 enrich 逻辑
|
# 若篇幅过短,二次扩写
|
||||||
if len(chapter_text) < 0.8 * word_number:
|
if len(chapter_text) < 0.8 * word_number:
|
||||||
logging.info("Chapter text seems shorter than 80% of desired length. Attempting to enrich content...")
|
logging.info("Chapter text is shorter than 80% of desired length. Enriching...")
|
||||||
chapter_text = enrich_chapter_text(
|
chapter_text = enrich_chapter_text(
|
||||||
chapter_text=chapter_text,
|
chapter_text=chapter_text,
|
||||||
word_number=word_number,
|
word_number=word_number,
|
||||||
@@ -495,12 +611,10 @@ def finalize_chapter(
|
|||||||
model_name=model_name,
|
model_name=model_name,
|
||||||
temperature=temperature
|
temperature=temperature
|
||||||
)
|
)
|
||||||
# 覆盖写回文件
|
|
||||||
clear_file_content(chapter_file)
|
clear_file_content(chapter_file)
|
||||||
save_string_to_txt(chapter_text, chapter_file)
|
save_string_to_txt(chapter_text, chapter_file)
|
||||||
logging.info("Chapter text has been enriched and updated.")
|
|
||||||
|
|
||||||
# 2) 更新全局摘要
|
# 更新全局摘要
|
||||||
model = ChatOpenAI(
|
model = ChatOpenAI(
|
||||||
model=model_name,
|
model=model_name,
|
||||||
api_key=api_key,
|
api_key=api_key,
|
||||||
@@ -513,31 +627,21 @@ def finalize_chapter(
|
|||||||
chapter_text=chapter_text,
|
chapter_text=chapter_text,
|
||||||
global_summary=old_summary
|
global_summary=old_summary
|
||||||
)
|
)
|
||||||
response = model.invoke(prompt)
|
return invoke_with_cleaning(model, prompt) or old_summary
|
||||||
if not response:
|
|
||||||
logging.warning("update_global_summary: No response.")
|
|
||||||
return old_summary
|
|
||||||
debug_log(prompt, response.content)
|
|
||||||
return response.content.strip()
|
|
||||||
|
|
||||||
new_global_summary = update_global_summary(chapter_text, old_global_summary)
|
new_global_summary = update_global_summary(chapter_text, old_global_summary)
|
||||||
|
|
||||||
# 3) 更新角色状态
|
# 更新角色状态
|
||||||
def update_character_state(chapter_text: str, old_state: str) -> str:
|
def update_character_state(chapter_text: str, old_state: str) -> str:
|
||||||
prompt = update_character_state_prompt.format(
|
prompt = update_character_state_prompt.format(
|
||||||
chapter_text=chapter_text,
|
chapter_text=chapter_text,
|
||||||
old_state=old_state
|
old_state=old_state
|
||||||
)
|
)
|
||||||
response = model.invoke(prompt)
|
return invoke_with_cleaning(model, prompt) or old_state
|
||||||
if not response:
|
|
||||||
logging.warning("update_character_state: No response.")
|
|
||||||
return old_state
|
|
||||||
debug_log(prompt, response.content)
|
|
||||||
return response.content.strip()
|
|
||||||
|
|
||||||
new_char_state = update_character_state(chapter_text, old_char_state)
|
new_char_state = update_character_state(chapter_text, old_char_state)
|
||||||
|
|
||||||
# ============ 新增2: 更新剧情要点 =============
|
# 更新剧情要点
|
||||||
new_plot_arcs = update_plot_arcs(
|
new_plot_arcs = update_plot_arcs(
|
||||||
chapter_text=chapter_text,
|
chapter_text=chapter_text,
|
||||||
old_plot_arcs=old_plot_arcs,
|
old_plot_arcs=old_plot_arcs,
|
||||||
@@ -547,7 +651,7 @@ def finalize_chapter(
|
|||||||
temperature=temperature
|
temperature=temperature
|
||||||
)
|
)
|
||||||
|
|
||||||
# 4) 覆盖写入角色状态文件、全局摘要文件、剧情要点文件
|
# 写回文件
|
||||||
clear_file_content(character_state_file)
|
clear_file_content(character_state_file)
|
||||||
save_string_to_txt(new_char_state, character_state_file)
|
save_string_to_txt(new_char_state, character_state_file)
|
||||||
|
|
||||||
@@ -557,10 +661,16 @@ def finalize_chapter(
|
|||||||
clear_file_content(plot_arcs_file)
|
clear_file_content(plot_arcs_file)
|
||||||
save_string_to_txt(new_plot_arcs, plot_arcs_file)
|
save_string_to_txt(new_plot_arcs, plot_arcs_file)
|
||||||
|
|
||||||
# 5) 更新向量检索库
|
# 更新向量库
|
||||||
update_vector_store(api_key, base_url, chapter_text)
|
update_vector_store(
|
||||||
|
api_key=api_key,
|
||||||
|
base_url=base_url,
|
||||||
|
new_chapter=chapter_text,
|
||||||
|
interface_format=interface_format,
|
||||||
|
embedding_model_name=embedding_model_name
|
||||||
|
)
|
||||||
|
|
||||||
logging.info(f"Chapter {novel_number} has been finalized (summary & state updated, plot arcs updated, vector store updated).")
|
logging.info(f"Chapter {novel_number} has been finalized.")
|
||||||
|
|
||||||
def enrich_chapter_text(
|
def enrich_chapter_text(
|
||||||
chapter_text: str,
|
chapter_text: str,
|
||||||
@@ -572,7 +682,6 @@ def enrich_chapter_text(
|
|||||||
) -> str:
|
) -> str:
|
||||||
"""
|
"""
|
||||||
当章节篇幅不足时,调用此函数对章节文本进行二次扩写。
|
当章节篇幅不足时,调用此函数对章节文本进行二次扩写。
|
||||||
可以让模型补充场景描写、角色心理等,保证与现有文本风格一致。
|
|
||||||
"""
|
"""
|
||||||
model = ChatOpenAI(
|
model = ChatOpenAI(
|
||||||
model=model_name,
|
model=model_name,
|
||||||
@@ -580,25 +689,26 @@ def enrich_chapter_text(
|
|||||||
base_url=base_url,
|
base_url=base_url,
|
||||||
temperature=temperature
|
temperature=temperature
|
||||||
)
|
)
|
||||||
prompt = f"""\
|
prompt = f"""以下是当前章节文本,可能篇幅较短,请在保持剧情连贯的前提下进行扩写,使其更充实、生动,并尽量靠近目标 {word_number} 字数。
|
||||||
以下是当前章节文本,可能篇幅较短,请在保持剧情连贯的前提下进行扩写,使其更充实、生动,并尽量靠近目标 {word_number} 字数。
|
|
||||||
|
|
||||||
原章节内容:
|
原章节内容:
|
||||||
{chapter_text}
|
{chapter_text}"""
|
||||||
"""
|
enriched_text = invoke_with_cleaning(model, prompt)
|
||||||
response = model.invoke(prompt)
|
return enriched_text if enriched_text else chapter_text
|
||||||
if not response:
|
|
||||||
logging.warning("enrich_chapter_text: No response.")
|
|
||||||
return chapter_text # 无响应时就返回原文
|
|
||||||
debug_log(prompt, response.content)
|
|
||||||
return response.content.strip()
|
|
||||||
|
|
||||||
# ============ 导入外部知识文本 ============
|
# ============ 导入外部知识文本 ============
|
||||||
|
def import_knowledge_file(
|
||||||
def import_knowledge_file(api_key: str, base_url: str, file_path: str) -> None:
|
api_key: str,
|
||||||
|
base_url: str,
|
||||||
|
interface_format: str,
|
||||||
|
embedding_model_name: str,
|
||||||
|
file_path: str,
|
||||||
|
embedding_base_url: str = ""
|
||||||
|
) -> None:
|
||||||
"""
|
"""
|
||||||
将用户选定的文本文件导入到向量库,以便在写作时检索。
|
将用户选定的文本文件导入到向量库,以便在写作时检索。
|
||||||
"""
|
"""
|
||||||
|
logging.info(f"开始导入知识库文件: {file_path}, 接口格式: {interface_format}, 模型: {embedding_model_name}")
|
||||||
if not os.path.exists(file_path):
|
if not os.path.exists(file_path):
|
||||||
logging.warning(f"知识库文件不存在: {file_path}")
|
logging.warning(f"知识库文件不存在: {file_path}")
|
||||||
return
|
return
|
||||||
@@ -608,12 +718,22 @@ def import_knowledge_file(api_key: str, base_url: str, file_path: str) -> None:
|
|||||||
logging.warning("知识库文件内容为空。")
|
logging.warning("知识库文件内容为空。")
|
||||||
return
|
return
|
||||||
|
|
||||||
|
nltk.download('punkt', quiet=True)
|
||||||
|
nltk.download('punkt_tab', quiet=True)
|
||||||
|
|
||||||
paragraphs = advanced_split_content(content)
|
paragraphs = advanced_split_content(content)
|
||||||
|
|
||||||
store = load_vector_store(api_key, base_url)
|
store = load_vector_store(api_key, base_url, interface_format, embedding_model_name, embedding_base_url)
|
||||||
if not store:
|
if not store:
|
||||||
logging.info("Vector store does not exist. Initializing a new one for knowledge import...")
|
logging.info("Vector store does not exist. Initializing a new one for knowledge import...")
|
||||||
init_vector_store(api_key, base_url, paragraphs)
|
init_vector_store(
|
||||||
|
api_key,
|
||||||
|
base_url,
|
||||||
|
interface_format,
|
||||||
|
embedding_model_name,
|
||||||
|
paragraphs,
|
||||||
|
embedding_base_url
|
||||||
|
)
|
||||||
return
|
return
|
||||||
|
|
||||||
docs = [Document(page_content=p) for p in paragraphs]
|
docs = [Document(page_content=p) for p in paragraphs]
|
||||||
@@ -625,11 +745,10 @@ def advanced_split_content(content: str,
|
|||||||
similarity_threshold: float = 0.7,
|
similarity_threshold: float = 0.7,
|
||||||
max_length: int = 500) -> List[str]:
|
max_length: int = 500) -> List[str]:
|
||||||
"""
|
"""
|
||||||
将文本先按句子切分,然后根据语义相似度进行合并,最后根据max_length进行二次切分。
|
将文本先按句子切分,然后根据语义相似度进行合并,最后按max_length二次切分。
|
||||||
|
可根据需要微调此逻辑。
|
||||||
"""
|
"""
|
||||||
nltk.download('punkt_tab', quiet=True)
|
|
||||||
sentences = nltk.sent_tokenize(content)
|
sentences = nltk.sent_tokenize(content)
|
||||||
|
|
||||||
if not sentences:
|
if not sentences:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user