AI_NovelGenerator/novel_generator.py

import os
import logging
from typing import Dict, List, Optional
try:
    from typing import TypedDict  # Python 3.8+ 直接可用；若是3.7可改用 typing_extensions
except ImportError:
    from typing_extensions import TypedDict

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.docstore.document import Document

from utils import (
    read_file, append_text_to_file, clear_file_content,
    save_string_to_txt
)
from prompt_definitions import (
    set_prompt, character_prompt, dark_lines_prompt,
    finalize_setting_prompt, novel_directory_prompt,
    summary_prompt, update_character_state_prompt,
    chapter_outline_prompt, chapter_write_prompt
)

# ============ 日志配置（可选） ============
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

# ============ 向量检索相关函数（Chroma） ============

VECTOR_STORE_DIR = "vectorstore"

def init_vector_store(api_key: str, base_url: str, texts: List[str]) -> Chroma:
    """
    初始化并返回一个Chroma向量库，将传入的文本进行嵌入并保存到本地目录。
    如果不存在该目录，会自动创建。
    """
    embeddings = OpenAIEmbeddings(
        openai_api_key=api_key,
        openai_api_base=base_url  # <-- 这里用传进来的 base_url
    )
    documents = [Document(page_content=t) for t in texts]
    vectorstore = Chroma.from_documents(
        documents,
        embedding=embeddings,
        persist_directory=VECTOR_STORE_DIR
    )
    vectorstore.persist()
    return vectorstore


def load_vector_store(api_key: str, base_url: str) -> Optional[Chroma]:
    """读取已存在的向量库。若不存在则返回 None。"""
    if not os.path.exists(VECTOR_STORE_DIR):
        return None
    embeddings = OpenAIEmbeddings(
        openai_api_key=api_key,
        openai_api_base=base_url  # <-- 使用 base_url
    )
    return Chroma(persist_directory=VECTOR_STORE_DIR, embedding_function=embeddings)


def update_vector_store(api_key: str, base_url: str, new_chapter: str) -> None:
    """将最新章节文本插入到向量库里，用于后续检索参考。若库不存在则初始化。"""
    store = load_vector_store(api_key, base_url)
    if not store:
        logging.info("Vector store does not exist. Initializing a new one...")
        init_vector_store(api_key, base_url, [new_chapter])
        return

    new_doc = Document(page_content=new_chapter)
    store.add_documents([new_doc])
    store.persist()


def get_relevant_context_from_vector_store(api_key: str, base_url: str, query: str, k: int = 2) -> str:
    """
    从向量库中检索与 query 最相关的 k 条文本，拼接后返回。
    若向量库不存在则返回空字符串。
    """
    store = load_vector_store(api_key, base_url)
    if not store:
        logging.warning("Vector store not found. Returning empty context.")
        return ""
    docs = store.similarity_search(query, k=k)
    combined = "\n".join([d.page_content for d in docs])
    return combined


# ============ 多步生成：设置 & 目录 ============

class OverallState(TypedDict):
    topic: str
    genre: str
    number_of_chapters: int
    word_number: int
    novel_setting_base: str
    character_setting: str
    dark_lines: str
    final_novel_setting: str
    novel_directory: str


def Novel_novel_directory_generate(
    api_key: str,
    base_url: str,
    llm_model: str,
    topic: str,
    genre: str,
    number_of_chapters: int,
    word_number: int,
    filepath: str
) -> None:
    """
    使用多步流程，生成 Novel_setting.txt 与 Novel_directory.txt 并保存到 filepath。

    :param api_key:       OpenAI API key
    :param base_url:      OpenAI API base url
    :param llm_model:     所使用的 LLM 模型名称
    :param topic:         小说主题
    :param genre:         小说类型
    :param number_of_chapters: 章节数
    :param word_number:   单章目标字数
    :param filepath:      存放生成文件的目录路径
    """
    # 确保文件夹存在
    os.makedirs(filepath, exist_ok=True)

    model = ChatOpenAI(
        model=llm_model,
        api_key=api_key,
        base_url=base_url
    )

    def generate_base_setting(state: OverallState) -> Dict[str, str]:
        prompt = set_prompt.format(
            topic=state["topic"],
            genre=state["genre"],
            number_of_chapters=state["number_of_chapters"],
            word_number=state["word_number"]
        )
        response = model.invoke(prompt)
        if not response:
            logging.warning("generate_base_setting: No response.")
            return {"novel_setting_base": ""}
        return {"novel_setting_base": response.content.strip()}

    def generate_character_setting(state: OverallState) -> Dict[str, str]:
        prompt = character_prompt.format(
            novel_setting=state["novel_setting_base"]
        )
        response = model.invoke(prompt)
        if not response:
            logging.warning("generate_character_setting: No response.")
            return {"character_setting": ""}
        return {"character_setting": response.content.strip()}

    def generate_dark_lines(state: OverallState) -> Dict[str, str]:
        prompt = dark_lines_prompt.format(
            character_info=state["character_setting"]
        )
        response = model.invoke(prompt)
        if not response:
            logging.warning("generate_dark_lines: No response.")
            return {"dark_lines": ""}
        return {"dark_lines": response.content.strip()}

    def finalize_novel_setting(state: OverallState) -> Dict[str, str]:
        prompt = finalize_setting_prompt.format(
            novel_setting_base=state["novel_setting_base"],
            character_setting=state["character_setting"],
            dark_lines=state["dark_lines"]
        )
        response = model.invoke(prompt)
        if not response:
            logging.warning("finalize_novel_setting: No response.")
            return {"final_novel_setting": ""}
        return {"final_novel_setting": response.content.strip()}

    def generate_novel_directory(state: OverallState) -> Dict[str, str]:
        prompt = novel_directory_prompt.format(
            final_novel_setting=state["final_novel_setting"],
            number_of_chapters=state["number_of_chapters"]
        )
        response = model.invoke(prompt)
        if not response:
            logging.warning("generate_novel_directory: No response.")
            return {"novel_directory": ""}
        return {"novel_directory": response.content.strip()}

    # 构建状态图
    graph = StateGraph(OverallState)
    graph.add_node("generate_base_setting", generate_base_setting)
    graph.add_node("generate_character_setting", generate_character_setting)
    graph.add_node("generate_dark_lines", generate_dark_lines)
    graph.add_node("finalize_novel_setting", finalize_novel_setting)
    graph.add_node("generate_novel_directory", generate_novel_directory)

    # 注意修正此处节点名称
    graph.add_edge(START, "generate_base_setting")
    graph.add_edge("generate_base_setting", "generate_character_setting")
    graph.add_edge("generate_character_setting", "generate_dark_lines")
    graph.add_edge("generate_dark_lines", "finalize_novel_setting")
    graph.add_edge("finalize_novel_setting", "generate_novel_directory")
    graph.add_edge("generate_novel_directory", END)

    app = graph.compile()

    input_params = {
        "topic": topic,
        "genre": genre,
        "number_of_chapters": number_of_chapters,
        "word_number": word_number
    }
    result = app.invoke(input_params)

    if not result:
        logging.warning("Novel_novel_directory_generate: invoke() 结果为空，生成失败。")
        return

    final_novel_setting = result.get("final_novel_setting", "")
    final_novel_directory = result.get("novel_directory", "")

    if not final_novel_setting or not final_novel_directory:
        logging.warning("生成失败：缺少 final_novel_setting 或 novel_directory。")
        return

    # 写入文件
    filename_set = os.path.join(filepath, "Novel_setting.txt")
    filename_novel_directory = os.path.join(filepath, "Novel_directory.txt")

    # 清理文本（去除多余 # 或 * 等）
    def clean_text(txt: str) -> str:
        return txt.replace('#', '').replace('*', '')

    final_novel_setting_cleaned = clean_text(final_novel_setting)
    final_novel_directory_cleaned = clean_text(final_novel_directory)

    # 以追加方式保存；如果希望覆盖可改为 save_string_to_txt()
    append_text_to_file(final_novel_setting_cleaned, filename_set)
    append_text_to_file(final_novel_directory_cleaned, filename_novel_directory)

    logging.info("Novel settings and directory generated successfully.")


# ============ 生成章节（每章独立文件） ============

def generate_chapter_with_state(
    novel_settings: str,
    novel_novel_directory: str,
    api_key: str,
    base_url: str,
    model_name: str,
    novel_number: int,
    filepath: str,
    word_number: int,
    lastchapter: str
) -> str:
    """
    多步流程:
      1) 更新/创建全局摘要
      2) 更新/生成角色状态文档
      3) 向量检索获取往期上下文
      4) 大纲 -> 正文
      5) 写入 chapter_{novel_number}.txt, 更新 lastchapter.txt
      6) 更新向量库

    :param novel_settings:        最终的作品设定（字符串）
    :param novel_novel_directory: 小说目录信息（此处暂时未使用，可根据需求做扩展）
    :param api_key:              OpenAI API Key
    :param base_url:             OpenAI Base URL
    :param model_name:           LLM 模型名称
    :param novel_number:         当前要生成的章节号
    :param filepath:             文件存放的目录
    :param word_number:          单章目标字数
    :param lastchapter:          上一章内容（若为空字符串，表示无上一章）
    :return:                     本章生成的正文内容
    """
    # 确保文件夹存在
    os.makedirs(filepath, exist_ok=True)

    model = ChatOpenAI(
        model=model_name,
        api_key=api_key,
        base_url=base_url,
        temperature=0.9
    )

    # --- 文件路径定义 ---
    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")
    lastchapter_file = os.path.join(filepath, "lastchapter.txt")
    character_state_file = os.path.join(filepath, "character_state.txt")
    global_summary_file = os.path.join(filepath, "global_summary.txt")

    old_char_state = read_file(character_state_file)
    old_global_summary = read_file(global_summary_file)

    # 1) 更新全局摘要
    def update_global_summary(chapter_text: str, old_summary: str) -> str:
        prompt = summary_prompt.format(
            chapter_text=chapter_text,
            global_summary=old_summary
        )
        response = model.invoke(prompt)
        if not response:
            logging.warning("update_global_summary: No response.")
            return old_summary
        return response.content.strip()

    if lastchapter.strip():
        new_global_summary = update_global_summary(lastchapter, old_global_summary)
    else:
        new_global_summary = old_global_summary

    # 2) 更新角色状态文档
    def update_character_state(chapter_text: str, old_state: str) -> str:
        prompt = update_character_state_prompt.format(
            chapter_text=chapter_text,
            old_state=old_state
        )
        response = model.invoke(prompt)
        if not response:
            logging.warning("update_character_state: No response.")
            return old_state
        return response.content.strip()

    if lastchapter.strip():
        new_char_state = update_character_state(lastchapter, old_char_state)
    else:
        new_char_state = old_char_state

    # 3) 从向量库检索上下文
    relevant_context = get_relevant_context_from_vector_store(
        api_key, base_url, "回顾剧情", k=2  # <-- 多传一个 base_url
    )

    # 4) 生成大纲
    def outline_chapter(
        novel_setting: str,
        char_state: str,
        global_summary: str,
        chap_num: int,
        extra_context: str
    ) -> str:
        prompt = chapter_outline_prompt.format(
            novel_setting=novel_setting,
            character_state=char_state + "\n\n【历史上下文】\n" + extra_context,
            global_summary=global_summary,
            novel_number=chap_num
        )
        response = model.invoke(prompt)
        if not response:
            logging.warning("outline_chapter: No response.")
            return ""
        return response.content.strip()

    chap_outline = outline_chapter(
        novel_settings, new_char_state, new_global_summary, novel_number, relevant_context
    )

    # 5) 生成正文
    def write_chapter(
        novel_setting: str,
        char_state: str,
        global_summary: str,
        outline: str,
        wnum: int,
        extra_context: str
    ) -> str:
        prompt = chapter_write_prompt.format(
            novel_setting=novel_setting,
            character_state=char_state + "\n\n【历史上下文】\n" + extra_context,
            global_summary=global_summary,
            chapter_outline=outline,
            word_number=wnum
        )
        response = model.invoke(prompt)
        if not response:
            logging.warning("write_chapter: No response.")
            return ""
        return response.content.strip()

    chapter_content = write_chapter(
        novel_settings,
        new_char_state,
        new_global_summary,
        chap_outline,
        word_number,
        relevant_context
    )

    # 写入文件并更新记录
    if chapter_content:
        save_string_to_txt(chapter_content, chapter_file)

        # 更新 lastchapter.txt
        clear_file_content(lastchapter_file)
        save_string_to_txt(chapter_content, lastchapter_file)

        # 更新角色状态、全局摘要
        clear_file_content(character_state_file)
        save_string_to_txt(new_char_state, character_state_file)

        clear_file_content(global_summary_file)
        save_string_to_txt(new_global_summary, global_summary_file)

        # 6) 更新向量检索库
        update_vector_store(api_key, base_url, chapter_content)
        logging.info(f"Chapter {novel_number} generated successfully.")
    else:
        logging.warning(f"Chapter {novel_number} generation failed.")

    return chapter_content