diff --git a/llm_adapters.py b/llm_adapters.py index 714e72a..64695fb 100644 --- a/llm_adapters.py +++ b/llm_adapters.py @@ -5,9 +5,6 @@ from typing import Optional from langchain_openai import ChatOpenAI def ensure_openai_base_url_has_v1(url: str) -> str: - """ - 若用户输入的 url 不包含 '/v1'，则在末尾追加 '/v1'。 - """ import re url = url.strip() if not url: @@ -28,19 +25,21 @@ class DeepSeekAdapter(BaseLLMAdapter): """ 适配官方/OpenAI兼容接口（使用 langchain.ChatOpenAI） """ - def __init__(self, api_key: str, base_url: str, model_name: str, max_tokens: int, temperature: float = 0.7): + def __init__(self, api_key: str, base_url: str, model_name: str, max_tokens: int, temperature: float = 0.7, timeout: Optional[int] = 600): self.base_url = ensure_openai_base_url_has_v1(base_url) self.api_key = api_key self.model_name = model_name self.max_tokens = max_tokens self.temperature = temperature + self.timeout = timeout self._client = ChatOpenAI( model=self.model_name, api_key=self.api_key, base_url=self.base_url, max_tokens=self.max_tokens, - temperature=self.temperature + temperature=self.temperature, + timeout=self.timeout ) def invoke(self, prompt: str) -> str: @@ -54,19 +53,21 @@ class OpenAIAdapter(BaseLLMAdapter): """ 适配官方/OpenAI兼容接口（使用 langchain.ChatOpenAI） """ - def __init__(self, api_key: str, base_url: str, model_name: str, max_tokens: int, temperature: float = 0.7): + def __init__(self, api_key: str, base_url: str, model_name: str, max_tokens: int, temperature: float = 0.7, timeout: Optional[int] = 600): self.base_url = ensure_openai_base_url_has_v1(base_url) self.api_key = api_key self.model_name = model_name self.max_tokens = max_tokens self.temperature = temperature + self.timeout = timeout self._client = ChatOpenAI( model=self.model_name, api_key=self.api_key, base_url=self.base_url, max_tokens=self.max_tokens, - temperature=self.temperature + temperature=self.temperature, + timeout=self.timeout ) def invoke(self, prompt: str) -> str: @@ -81,19 +82,21 @@ class OllamaAdapter(BaseLLMAdapter): Ollama 同样有一个 OpenAI-like /v1/chat 接口，可直接使用 ChatOpenAI。 但是通常 Ollama 默认本地服务在 http://localhost:11434，如果符合OpenAI风格即可直接传参。 """ - def __init__(self, api_key: str, base_url: str, model_name: str, max_tokens: int, temperature: float = 0.7): + def __init__(self, api_key: str, base_url: str, model_name: str, max_tokens: int, temperature: float = 0.7, timeout: Optional[int] = 600): self.base_url = ensure_openai_base_url_has_v1(base_url) self.api_key = api_key self.model_name = model_name self.max_tokens = max_tokens self.temperature = temperature + self.timeout = timeout self._client = ChatOpenAI( model=self.model_name, api_key=self.api_key, base_url=self.base_url, max_tokens=self.max_tokens, - temperature=self.temperature + temperature=self.temperature, + timeout=self.timeout ) def invoke(self, prompt: str) -> str: @@ -104,19 +107,21 @@ class OllamaAdapter(BaseLLMAdapter): return response.content class MLStudioAdapter(BaseLLMAdapter): - def __init__(self, api_key: str, base_url: str, model_name: str, max_tokens: int, temperature: float = 0.7): + def __init__(self, api_key: str, base_url: str, model_name: str, max_tokens: int, temperature: float = 0.7, timeout: Optional[int] = 600): self.base_url = ensure_openai_base_url_has_v1(base_url) self.api_key = api_key self.model_name = model_name self.max_tokens = max_tokens self.temperature = temperature + self.timeout = timeout self._client = ChatOpenAI( model=self.model_name, api_key=self.api_key, base_url=self.base_url, max_tokens=self.max_tokens, - temperature=self.temperature + temperature=self.temperature, + timeout=self.timeout ) def invoke(self, prompt: str) -> str: diff --git a/novel_generator.py b/novel_generator.py index 6c4e47f..381055d 100644 --- a/novel_generator.py +++ b/novel_generator.py @@ -5,6 +5,7 @@ import logging import re import time import traceback +import json from typing import List, Optional, Tuple from langchain_chroma import Chroma @@ -46,6 +47,69 @@ from embedding_adapters import create_embedding_adapter logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") +# ============ 进度文件管理 ============ + +PROGRESS_FILE = "progress.json" + +def load_progress() -> dict: + """ + 简易进度文件读取，如果不存在则返回默认空字典。 + 你也可以在这里定制更多的进度信息。 + """ + if not os.path.exists(PROGRESS_FILE): + return { + "architecture_done": False, + "blueprint_done": False, + "blueprint_chunk_index": 1, # 若有分块生成，则记录当前分块的起始 + # 也可以记录已完成的章节 + "chapters_generated": [], # 已经生成草稿的章节列表 + "chapters_finalized": [] # 已经定稿的章节列表 + } + try: + with open(PROGRESS_FILE, "r", encoding="utf-8") as f: + return json.load(f) + except Exception: + return { + "architecture_done": False, + "blueprint_done": False, + "blueprint_chunk_index": 1, + "chapters_generated": [], + "chapters_finalized": [] + } + +def save_progress(progress: dict): + """ + 将进度写入到 progress.json 中。 + """ + with open(PROGRESS_FILE, "w", encoding="utf-8") as f: + json.dump(progress, f, ensure_ascii=False, indent=2) + + +# ============ 通用的重试封装 ============ + +def call_with_retry(func, max_retries=3, sleep_time=2, fallback_return=None, **kwargs): + """ + 通用的重试机制封装。 + :param func: 要执行的函数 + :param max_retries: 最大重试次数 + :param sleep_time: 重试前的等待秒数 + :param fallback_return: 如果多次重试仍失败时的返回值 + :param kwargs: 传给func的命名参数 + :return: func的结果，若失败则返回 fallback_return + """ + for attempt in range(1, max_retries + 1): + try: + return func(**kwargs) + except Exception as e: + logging.warning(f"[call_with_retry] Attempt {attempt} failed with error: {e}") + traceback.print_exc() + if attempt < max_retries: + time.sleep(sleep_time) + else: + logging.error("Max retries reached, returning fallback_return.") + return fallback_return + + # ============ 工具函数 ============ def remove_think_tags(text: str) -> str: @@ -61,10 +125,16 @@ def debug_log(prompt: str, response_content: str): ) def invoke_with_cleaning(llm_adapter, prompt: str) -> str: - """通用封装：调用 LLM，并移除 ... 文本，记录日志后返回""" - response = llm_adapter.invoke(prompt) + """ + 对 LLM 的调用增加了重试封装， + 如果多次失败，则返回空字符串以继续流程，而不是中断。 + """ + def _invoke(prompt): + return llm_adapter.invoke(prompt) + + response = call_with_retry(func=_invoke, max_retries=3, fallback_return="", prompt=prompt) if not response: - logging.warning("No response from model.") + logging.warning("No response from model after retry. Return empty.") return "" cleaned_text = remove_think_tags(response) debug_log(prompt, cleaned_text) @@ -76,6 +146,7 @@ def invoke_with_cleaning(llm_adapter, prompt: str) -> str: def get_vectorstore_dir(filepath: str) -> str: return os.path.join(filepath, "vectorstore") + # ============ 清空向量库 ============ def clear_vector_store(filepath: str) -> bool: @@ -100,35 +171,52 @@ def init_vector_store( embedding_adapter, texts: List[str], filepath: str -) -> Chroma: +) -> Optional[Chroma]: """ 在 filepath 下创建/加载一个 Chroma 向量库并插入 texts。 - 这里 embedding_adapter 是一个实现了 embed_documents(texts) 的对象 + 如果Embedding失败，则返回 None，不中断任务。 """ + from langchain.embeddings.base import Embeddings as LCEmbeddings + store_dir = get_vectorstore_dir(filepath) os.makedirs(store_dir, exist_ok=True) documents = [Document(page_content=str(t)) for t in texts] - from langchain.embeddings.base import Embeddings as LCEmbeddings + # 包一层try，如果embedding在初始化或插入过程中报错，则跳过 + try: + class LCEmbeddingWrapper(LCEmbeddings): + def embed_documents(self, doc_texts: List[str]) -> List[List[float]]: + return call_with_retry( + func=embedding_adapter.embed_documents, + max_retries=3, + fallback_return=[], + doc_texts=doc_texts + ) - class LCEmbeddingWrapper(LCEmbeddings): - def embed_documents(self, doc_texts: List[str]) -> List[List[float]]: - return embedding_adapter.embed_documents(doc_texts) + def embed_query(self, query_text: str) -> List[float]: + res = call_with_retry( + func=embedding_adapter.embed_query, + max_retries=3, + fallback_return=[], + query_text=query_text + ) + return res - def embed_query(self, query_text: str) -> List[float]: - return embedding_adapter.embed_query(query_text) + chroma_embedding = LCEmbeddingWrapper() - chroma_embedding = LCEmbeddingWrapper() - - vectorstore = Chroma.from_documents( - documents, - embedding=chroma_embedding, - persist_directory=store_dir, - client_settings=Settings(anonymized_telemetry=False), - collection_name="novel_collection" - ) - return vectorstore + vectorstore = Chroma.from_documents( + documents, + embedding=chroma_embedding, + persist_directory=store_dir, + client_settings=Settings(anonymized_telemetry=False), + collection_name="novel_collection" + ) + return vectorstore + except Exception as e: + logging.warning(f"Init vector store failed: {e}") + traceback.print_exc() + return None def load_vector_store( embedding_adapter, @@ -136,6 +224,7 @@ def load_vector_store( ) -> Optional[Chroma]: """ 读取已存在的 Chroma 向量库。若不存在则返回 None。 + 如果加载失败（embedding 或IO问题），则返回 None。 """ store_dir = get_vectorstore_dir(filepath) if not os.path.exists(store_dir): @@ -144,21 +233,37 @@ def load_vector_store( from langchain.embeddings.base import Embeddings as LCEmbeddings - class LCEmbeddingWrapper(LCEmbeddings): - def embed_documents(self, doc_texts: List[str]) -> List[List[float]]: - return embedding_adapter.embed_documents(doc_texts) + try: + class LCEmbeddingWrapper(LCEmbeddings): + def embed_documents(self, doc_texts: List[str]) -> List[List[float]]: + return call_with_retry( + func=embedding_adapter.embed_documents, + max_retries=3, + fallback_return=[], + doc_texts=doc_texts + ) - def embed_query(self, query_text: str) -> List[float]: - return embedding_adapter.embed_query(query_text) + def embed_query(self, query_text: str) -> List[float]: + res = call_with_retry( + func=embedding_adapter.embed_query, + max_retries=3, + fallback_return=[], + query_text=query_text + ) + return res - chroma_embedding = LCEmbeddingWrapper() + chroma_embedding = LCEmbeddingWrapper() - return Chroma( - persist_directory=store_dir, - embedding_function=chroma_embedding, - client_settings=Settings(anonymized_telemetry=False), - collection_name="novel_collection" - ) + return Chroma( + persist_directory=store_dir, + embedding_function=chroma_embedding, + client_settings=Settings(anonymized_telemetry=False), + collection_name="novel_collection" + ) + except Exception as e: + logging.warning(f"Failed to load vector store: {e}") + traceback.print_exc() + return None # ============ 文本分段工具 ============ @@ -218,6 +323,7 @@ def split_text_for_vectorstore(chapter_text: str, return final_segments + # ============ 更新向量库 ============ def update_vector_store( @@ -226,7 +332,8 @@ def update_vector_store( filepath: str ): """ - 将最新章节文本插入到向量库中。若库不存在则初始化。 + 将最新章节文本插入到向量库中。 + 若库不存在则初始化；若初始化/更新失败，则跳过。 """ splitted_texts = split_text_for_vectorstore(new_chapter) if not splitted_texts: @@ -235,13 +342,23 @@ def update_vector_store( store = load_vector_store(embedding_adapter, filepath) if not store: - logging.info("Vector store does not exist. Initializing a new one for new chapter...") - init_vector_store(embedding_adapter, splitted_texts, filepath) + logging.info("Vector store does not exist or failed to load. Initializing a new one for new chapter...") + store = init_vector_store(embedding_adapter, splitted_texts, filepath) + if not store: + logging.warning("Init vector store failed, skip embedding.") + else: + logging.info("New vector store created successfully.") return - 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.") + # 如果已有store，则直接往里插入 + try: + 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.") + except Exception as e: + logging.warning(f"Failed to update vector store: {e}") + traceback.print_exc() + # ============ 向量检索上下文 ============ @@ -253,19 +370,25 @@ def get_relevant_context_from_vector_store( ) -> str: """ 从向量库中检索与 query 最相关的 k 条文本，拼接后返回。 + 如果向量库加载/检索失败，则返回空字符串。 """ store = load_vector_store(embedding_adapter, filepath) if not store: - logging.info("No vector store found. Returning empty context.") + logging.info("No vector store found or load failed. Returning empty context.") return "" - docs = store.similarity_search(query, k=k) - if not docs: - logging.info(f"No relevant documents found for query '{query}'. Returning empty context.") + try: + 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]) + return combined + except Exception as e: + logging.warning(f"Similarity search failed: {e}") + traceback.print_exc() return "" - combined = "\n".join([d.page_content for d in docs]) - return combined # ============ 从目录中获取最近 n 章文本 ============ @@ -281,6 +404,7 @@ def get_last_n_chapters_text(chapters_dir: str, current_chapter_num: int, n: int texts.append("") return texts + # ============ 提炼(短期摘要, 下一章关键字) ============ def summarize_recent_chapters( @@ -350,7 +474,13 @@ def Novel_architecture_generate( 3. world_building_prompt 4. plot_architecture_prompt 最终输出 Novel_architecture.txt + 如果已生成，则不重复执行（利用 progress.json 中的标记）。 """ + progress = load_progress() + if progress.get("architecture_done", False): + logging.info("Novel architecture generation is already done. Skip.") + return + os.makedirs(filepath, exist_ok=True) llm_adapter = create_llm_adapter( @@ -405,6 +535,10 @@ def Novel_architecture_generate( save_string_to_txt(final_content, arch_file) logging.info("Novel_architecture.txt has been generated successfully.") + # 更新进度 + progress["architecture_done"] = True + save_progress(progress) + # ============ 计算分块大小的工具函数 ============ @@ -447,7 +581,14 @@ def Chapter_blueprint_generate( 2) 先生成 [1..chunk_size] 的章节 3) 将生成的文本作为已有目录传入，继续生成 [chunk_size+1..] 的章节 4) 最后汇总全部章节目录写入 Novel_directory.txt + + 过程中若发生错误，会进行一定次数重试；若仍失败则保留已生成的结果，方便下次中断续作。 """ + progress = load_progress() + if progress.get("blueprint_done", False): + logging.info("Chapter blueprint generation is already done. Skip.") + return + 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.") @@ -486,11 +627,14 @@ def Chapter_blueprint_generate( clear_file_content(filename_dir) save_string_to_txt(blueprint_text, filename_dir) logging.info("Novel_directory.txt (chapter blueprint) has been generated successfully (single-shot).") + + progress["blueprint_done"] = True + save_progress(progress) return # 否则，分块生成 final_blueprint = "" - current_start = 1 + current_start = progress.get("blueprint_chunk_index", 1) # 若之前中断，则从上一次的 chunk index 开始 while current_start <= number_of_chapters: current_end = min(current_start + chunk_size - 1, number_of_chapters) @@ -515,20 +659,29 @@ def Chapter_blueprint_generate( else: final_blueprint = chunk_result + # 更新下一个块 current_start = current_end + 1 + # 将当前的 final_blueprint 写入文件，以便中断后保留 + filename_dir = os.path.join(filepath, "Novel_directory.txt") + clear_file_content(filename_dir) + save_string_to_txt(final_blueprint.strip(), filename_dir) + + # 更新进度，以便中断后能接着来 + progress["blueprint_chunk_index"] = current_start + save_progress(progress) + if not final_blueprint.strip(): logging.warning("All chunked generation results are empty, cannot create blueprint.") return - filename_dir = os.path.join(filepath, "Novel_directory.txt") - clear_file_content(filename_dir) - save_string_to_txt(final_blueprint.strip(), filename_dir) - + # 生成完成 logging.info("Novel_directory.txt (chapter blueprint) has been generated successfully (chunked).") + progress["blueprint_done"] = True + save_progress(progress) -# ============ 3) 生成章节草稿（分「第一章」与「后续章节」） ============ +# ============ 3) 生成章节草稿 ============ def generate_chapter_draft( api_key: str, @@ -555,7 +708,16 @@ def generate_chapter_draft( 根据 novel_number 判断是否为第一章。 - 若是第一章，则使用 first_chapter_draft_prompt - 否则使用 next_chapter_draft_prompt + 生成草稿后存入 chapters/chapter_{novel_number}.txt """ + progress = load_progress() + if novel_number in progress.get("chapters_generated", []): + logging.info(f"Chapter {novel_number} draft already generated. Skip.") + # 直接返回已有内容 + chapters_dir = os.path.join(filepath, "chapters") + chapter_file = os.path.join(chapters_dir, f"chapter_{novel_number}.txt") + return read_file(chapter_file) + arch_file = os.path.join(filepath, "Novel_architecture.txt") novel_architecture_text = read_file(arch_file) @@ -582,7 +744,7 @@ def generate_chapter_draft( chapters_dir = os.path.join(filepath, "chapters") os.makedirs(chapters_dir, exist_ok=True) - # 如果是第一章，不需要前情检索与前章结尾 + # 根据是否是第一章，选择不同的 Prompt if novel_number == 1: # 使用第一章提示词 prompt_text = first_chapter_draft_prompt.format( @@ -603,7 +765,6 @@ def generate_chapter_draft( novel_setting=novel_architecture_text ) - else: # 若不是第一章，则先获取最近几章文本，并做摘要与检索 recent_3_texts = get_last_n_chapters_text(chapters_dir, novel_number, n=3) @@ -627,7 +788,7 @@ def generate_chapter_draft( previous_chapter_excerpt = text_block break - # 从向量库检索上下文 + # 从向量库检索上下文（若失败则为空，不中断） embedding_adapter = create_embedding_adapter( embedding_interface_format, embedding_api_key, @@ -687,6 +848,11 @@ def generate_chapter_draft( save_string_to_txt(chapter_content, chapter_file) logging.info(f"[Draft] Chapter {novel_number} generated as a draft.") + + # 更新进度 + progress["chapters_generated"].append(novel_number) + save_progress(progress) + return chapter_content @@ -707,6 +873,11 @@ def finalize_chapter( interface_format: str, max_tokens: int ): + progress = load_progress() + if novel_number in progress.get("chapters_finalized", []): + logging.info(f"Chapter {novel_number} is already finalized. Skip.") + return + 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() @@ -755,7 +926,7 @@ def finalize_chapter( clear_file_content(character_state_file) save_string_to_txt(new_char_state, character_state_file) - # 更新向量库 + # 更新向量库（若失败则跳过） embedding_adapter = create_embedding_adapter( embedding_interface_format, embedding_api_key, @@ -766,6 +937,11 @@ def finalize_chapter( logging.info(f"Chapter {novel_number} has been finalized.") + # 更新进度 + progress["chapters_finalized"].append(novel_number) + save_progress(progress) + + def enrich_chapter_text( chapter_text: str, word_number: int, @@ -861,9 +1037,17 @@ def import_knowledge_file( store = load_vector_store(embedding_adapter, filepath) if not store: - logging.info("Vector store does not exist. Initializing a new one for knowledge import...") - init_vector_store(embedding_adapter, paragraphs, filepath) + logging.info("Vector store does not exist or load failed. Initializing a new one for knowledge import...") + store = init_vector_store(embedding_adapter, paragraphs, filepath) + if store: + logging.info("知识库文件已成功导入至向量库(新初始化)。") + else: + logging.warning("知识库导入失败，跳过。") else: - docs = [Document(page_content=str(p)) for p in paragraphs] - store.add_documents(docs) - logging.info("知识库文件已成功导入至向量库。") + try: + docs = [Document(page_content=str(p)) for p in paragraphs] + store.add_documents(docs) + logging.info("知识库文件已成功导入至向量库(追加模式)。") + except Exception as e: + logging.warning(f"知识库导入失败: {e}") + traceback.print_exc()