AI_NovelGenerator/novel_generator.py

# novel_generator.py
# -*- coding: utf-8 -*-
import os
import logging
import re
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

# 
import nltk
import math
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

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
    )
    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
    )
    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,
    temperature: float = 0.7
) -> 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:      存放生成文件的目录路径
    :param temperature:   生成温度
    """
    # 确保文件夹存在
    os.makedirs(filepath, exist_ok=True)

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

    def debug_log(prompt: str, response_content: str):
        """在控制台打印或记录下每次Prompt与Response，[调试]"""
        logging.info(f"\n[Prompt >>>] {prompt}\n")
        logging.info(f"[Response <<<] {response_content}\n")

    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": ""}
        debug_log(prompt, response.content)
        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": ""}
        debug_log(prompt, response.content)
        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": ""}
        debug_log(prompt, response.content)
        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": ""}
        debug_log(prompt, response.content)
        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": ""}
        debug_log(prompt, response.content)
        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)

    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.")

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

CHINESE_NUM_MAP = {
    '零': 0, '○': 0, '〇': 0,
    '一': 1, '二': 2, '三': 3, '四': 4, '五': 5,
    '六': 6, '七': 7, '八': 8, '九': 9,
    '十': 10, '百': 100, '千': 1000, '万': 10000
}

def chinese_to_arabic(chinese_str: str) -> int:
    """
    只能处理到万(10000)以内的中文数字，正常小说章节应该够用了
    """
    total = 0
    current_unit = 1  # 记录当前单位
    tmp_val = 0       # 暂存本轮数字

    for char in reversed(chinese_str):
        if char in CHINESE_NUM_MAP:
            val = CHINESE_NUM_MAP[char]
            if val >= 10:
                if val > current_unit:
                    # 如 100, 1000, 10000
                    current_unit = val
                else:
                    # 比如 “十二” -> 2 * 10 + 1
                    # 如果 val <= current_unit, 那么相当于在这个单位下加
                    total += tmp_val * val
                    tmp_val = 0
            else:
                # 0~9
                tmp_val = tmp_val + val * current_unit
        else:
            # 非中文数字字符，视情况决定怎么处理，这里直接跳过
            pass
    
    total += tmp_val
    return total

def parse_chapter_title_from_directory(novel_directory_text: str, 
                                       novel_number: int, 
                                       range_size: int = 1) -> str:
    """
    从小说目录文本中，提取指定章节（以及前后几章）的目录信息。
    range_size=1，表示获取当前章节、前一章和后一章的目录信息(若存在)。
    支持多种常见的章节格式。
    """

    lines = novel_directory_text.splitlines()

    # 可以根据需求自行扩展，这里列举了几种常见的章节标题格式，如果模型实在不听话，可以适当调整
    # 每个pattern都应该捕获两个组：
    # 1. chapter_num_str：章节数字(可能是中文也可能是阿拉伯数字)
    # 2. chapter_title ：章节标题（.*）
    patterns = [
        # 1) 第12章 标题
        r"^第\s*([\d]+)\s*章[:：]?\s*(.*)$",
        # 2) 第十二章 标题（中文数字）
        r"^第\s*([零○〇一二三四五六七八九十百千万]+)\s*章[:：]?\s*(.*)$",
        # 3) Chapter 12 标题
        r"^Chapter\s+(\d+)\s*[:：]?\s*(.*)$",
        # 4) Ch 12 标题
        r"^Ch\s+(\d+)\s*[:：]?\s*(.*)$",
        # 5) 第12节 标题
        r"^第\s*([\d]+)\s*节[:：]?\s*(.*)$",
        # 6) 第12话 标题
        r"^第\s*([\d]+)\s*话[:：]?\s*(.*)$",
        # ... 更多模式 ...
    ]

    # 用来存储匹配结果： chapter_num -> title
    directory_map = {}

    for line in lines:
        line = line.strip()
        if not line:
            continue
        
        # 依次尝试每一种pattern
        matched = False
        for pat in patterns:
            match = re.match(pat, line, flags=re.IGNORECASE)
            if match:
                chapter_num_str = match.group(1)
                chapter_title   = match.group(2).strip()

                # 如果是中文数字，需要转换
                # 如果是阿拉伯数字，直接转 int 即可
                if re.match(r"^[零○〇一二三四五六七八九十百千万]+$", chapter_num_str):
                    chapter_num = chinese_to_arabic(chapter_num_str)
                else:
                    chapter_num = int(chapter_num_str)

                directory_map[chapter_num] = chapter_title
                matched = True
                break
        
        # 如果已经匹配到其中一个pattern，就不需要继续匹配剩余pattern
        if matched:
            continue

    # 收集需要的章节范围
    chapters_info = []
    for cnum in range(novel_number - range_size, novel_number + range_size + 1):
        if cnum in directory_map:
            if cnum == novel_number:
                chapters_info.append(f"【当前】第{cnum}章：{directory_map[cnum]}")
            else:
                chapters_info.append(f"第{cnum}章：{directory_map[cnum]}")

    if chapters_info:
        return "\n".join(chapters_info)
    return ""

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,
    user_guidance: str = "",
    temperature: float = 0.7
) -> str:
    """
    多步流程:
      1) 更新/创建全局摘要
      2) 更新/生成角色状态文档
      3) 向量检索获取往期上下文
      4) 从Novel_directory.txt中获取当前(和前后几章)的目录信息
      5) 大纲 -> 正文（可结合用户给出的额外指导）
      6) 写入 chapter_{novel_number}.txt, 更新 last_chapter.txt
      7) 更新向量库

    :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:          上一章内容（若为空字符串，表示无上一章）
    :param user_guidance:        用户对当前章节的额外指导或想法
    :param temperature:          生成温度
    :return:                     本章生成的正文内容
    """
    # 确保文件夹存在
    os.makedirs(filepath, exist_ok=True)

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

    # 调试输出函数
    def debug_log(prompt: str, response_content: str):
        """在控制台打印或记录下每次的 Prompt 与 Response，便于观察生成过程。"""
        logging.info(f"\n[Prompt >>>]\n{prompt}\n")
        logging.info(f"[Response <<<]\n{response_content}\n")

    # --- 文件路径定义 ---
    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, "last_chapter.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
        debug_log(prompt, response.content)
        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
        debug_log(prompt, response.content)
        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
    )

    # 4) 解析本章及前后章节目录信息
    this_and_related_chapters = parse_chapter_title_from_directory(novel_novel_directory, novel_number, range_size=1)

    # 5) 生成大纲
    def outline_chapter(
        novel_setting: str,
        char_state: str,
        global_summary: str,
        chap_num: int,
        extra_context: str,
        directory_hint: str,
        user_guide: str
    ) -> str:
        """
        将目录提示以及用户额外指导内容一起放入 Prompt 中。
        """
        # 适度修改章节提纲提示词，以整合目录信息 & 用户指导
        outline_prompt = (
            chapter_outline_prompt
            + "\n\n【目录参考】\n" + directory_hint
            + "\n\n【用户指导】\n" + user_guide
        ).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(outline_prompt)
        if not response:
            logging.warning("outline_chapter: No response.")
            return ""
        debug_log(outline_prompt, response.content)
        return response.content.strip()

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

    # 6) 生成正文
    def write_chapter(
        novel_setting: str,
        char_state: str,
        global_summary: str,
        outline: str,
        wnum: int,
        extra_context: str,
        directory_hint: str,
        user_guide: str
    ) -> str:
        # 同理，整合目录信息和用户指导
        writing_prompt = (
            chapter_write_prompt
            + "\n\n【目录参考】\n" + directory_hint
            + "\n\n【用户指导】\n" + user_guide
        ).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(writing_prompt)
        if not response:
            logging.warning("write_chapter: No response.")
            return ""
        debug_log(writing_prompt, response.content)
        return response.content.strip()

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

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

        # 更新 last_chapter.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)

        # 7) 更新向量检索库
        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

def import_knowledge_file(api_key: str, base_url: str, file_path: str) -> None:
    """
    将用户选定的文本文件导入到向量库，以便在写作时检索。
    可以在UI中提供按钮来调用此函数。
    """

    # 1. 检查文件路径是否有效
    if not os.path.exists(file_path):
        logging.warning(f"知识库文件不存在: {file_path}")
        return

    # 2. 读取文件内容
    content = read_file(file_path)
    if not content.strip():
        logging.warning("知识库文件内容为空。")
        return

    # 3. 对内容进行高级切分处理
    paragraphs = advanced_split_content(content)

    # 4. 加载或初始化向量存储
    store = load_vector_store(api_key, base_url)
    if not store:
        logging.info("Vector store does not exist. Initializing a new one for knowledge import...")
        init_vector_store(api_key, base_url, paragraphs)
        return

    # 5. 创建Document对象并更新到向量库
    docs = [Document(page_content=p) for p in paragraphs]
    store.add_documents(docs)
    store.persist()
    logging.info("知识库文件已成功导入至向量库。")


def advanced_split_content(content: str,
                           similarity_threshold: float = 0.7,
                           max_length: int = 500) -> List[str]:
    """
    将文本先按句子切分，然后根据语义相似度进行合并，最后根据max_length进行二次切分。

    :param content: 原始文本内容
    :param similarity_threshold: 相邻句子合并的语义相似度阈值，小于此值则会开启新的段落
    :param max_length: 每个段落的最大长度（按字符数计算，超过则进一步拆分）
    :return: 切分好的段落列表
    """

    # 1. 按句子切分
    nltk.download('punkt', quiet=True)  # 确保 punkt 数据可用
    sentences = nltk.sent_tokenize(content)

    if not sentences:
        return []

    # 2. 加载 SentenceTransformer 模型，用于计算语义相似度
    model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
    embeddings = model.encode(sentences)

    # 3. 根据相邻句子的语义相似度合并段落
    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 = (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))

    # 4. 根据最大长度 max_length 做二次拆分，避免段落过长
    final_segments = []
    for para in merged_paragraphs:
        # 如果段落长度超过max_length，进一步切分
        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


def split_by_length(text: str, max_length: int = 500) -> List[str]:
    """
    将文本按照max_length进行拆分，以避免段落过长。
    这里以字符数为单位进行简单的拆分，也可以改为按词数或token数等。
    """
    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