120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
#novel_generator/finalization.py
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
定稿章节和扩写章节（finalize_chapter、enrich_chapter_text）
|
||
"""
|
||
import os
|
||
import logging
|
||
from llm_adapters import create_llm_adapter
|
||
from embedding_adapters import create_embedding_adapter
|
||
from prompt_definitions import summary_prompt, update_character_state_prompt
|
||
from novel_generator.common import invoke_with_cleaning
|
||
from utils import read_file, clear_file_content, save_string_to_txt
|
||
from novel_generator.vectorstore_utils import update_vector_store
|
||
|
||
def finalize_chapter(
|
||
novel_number: int,
|
||
word_number: int,
|
||
api_key: str,
|
||
base_url: str,
|
||
model_name: str,
|
||
temperature: float,
|
||
filepath: str,
|
||
embedding_api_key: str,
|
||
embedding_url: str,
|
||
embedding_interface_format: str,
|
||
embedding_model_name: str,
|
||
interface_format: str,
|
||
max_tokens: int,
|
||
timeout: int = 600
|
||
):
|
||
"""
|
||
对指定章节做最终处理：更新全局摘要、更新角色状态、插入向量库等。
|
||
默认无需再做扩写操作，若有需要可在外部调用 enrich_chapter_text 处理后再定稿。
|
||
"""
|
||
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()
|
||
if not chapter_text:
|
||
logging.warning(f"Chapter {novel_number} is empty, cannot finalize.")
|
||
return
|
||
|
||
global_summary_file = os.path.join(filepath, "global_summary.txt")
|
||
old_global_summary = read_file(global_summary_file)
|
||
character_state_file = os.path.join(filepath, "character_state.txt")
|
||
old_character_state = read_file(character_state_file)
|
||
|
||
llm_adapter = create_llm_adapter(
|
||
interface_format=interface_format,
|
||
base_url=base_url,
|
||
model_name=model_name,
|
||
api_key=api_key,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
timeout=timeout
|
||
)
|
||
|
||
prompt_summary = summary_prompt.format(
|
||
chapter_text=chapter_text,
|
||
global_summary=old_global_summary
|
||
)
|
||
new_global_summary = invoke_with_cleaning(llm_adapter, prompt_summary)
|
||
if not new_global_summary.strip():
|
||
new_global_summary = old_global_summary
|
||
|
||
prompt_char_state = update_character_state_prompt.format(
|
||
chapter_text=chapter_text,
|
||
old_state=old_character_state
|
||
)
|
||
new_char_state = invoke_with_cleaning(llm_adapter, prompt_char_state)
|
||
if not new_char_state.strip():
|
||
new_char_state = old_character_state
|
||
|
||
clear_file_content(global_summary_file)
|
||
save_string_to_txt(new_global_summary, global_summary_file)
|
||
clear_file_content(character_state_file)
|
||
save_string_to_txt(new_char_state, character_state_file)
|
||
|
||
update_vector_store(
|
||
embedding_adapter=create_embedding_adapter(
|
||
embedding_interface_format,
|
||
embedding_api_key,
|
||
embedding_url,
|
||
embedding_model_name
|
||
),
|
||
new_chapter=chapter_text,
|
||
filepath=filepath
|
||
)
|
||
|
||
logging.info(f"Chapter {novel_number} has been finalized.")
|
||
|
||
def enrich_chapter_text(
|
||
chapter_text: str,
|
||
word_number: int,
|
||
api_key: str,
|
||
base_url: str,
|
||
model_name: str,
|
||
temperature: float,
|
||
interface_format: str,
|
||
max_tokens: int,
|
||
timeout: int=600
|
||
) -> str:
|
||
"""
|
||
对章节文本进行扩写，使其更接近 word_number 字数，保持剧情连贯。
|
||
"""
|
||
llm_adapter = create_llm_adapter(
|
||
interface_format=interface_format,
|
||
base_url=base_url,
|
||
model_name=model_name,
|
||
api_key=api_key,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
timeout=timeout
|
||
)
|
||
prompt = f"""以下章节文本较短，请在保持剧情连贯的前提下进行扩写，使其更充实，接近 {word_number} 字左右：
|
||
原内容：
|
||
{chapter_text}
|
||
"""
|
||
enriched_text = invoke_with_cleaning(llm_adapter, prompt)
|
||
return enriched_text if enriched_text else chapter_text
|