core/frontend.py

import re
from rich.traceback import install
install(show_locals=False)

import gradio as gr
import yaml
import functools
import time
import sys
import os
import copy

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config import RENDER_SAVE_LOAD_BTN, RENDER_STOP_BTN
from core.backend import call_write, call_accept, match_quote_text
from core.frontend_copy import enable_copy_js, on_copy
from core.frontend_setting import new_setting, render_setting
from core.frontend_utils import (
    title, info,
    create_progress_md, create_text_md, messages2chatbot,
    init_writer, has_accept, is_running, try_cancel, writer_y_is_empty, writer_x_is_empty,
    cancellable, process_writer_to_backend, process_writer_from_backend,
    init_chapters_w, init_draft_w
)
from core.writer_utils import KeyPointMsg

from prompts.baseprompt import clean_txt_content, load_prompt


# 读取YAML文件
with open('prompts/idea-examples.yaml', 'r', encoding='utf-8') as file:
    examples_data = yaml.safe_load(file)

# 准备示例列表
examples = [[example['idea']] for example in examples_data['examples']]

with gr.Blocks(head=enable_copy_js) as demo:
    gr.HTML(title)
    with gr.Accordion("使用指南"):
        gr.Markdown(info)

    writer_state = gr.State(init_writer('', check_empty=False))
    setting_state = gr.State(new_setting())

    if RENDER_SAVE_LOAD_BTN:
        with gr.Row():
            save_button = gr.Button("保存状态")
            load_button = gr.Button("加载状态")
            save_file_name = gr.Textbox(value='states', placeholder='输入文件名', lines=1, label=None, show_label=False, container=False)

    def save_states(save_file_name, writer, setting):
        import json
        json_file_name = save_file_name + '.json'
        with open(json_file_name, 'w', encoding='utf-8') as f:
            json.dump({
                'writer': writer,
                'setting': setting
            }, f, ensure_ascii=False, indent=2)
        gr.Info(f"状态已保存到文件：{json_file_name}")

    def load_states(save_file_name):
        import json
        json_file_name = save_file_name + '.json'
        try:
            with open(json_file_name, 'r', encoding='utf-8') as f:
                states = json.load(f)
            gr.Info(f"状态文件已加载：{json_file_name}")
            states['setting']['render_time'] = time.time()
            # 为了确保setting被渲染，选择模型是不会赋值setting_state的
            # 需要保证setting界面持有的对象和setting_state是同一个
            return states['writer'], states['setting']
        except FileNotFoundError:
            raise gr.Error(f"未找到保存的状态文件：{json_file_name}")

    idea_textbox = gr.Textbox(placeholder='用一段话描述你要写的小说，或者从下方示例中选择一个创意...', lines=2, scale=1, label=None, show_label=False, container=False, max_length=1000)
    
    gr.Examples(
        label='示例',
        examples=examples,
        inputs=[idea_textbox],
    )

    with gr.Row():    
        outline_btn = gr.Button("创作大纲", scale=1, min_width=1, interactive = True, variant='primary')
        chapters_btn = gr.Button("创作剧情", scale=1, min_width=1, interactive = False, variant='secondary')
        draft_btn = gr.Button("创作正文", scale=1, min_width=1, interactive = False, variant='secondary')
        auto_checkbox = gr.Checkbox(label='一键生成', scale=1, value=False, visible=False)  # TODO: V1.10版本 “自动”尚不完善，暂不显示

    progress_md = create_progress_md(writer_state.value)
    text_md = create_text_md(writer_state.value)

    @gr.render(inputs=writer_state)
    def create_prompt_preview(writer):
        prompt_outputs = writer['prompt_outputs'] if 'prompt_outputs' in writer else []
        with gr.Accordion("Prompt预览", open=bool(prompt_outputs)):
            pause_on_prompt_finished_checkbox = gr.Checkbox(label='允许在LLM响应完成后，预览Prompt', scale=1, value=writer['pause_on_prompt_finished_flag']) 
    
            for i, prompt_output in enumerate(prompt_outputs, 1):
                with gr.Tab(f"Prompt {i}"):
                    gr.Chatbot(messages2chatbot(prompt_output['response_msgs']), type='messages')
            if not prompt_outputs:
                gr.Markdown('当前没有可预览的Prompt。')

            continue_btn = gr.Button('继续', visible=bool(prompt_outputs), variant='primary')

        def on_pause_on_prompt_finished(value):
            if value:
                gr.Info("在LLM响应完成后，将可以预览Prompt")
            writer['pause_on_prompt_finished_flag'] = value
        
        pause_on_prompt_finished_checkbox.change(on_pause_on_prompt_finished, [pause_on_prompt_finished_checkbox])

        def on_continue(writer):
            writer['pause_flag'] = False
            writer['prompt_outputs'] = []
            return writer
        
        continue_btn.click(on_continue, writer_state, writer_state)


    with gr.Row():
        rewrite_all_button = gr.Button("开始创作", min_width=100, scale=2, variant='secondary', interactive=False)
        suggestion_dropdown = gr.Dropdown(choices=[], min_width=100, scale=2, label=None, show_label=False, container=False, allow_custom_value=False)
        quote_checkbox = gr.Checkbox(label='允许引用', min_width=100, scale=2, value=False)
        gr.Textbox('窗口大小：', container=False, text_align='right', scale=1, min_width=100)
        chunk_length_dropdown = gr.Dropdown(choices=[], min_width=80, scale=1, label=None, show_label=False, container=False, allow_custom_value=False)

    quote_md = gr.Markdown(visible=False)
    
    def on_quote_checkbox_change(writer, value):
        if writer['current_w'] == 'outline_w':
            gr.Info("大纲创作不支持引用\n考虑在剧情和正文创作中使用吧~")
            return gr.update(value=False, visible=False)
        
        if value:
            gr.Info("允许引用（右键或Ctrl+C复制你想引用的文本）")
        writer['quote_span'] = None
        writer['quoted_text'] = ''
        return gr.update(value=None, visible=False)

    quote_checkbox.change(on_quote_checkbox_change, [writer_state, quote_checkbox], [quote_md])

    def on_chunk_length_change(writer, value):
        current_w_name = writer['current_w']
        writer[current_w_name]['y_chunk_length'] = value
        return gr.update(value=value)

    chunk_length_dropdown.change(on_chunk_length_change, [writer_state, chunk_length_dropdown], [chunk_length_dropdown])

    def on_copy_handle(text, writer, setting, quote_checkbox):
        # gr.Info(f"Copy: {text}")
        text = text.strip()

        if has_accept(writer):
            gr.Info('考虑先接受或拒绝修改哦~')
            return gr.update(visible=False)

        if len(text) < 10:
            gr.Info('选中的文本太短，无法引用')
            return gr.update(visible=False)
        
        if quote_checkbox:
            quote_span, quoted_text = match_quote_text(writer, setting, text)
            if quote_span:
                writer['quote_span'] = quote_span
                writer['quoted_text'] = quoted_text
                lines = quoted_text.split('\n')
                if len(lines) > 10:
                    lines[5:-5] = ['......']
                lines = ['```', ] + lines + ['```', ]
                quoted_text = '\n'.join(["> " + e for e in lines])
                return gr.update(value=quoted_text, visible=True)
            else:
                gr.Info('未找到匹配的引用文本')

        writer['quote_span'] = None
        writer['quoted_text'] = ''
        return gr.update(visible=False)

    on_copy(on_copy_handle, [writer_state, setting_state, quote_checkbox], [quote_md])
    

    suggestion_textbox = gr.Textbox(max_length=1000, placeholder='在这里输入你的意见，或者从右上单选框选择', lines=2, scale=1, label=None, show_label=False, container=False)

    with gr.Row():    
        accept_button = gr.Button("接受", scale=1, min_width=1, variant='secondary', interactive=False)
        pause_button = gr.Button("暂停", scale=1, min_width=1, variant='secondary', visible=RENDER_STOP_BTN)
        stop_button = gr.Button("取消", scale=1, min_width=1, variant='secondary')
        flash_button = gr.Button("刷新", scale=1, min_width=1, variant='secondary')

    def flash_interface(writer):
        current_w_name = writer['current_w']

        can_accept_flag = has_accept(writer) and not is_running(writer)
        can_write_flag = not writer_x_is_empty(writer, current_w_name) and not can_accept_flag

        match current_w_name:
            case 'outline_w':
                rewrite_all_button = gr.update(value='开始创作', variant='primary' if can_write_flag else 'secondary', interactive=can_write_flag)
            case 'chapters_w':
                rewrite_all_button = gr.update(value='开始创作', variant='primary' if can_write_flag else 'secondary', interactive=can_write_flag)
            case 'draft_w':
                rewrite_all_button = gr.update(value='开始创作', variant='primary' if can_write_flag else 'secondary', interactive=can_write_flag)

        accept_button = gr.update(variant='primary' if can_accept_flag else 'secondary', interactive=can_accept_flag)
        
        # 更新 chapters_btn 和 draft_btn 的 interactive 状态
        outline_btn = gr.update(
            variant='primary' if current_w_name == 'outline_w' else 'secondary'
            )
        chapters_btn = gr.update(
            interactive=not writer_y_is_empty(writer, 'outline_w'),
            variant='primary' if current_w_name == 'chapters_w' else 'secondary'
        )
        draft_btn = gr.update(
            interactive=not writer_y_is_empty(writer, 'chapters_w'),
            variant='primary' if current_w_name == 'draft_w' else 'secondary'
        )

        pause_button = gr.update(
            value="继续" if writer['pause_flag'] else "暂停",
            variant='secondary',
        )

        suggestion_choices = writer['suggestions'][current_w_name]
        # suggestion_choices = ['自动', ] + writer['suggestions'][current_w_name]  # TODO: V1.10版本 “自动”尚不完善，暂不显示
        if writer_y_is_empty(writer, current_w_name):
            suggestion_dropdown = gr.update(choices=suggestion_choices, value=suggestion_choices[0])
        else:
            suggestion_dropdown = gr.update(choices=suggestion_choices,)

        chunk_length_choices = writer['chunk_length'][current_w_name]
        if cur_chunk_length := writer[current_w_name].get('y_chunk_length', None):
            chunk_length_dropdown = gr.update(choices=chunk_length_choices, value=cur_chunk_length)
        else:
            chunk_length_dropdown = gr.update(choices=chunk_length_choices, value=chunk_length_choices[0])

        return (
            create_text_md(writer),
            create_progress_md(writer),
            rewrite_all_button,
            accept_button,
            outline_btn,
            chapters_btn,
            draft_btn,
            pause_button,
            suggestion_dropdown,
            chunk_length_dropdown
        )

    # 更新 flash_event 字典以包含新的输出
    flash_event = dict(
        fn=flash_interface, 
        inputs=[writer_state], 
        outputs=[
            text_md,
            progress_md,
            rewrite_all_button,
            accept_button,
            outline_btn,
            chapters_btn,
            draft_btn,
            pause_button,
            suggestion_dropdown,
            chunk_length_dropdown
        ]
    )
    
    flash_button.click(**flash_event)
    if RENDER_SAVE_LOAD_BTN:
        save_button.click(save_states, inputs=[save_file_name, writer_state, setting_state], outputs=[])
        load_button.click(load_states, inputs=[save_file_name], outputs=[writer_state, setting_state]).success(**flash_event)
    # stop_write_long_novel_button.click(on_cancel, inputs=[writer_state])
    stop_button.click(try_cancel, inputs=[writer_state]).success(**flash_event).success(lambda :gr.update(), None, writer_state)
    # TODO: stop_btn对writer_state的更新没有起效

    @cancellable
    def _on_write_all(writer, setting, auto_write=False, suggestion=None):
        current_w_name = writer['current_w']
           
        if writer_x_is_empty(writer, current_w_name):
            gr.Info('请先输入需要创作的内容！')
            return
        
        writer['prompt_outputs'].clear()

        if writer['quote_span']:
            quote_span, quoted_text = match_quote_text(writer, setting, writer['quoted_text'])
            if quote_span != writer['quote_span'] or quoted_text != writer['quoted_text']:
                raise gr.Error('引用文本不存在！')

        generator = call_write(process_writer_to_backend(writer), setting, auto_write, suggestion)

        new_writer = None
        while True:
            try:
                kp_msg = next(generator)
                if isinstance(kp_msg, KeyPointMsg):
                    # TODO: 由于KeyPointMsg的设计问题，这里的逻辑比较复杂，后续可以考虑优化
                    if kp_msg.is_prompt() and kp_msg.is_finished() and writer['pause_on_prompt_finished_flag']:
                        gr.Info('LLM响应完成，可以预览Prompt')  
                        writer['pause_flag'] = True
                        if new_writer is None: continue
                    elif kp_msg.is_title(): # TODO: 标题节点还未实现finish逻辑
                        # if new_writer is not None:
                        #     # 说明这是一个关键节点，进行保存
                        #     process_writer_from_backend(writer, new_writer)
                        #     yield create_text_md(writer), writer
                        #     gr.Info(f'已自动保存进度')
                        continue
                        # 关键节点保存的逻辑比较复杂，有bug，之后版本考虑提供
                    else:
                        continue
                else:
                    new_writer = kp_msg
                
                if writer['pause_flag']:
                    writer['prompt_outputs'] = copy.deepcopy(new_writer['prompt_outputs'])  
                    # 将prompt_outputs传递到writer_state中，使得暂停时能显示prompt, 需要序列化，否则writer会不断更新，导致prompt不断渲染
                    yield create_text_md(new_writer), writer

                    while writer['pause_flag'] and not writer['cancel_flag']:
                        time.sleep(0.1)
                else:
                    yield create_text_md(new_writer), gr.update()
            except StopIteration as e:
                # 这里处理最终状态
                process_writer_from_backend(writer, e.value)
                yield create_text_md(writer), writer
                if has_accept(writer):
                    gr.Info('创作完成！点击接受按钮接受修改。')
                else:
                    gr.Info('本次创作没有任何更改。')  # 通常因为审阅意见认为无需更改
                return
        
    def on_auto_write_all(writer, setting, auto_write):
        if auto_write:
            yield from _on_write_all(writer, setting, True)
        else:
            pass
            # suggestion = writer['suggestions'][writer['current_w']][0]
            # yield from _on_write_all(writer, setting, False, suggestion)

    writer_all_events = dict(
            fn=on_auto_write_all,
            queue=True,
            inputs=[writer_state, setting_state, auto_checkbox],
            outputs=[text_md, writer_state],
            concurrency_limit=10
    )

    def on_init_outline(idea, writer):
        if not idea.strip():
            gr.Info("先输入小说简介或从示例中选择一个")
            return gr.update()
        new_writer = init_writer(idea)
        writer.update({
            k:v for k, v in new_writer.items() if k in ['current_w', 'outline_w', 'prompt_outputs']
        })
        return writer
    
    outline_btn.click(on_init_outline, inputs=[idea_textbox, writer_state], outputs=[writer_state]).success(**writer_all_events).then(**flash_event)
    chapters_btn.click(lambda writer: init_chapters_w(writer), inputs=[writer_state], outputs=[writer_state]).success(**writer_all_events).then(**flash_event)
    draft_btn.click(lambda writer: init_draft_w(writer), inputs=[writer_state], outputs=[writer_state]).success(**writer_all_events).then(**flash_event)

    def on_select_suggestion(writer, setting, choice):
        if choice == '自动':
            return gr.update(value=choice, visible=False)

        current_w_name = writer['current_w']
        dirname = writer['suggestions_dirname'][current_w_name]
        suggestion = clean_txt_content(load_prompt(dirname, choice))
        if suggestion.startswith("user:\n"):
            suggestion = suggestion[len("user:\n"):]
        
        return gr.update(value=suggestion, visible=True)
    
    suggestion_dropdown.change(on_select_suggestion, inputs=[writer_state, setting_state, suggestion_dropdown], outputs=[suggestion_textbox])

    def on_write_all(writer, setting, suggestion):
        if not suggestion.strip():
            gr.Info('需要输入创作意见！')
            return
        yield from _on_write_all(writer, setting, False, suggestion)
        
    rewrite_all_button.click(
            on_write_all,
            queue=True,
            inputs=[writer_state, setting_state, suggestion_textbox],
            outputs=[text_md, writer_state],
            concurrency_limit=10
        ).then(**flash_event)    


    @cancellable
    def on_accept_write(writer, setting):
        current_w_name = writer['current_w']
        current_w = writer[current_w_name]
        
        if not current_w['apply_chunks']:
            raise gr.Error('请先进行创作！')
        
        new_writer = call_accept(process_writer_to_backend(writer), setting)
        process_writer_from_backend(writer, new_writer)
        yield create_text_md(writer), writer
    
    accept_button.click(fn=on_accept_write, inputs=[writer_state, setting_state], outputs=[text_md, writer_state]).then(**flash_event)

    def toggle_pause(writer):
        if not is_running(writer):
            gr.Info('当前没有正在进行的操作')
            return gr.update()
        
        writer['pause_flag'] = not writer['pause_flag']
        # gr.Info('已' + ('暂停' if writer['pause_flag'] else '继续') + '操作')
        return gr.update(value="暂停" if not writer['pause_flag'] else "继续")

    pause_button.click(
        toggle_pause,
        inputs=[writer_state],
        outputs=[pause_button]
    )

    @gr.render(inputs=setting_state)
    def _render_setting(setting):
        return render_setting(setting, setting_state)


demo.queue()
demo.launch(server_name="0.0.0.0", server_port=7860)
#demo.launch()