363 lines
16 KiB
Python
363 lines
16 KiB
Python
|
# ui/generation_handlers.py
|
||
|
|
# -*- coding: utf-8 -*-
|
||
|
|
import os
|
||
|
|
import threading
|
||
|
|
import tkinter as tk
|
||
|
|
from tkinter import messagebox
|
||
|
|
import customtkinter as ctk
|
||
|
|
import traceback
|
||
|
|
from utils import read_file, save_string_to_txt, clear_file_content
|
||
|
|
from novel_generator import (
|
||
|
|
Novel_architecture_generate,
|
||
|
|
Chapter_blueprint_generate,
|
||
|
|
generate_chapter_draft,
|
||
|
|
finalize_chapter,
|
||
|
|
import_knowledge_file,
|
||
|
|
clear_vector_store,
|
||
|
|
enrich_chapter_text
|
||
|
|
)
|
||
|
|
from consistency_checker import check_consistency
|
||
|
|
|
||
|
|
def generate_novel_architecture_ui(self):
|
||
|
|
filepath = self.filepath_var.get().strip()
|
||
|
|
if not filepath:
|
||
|
|
messagebox.showwarning("警告", "请先选择保存文件路径")
|
||
|
|
return
|
||
|
|
|
||
|
|
def task():
|
||
|
|
self.disable_button_safe(self.btn_generate_architecture)
|
||
|
|
try:
|
||
|
|
interface_format = self.interface_format_var.get().strip()
|
||
|
|
api_key = self.api_key_var.get().strip()
|
||
|
|
base_url = self.base_url_var.get().strip()
|
||
|
|
model_name = self.model_name_var.get().strip()
|
||
|
|
temperature = self.temperature_var.get()
|
||
|
|
max_tokens = self.max_tokens_var.get()
|
||
|
|
timeout_val = self.safe_get_int(self.timeout_var, 600)
|
||
|
|
|
||
|
|
topic = self.topic_text.get("0.0", "end").strip()
|
||
|
|
genre = self.genre_var.get().strip()
|
||
|
|
num_chapters = self.safe_get_int(self.num_chapters_var, 10)
|
||
|
|
word_number = self.safe_get_int(self.word_number_var, 3000)
|
||
|
|
|
||
|
|
self.safe_log("开始生成小说架构...")
|
||
|
|
Novel_architecture_generate(
|
||
|
|
interface_format=interface_format,
|
||
|
|
api_key=api_key,
|
||
|
|
base_url=base_url,
|
||
|
|
llm_model=model_name,
|
||
|
|
topic=topic,
|
||
|
|
genre=genre,
|
||
|
|
number_of_chapters=num_chapters,
|
||
|
|
word_number=word_number,
|
||
|
|
filepath=filepath,
|
||
|
|
temperature=temperature,
|
||
|
|
max_tokens=max_tokens,
|
||
|
|
timeout=timeout_val
|
||
|
|
)
|
||
|
|
self.safe_log("✅ 小说架构生成完成。请在 'Novel Architecture' 标签页查看或编辑。")
|
||
|
|
except Exception:
|
||
|
|
self.handle_exception("生成小说架构时出错")
|
||
|
|
finally:
|
||
|
|
self.enable_button_safe(self.btn_generate_architecture)
|
||
|
|
threading.Thread(target=task, daemon=True).start()
|
||
|
|
|
||
|
|
def generate_chapter_blueprint_ui(self):
|
||
|
|
filepath = self.filepath_var.get().strip()
|
||
|
|
if not filepath:
|
||
|
|
messagebox.showwarning("警告", "请先选择保存文件路径")
|
||
|
|
return
|
||
|
|
|
||
|
|
def task():
|
||
|
|
self.disable_button_safe(self.btn_generate_directory)
|
||
|
|
try:
|
||
|
|
interface_format = self.interface_format_var.get().strip()
|
||
|
|
api_key = self.api_key_var.get().strip()
|
||
|
|
base_url = self.base_url_var.get().strip()
|
||
|
|
model_name = self.model_name_var.get().strip()
|
||
|
|
number_of_chapters = self.safe_get_int(self.num_chapters_var, 10)
|
||
|
|
temperature = self.temperature_var.get()
|
||
|
|
max_tokens = self.max_tokens_var.get()
|
||
|
|
timeout_val = self.safe_get_int(self.timeout_var, 600)
|
||
|
|
|
||
|
|
self.safe_log("开始生成章节蓝图...")
|
||
|
|
Chapter_blueprint_generate(
|
||
|
|
interface_format=interface_format,
|
||
|
|
api_key=api_key,
|
||
|
|
base_url=base_url,
|
||
|
|
llm_model=model_name,
|
||
|
|
number_of_chapters=number_of_chapters,
|
||
|
|
filepath=filepath,
|
||
|
|
temperature=temperature,
|
||
|
|
max_tokens=max_tokens,
|
||
|
|
timeout=timeout_val
|
||
|
|
)
|
||
|
|
self.safe_log("✅ 章节蓝图生成完成。请在 'Chapter Blueprint' 标签页查看或编辑。")
|
||
|
|
except Exception:
|
||
|
|
self.handle_exception("生成章节蓝图时出错")
|
||
|
|
finally:
|
||
|
|
self.enable_button_safe(self.btn_generate_directory)
|
||
|
|
threading.Thread(target=task, daemon=True).start()
|
||
|
|
|
||
|
|
def generate_chapter_draft_ui(self):
|
||
|
|
filepath = self.filepath_var.get().strip()
|
||
|
|
if not filepath:
|
||
|
|
messagebox.showwarning("警告", "请先配置保存文件路径。")
|
||
|
|
return
|
||
|
|
|
||
|
|
def task():
|
||
|
|
self.disable_button_safe(self.btn_generate_chapter)
|
||
|
|
try:
|
||
|
|
interface_format = self.interface_format_var.get().strip()
|
||
|
|
api_key = self.api_key_var.get().strip()
|
||
|
|
base_url = self.base_url_var.get().strip()
|
||
|
|
model_name = self.model_name_var.get().strip()
|
||
|
|
temperature = self.temperature_var.get()
|
||
|
|
max_tokens = self.max_tokens_var.get()
|
||
|
|
timeout_val = self.safe_get_int(self.timeout_var, 600)
|
||
|
|
|
||
|
|
chap_num = self.safe_get_int(self.chapter_num_var, 1)
|
||
|
|
word_number = self.safe_get_int(self.word_number_var, 3000)
|
||
|
|
user_guidance = self.user_guide_text.get("0.0", "end").strip()
|
||
|
|
|
||
|
|
char_inv = self.characters_involved_var.get().strip()
|
||
|
|
key_items = self.key_items_var.get().strip()
|
||
|
|
scene_loc = self.scene_location_var.get().strip()
|
||
|
|
time_constr = self.time_constraint_var.get().strip()
|
||
|
|
|
||
|
|
embedding_api_key = self.embedding_api_key_var.get().strip()
|
||
|
|
embedding_url = self.embedding_url_var.get().strip()
|
||
|
|
embedding_interface_format = self.embedding_interface_format_var.get().strip()
|
||
|
|
embedding_model_name = self.embedding_model_name_var.get().strip()
|
||
|
|
embedding_k = self.safe_get_int(self.embedding_retrieval_k_var, 4)
|
||
|
|
|
||
|
|
self.safe_log(f"开始生成第{chap_num}章草稿...")
|
||
|
|
draft_text = generate_chapter_draft(
|
||
|
|
api_key=api_key,
|
||
|
|
base_url=base_url,
|
||
|
|
model_name=model_name,
|
||
|
|
filepath=filepath,
|
||
|
|
novel_number=chap_num,
|
||
|
|
word_number=word_number,
|
||
|
|
temperature=temperature,
|
||
|
|
user_guidance=user_guidance,
|
||
|
|
characters_involved=char_inv,
|
||
|
|
key_items=key_items,
|
||
|
|
scene_location=scene_loc,
|
||
|
|
time_constraint=time_constr,
|
||
|
|
embedding_api_key=embedding_api_key,
|
||
|
|
embedding_url=embedding_url,
|
||
|
|
embedding_interface_format=embedding_interface_format,
|
||
|
|
embedding_model_name=embedding_model_name,
|
||
|
|
embedding_retrieval_k=embedding_k,
|
||
|
|
interface_format=interface_format,
|
||
|
|
max_tokens=max_tokens,
|
||
|
|
timeout=timeout_val
|
||
|
|
)
|
||
|
|
if draft_text:
|
||
|
|
self.safe_log(f"✅ 第{chap_num}章草稿生成完成。请在左侧查看或编辑。")
|
||
|
|
self.master.after(0, lambda: self.show_chapter_in_textbox(draft_text))
|
||
|
|
else:
|
||
|
|
self.safe_log("⚠️ 本章草稿生成失败或无内容。")
|
||
|
|
except Exception:
|
||
|
|
self.handle_exception("生成章节草稿时出错")
|
||
|
|
finally:
|
||
|
|
self.enable_button_safe(self.btn_generate_chapter)
|
||
|
|
threading.Thread(target=task, daemon=True).start()
|
||
|
|
|
||
|
|
def finalize_chapter_ui(self):
|
||
|
|
filepath = self.filepath_var.get().strip()
|
||
|
|
if not filepath:
|
||
|
|
messagebox.showwarning("警告", "请先配置保存文件路径。")
|
||
|
|
return
|
||
|
|
|
||
|
|
def task():
|
||
|
|
self.disable_button_safe(self.btn_finalize_chapter)
|
||
|
|
try:
|
||
|
|
interface_format = self.interface_format_var.get().strip()
|
||
|
|
api_key = self.api_key_var.get().strip()
|
||
|
|
base_url = self.base_url_var.get().strip()
|
||
|
|
model_name = self.model_name_var.get().strip()
|
||
|
|
temperature = self.temperature_var.get()
|
||
|
|
max_tokens = self.max_tokens_var.get()
|
||
|
|
timeout_val = self.safe_get_int(self.timeout_var, 600)
|
||
|
|
|
||
|
|
embedding_api_key = self.embedding_api_key_var.get().strip()
|
||
|
|
embedding_url = self.embedding_url_var.get().strip()
|
||
|
|
embedding_interface_format = self.embedding_interface_format_var.get().strip()
|
||
|
|
embedding_model_name = self.embedding_model_name_var.get().strip()
|
||
|
|
|
||
|
|
chap_num = self.safe_get_int(self.chapter_num_var, 1)
|
||
|
|
word_number = self.safe_get_int(self.word_number_var, 3000)
|
||
|
|
|
||
|
|
self.safe_log(f"开始定稿第{chap_num}章...")
|
||
|
|
|
||
|
|
chapters_dir = os.path.join(filepath, "chapters")
|
||
|
|
os.makedirs(chapters_dir, exist_ok=True)
|
||
|
|
chapter_file = os.path.join(chapters_dir, f"chapter_{chap_num}.txt")
|
||
|
|
|
||
|
|
edited_text = self.chapter_result.get("0.0", "end").strip()
|
||
|
|
|
||
|
|
if len(edited_text) < 0.7 * word_number:
|
||
|
|
ask = messagebox.askyesno("字数不足", f"当前章节字数 ({len(edited_text)}) 低于目标字数({word_number})的70%，是否要尝试扩写？")
|
||
|
|
if ask:
|
||
|
|
self.safe_log("正在扩写章节内容...")
|
||
|
|
enriched = enrich_chapter_text(
|
||
|
|
chapter_text=edited_text,
|
||
|
|
word_number=word_number,
|
||
|
|
api_key=api_key,
|
||
|
|
base_url=base_url,
|
||
|
|
model_name=model_name,
|
||
|
|
temperature=temperature,
|
||
|
|
interface_format=interface_format,
|
||
|
|
max_tokens=max_tokens,
|
||
|
|
timeout=timeout_val
|
||
|
|
)
|
||
|
|
edited_text = enriched
|
||
|
|
self.master.after(0, lambda: self.chapter_result.delete("0.0", "end"))
|
||
|
|
self.master.after(0, lambda: self.chapter_result.insert("0.0", edited_text))
|
||
|
|
clear_file_content(chapter_file)
|
||
|
|
save_string_to_txt(edited_text, chapter_file)
|
||
|
|
|
||
|
|
finalize_chapter(
|
||
|
|
novel_number=chap_num,
|
||
|
|
word_number=word_number,
|
||
|
|
api_key=api_key,
|
||
|
|
base_url=base_url,
|
||
|
|
model_name=model_name,
|
||
|
|
temperature=temperature,
|
||
|
|
filepath=filepath,
|
||
|
|
embedding_api_key=embedding_api_key,
|
||
|
|
embedding_url=embedding_url,
|
||
|
|
embedding_interface_format=embedding_interface_format,
|
||
|
|
embedding_model_name=embedding_model_name,
|
||
|
|
interface_format=interface_format,
|
||
|
|
max_tokens=max_tokens,
|
||
|
|
timeout=timeout_val
|
||
|
|
)
|
||
|
|
self.safe_log(f"✅ 第{chap_num}章定稿完成（已更新全局摘要、角色状态、向量库）。")
|
||
|
|
|
||
|
|
final_text = read_file(chapter_file)
|
||
|
|
self.master.after(0, lambda: self.show_chapter_in_textbox(final_text))
|
||
|
|
except Exception:
|
||
|
|
self.handle_exception("定稿章节时出错")
|
||
|
|
finally:
|
||
|
|
self.enable_button_safe(self.btn_finalize_chapter)
|
||
|
|
threading.Thread(target=task, daemon=True).start()
|
||
|
|
|
||
|
|
def do_consistency_check(self):
|
||
|
|
filepath = self.filepath_var.get().strip()
|
||
|
|
if not filepath:
|
||
|
|
messagebox.showwarning("警告", "请先配置保存文件路径。")
|
||
|
|
return
|
||
|
|
|
||
|
|
def task():
|
||
|
|
self.disable_button_safe(self.btn_check_consistency)
|
||
|
|
try:
|
||
|
|
api_key = self.api_key_var.get().strip()
|
||
|
|
base_url = self.base_url_var.get().strip()
|
||
|
|
model_name = self.model_name_var.get().strip()
|
||
|
|
temperature = self.temperature_var.get()
|
||
|
|
interface_format = self.interface_format_var.get()
|
||
|
|
max_tokens = self.max_tokens_var.get()
|
||
|
|
timeout = self.timeout_var.get()
|
||
|
|
|
||
|
|
chap_num = self.safe_get_int(self.chapter_num_var, 1)
|
||
|
|
chap_file = os.path.join(filepath, "chapters", f"chapter_{chap_num}.txt")
|
||
|
|
chapter_text = read_file(chap_file)
|
||
|
|
|
||
|
|
if not chapter_text.strip():
|
||
|
|
self.safe_log("⚠️ 当前章节文件为空或不存在，无法审校。")
|
||
|
|
return
|
||
|
|
|
||
|
|
self.safe_log("开始一致性审校...")
|
||
|
|
result = check_consistency(
|
||
|
|
novel_setting="",
|
||
|
|
character_state=read_file(os.path.join(filepath, "character_state.txt")),
|
||
|
|
global_summary=read_file(os.path.join(filepath, "global_summary.txt")),
|
||
|
|
chapter_text=chapter_text,
|
||
|
|
api_key=api_key,
|
||
|
|
base_url=base_url,
|
||
|
|
model_name=model_name,
|
||
|
|
temperature=temperature,
|
||
|
|
interface_format=interface_format,
|
||
|
|
max_tokens=max_tokens,
|
||
|
|
timeout=timeout,
|
||
|
|
plot_arcs=""
|
||
|
|
)
|
||
|
|
self.safe_log("审校结果：")
|
||
|
|
self.safe_log(result)
|
||
|
|
except Exception:
|
||
|
|
self.handle_exception("审校时出错")
|
||
|
|
finally:
|
||
|
|
self.enable_button_safe(self.btn_check_consistency)
|
||
|
|
threading.Thread(target=task, daemon=True).start()
|
||
|
|
|
||
|
|
def import_knowledge_handler(self):
|
||
|
|
selected_file = tk.filedialog.askopenfilename(
|
||
|
|
title="选择要导入的知识库文件",
|
||
|
|
filetypes=[("Text Files", "*.txt"), ("All Files", "*.*")]
|
||
|
|
)
|
||
|
|
if selected_file:
|
||
|
|
def task():
|
||
|
|
self.disable_button_safe(self.btn_import_knowledge)
|
||
|
|
try:
|
||
|
|
emb_api_key = self.embedding_api_key_var.get().strip()
|
||
|
|
emb_url = self.embedding_url_var.get().strip()
|
||
|
|
emb_format = self.embedding_interface_format_var.get().strip()
|
||
|
|
emb_model = self.embedding_model_name_var.get().strip()
|
||
|
|
|
||
|
|
self.safe_log(f"开始导入知识库文件: {selected_file}")
|
||
|
|
import_knowledge_file(
|
||
|
|
embedding_api_key=emb_api_key,
|
||
|
|
embedding_url=emb_url,
|
||
|
|
embedding_interface_format=emb_format,
|
||
|
|
embedding_model_name=emb_model,
|
||
|
|
file_path=selected_file,
|
||
|
|
filepath=self.filepath_var.get().strip()
|
||
|
|
)
|
||
|
|
self.safe_log("✅ 知识库文件导入完成。")
|
||
|
|
except Exception:
|
||
|
|
self.handle_exception("导入知识库时出错")
|
||
|
|
finally:
|
||
|
|
self.enable_button_safe(self.btn_import_knowledge)
|
||
|
|
threading.Thread(target=task, daemon=True).start()
|
||
|
|
|
||
|
|
def clear_vectorstore_handler(self):
|
||
|
|
filepath = self.filepath_var.get().strip()
|
||
|
|
if not filepath:
|
||
|
|
messagebox.showwarning("警告", "请先配置保存文件路径。")
|
||
|
|
return
|
||
|
|
|
||
|
|
first_confirm = messagebox.askyesno("警告", "确定要清空本地向量库吗？此操作不可恢复！")
|
||
|
|
if first_confirm:
|
||
|
|
second_confirm = messagebox.askyesno("二次确认", "你确定真的要删除所有向量数据吗？此操作不可恢复！")
|
||
|
|
if second_confirm:
|
||
|
|
if clear_vector_store(filepath):
|
||
|
|
self.log("已清空向量库。")
|
||
|
|
else:
|
||
|
|
self.log(f"未能清空向量库，请关闭程序后手动删除 {filepath} 下的 vectorstore 文件夹。")
|
||
|
|
|
||
|
|
def show_plot_arcs_ui(self):
|
||
|
|
filepath = self.filepath_var.get().strip()
|
||
|
|
if not filepath:
|
||
|
|
messagebox.showwarning("警告", "请先在主Tab中设置保存文件路径")
|
||
|
|
return
|
||
|
|
|
||
|
|
plot_arcs_file = os.path.join(filepath, "plot_arcs.txt")
|
||
|
|
if not os.path.exists(plot_arcs_file):
|
||
|
|
messagebox.showinfo("剧情要点", "当前还未生成任何剧情要点或冲突记录。")
|
||
|
|
return
|
||
|
|
|
||
|
|
arcs_text = read_file(plot_arcs_file).strip()
|
||
|
|
if not arcs_text:
|
||
|
|
arcs_text = "当前没有记录的剧情要点或冲突。"
|
||
|
|
|
||
|
|
top = ctk.CTkToplevel(self.master)
|
||
|
|
top.title("剧情要点/未解决冲突")
|
||
|
|
top.geometry("600x400")
|
||
|
|
text_area = ctk.CTkTextbox(top, wrap="word", font=("Microsoft YaHei", 12))
|
||
|
|
text_area.pack(fill="both", expand=True, padx=10, pady=10)
|
||
|
|
text_area.insert("0.0", arcs_text)
|
||
|
|
text_area.configure(state="disabled")
|