app.py

import json
import time

from flask import Flask, request, Response, jsonify
from flask_cors import CORS
app = Flask(__name__)
CORS(app)

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

from prompts.baseprompt import clean_txt_content, load_prompt

from core.writer_utils import KeyPointMsg
from core.draft_writer import DraftWriter
from core.plot_writer import PlotWriter
from core.outline_writer import OutlineWriter

from setting import setting_bp
from summary import process_novel
from backend_utils import get_model_config_from_provider_model
from config import MAX_NOVEL_SUMMARY_LENGTH, MAX_THREAD_NUM, ENABLE_ONLINE_DEMO


app.register_blueprint(setting_bp)

# 添加配置
BACKEND_HOST = os.environ.get('BACKEND_HOST', '0.0.0.0')
BACKEND_PORT = int(os.environ.get('BACKEND_PORT', 7869))


@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({
        'status': 'healthy',
        'timestamp': int(time.time())
    }), 200


def load_novel_writer(writer_mode, chunk_list, global_context, x_chunk_length, y_chunk_length, main_model, sub_model, max_thread_num) -> DraftWriter:
    kwargs = dict(
        xy_pairs=chunk_list,
        model=get_model_config_from_provider_model(main_model),
        sub_model=get_model_config_from_provider_model(sub_model),
    )

    kwargs['x_chunk_length'] = x_chunk_length
    kwargs['y_chunk_length'] = y_chunk_length
    kwargs['max_thread_num'] = max_thread_num
    match writer_mode:
        case 'draft':
            kwargs['global_context'] = {}
            novel_writer = DraftWriter(**kwargs)
        case 'outline':
            kwargs['global_context'] = {'summary': global_context}
            novel_writer = OutlineWriter(**kwargs)
        case 'plot':
            kwargs['global_context'] = {'chapter': global_context}
            novel_writer = PlotWriter(**kwargs)
        case _:
            raise ValueError(f"unknown writer: {writer_mode}")
            
    return novel_writer





prompt_names = dict(
    outline = ['新建章节', '扩写章节', '润色章节'],
    plot = ['新建剧情', '扩写剧情', '润色剧情'],
    draft = ['新建正文', '扩写正文', '润色正文'],
)

prompt_dirname = dict(
    outline = 'prompts/创作章节',
    plot = 'prompts/创作剧情',
    draft = 'prompts/创作正文',
)


PROMPTS = {}
for type_name, dirname in prompt_dirname.items():
    PROMPTS[type_name] = {'prompt_names': prompt_names[type_name]}
    for name in prompt_names[type_name]:
        content = clean_txt_content(load_prompt(dirname, name))
        if content.startswith("user:\n"):
            content = content[len("user:\n"):]
        PROMPTS[type_name][name] = {'content': content}


@app.route('/prompts', methods=['GET'])
def get_prompts():
    return jsonify(PROMPTS)

def get_delta_chunks(prev_chunks, curr_chunks):
    """Calculate delta between previous and current chunks"""
    if not prev_chunks or len(prev_chunks) != len(curr_chunks):
        return "init", curr_chunks
    
    # Check if all strings in current chunks start with their corresponding previous strings
    is_delta = True
    for prev_chunk, curr_chunk in zip(prev_chunks, curr_chunks):
        if len(prev_chunk) != len(curr_chunk):
            is_delta = False
            break
        for prev_str, curr_str in zip(prev_chunk, curr_chunk):
            if not curr_str.startswith(prev_str):
                is_delta = False
                break
        if not is_delta:
            break
    
    if not is_delta:
        return "init", curr_chunks
    
    # Calculate deltas
    delta_chunks = []
    for prev_chunk, curr_chunk in zip(prev_chunks, curr_chunks):
        delta_chunk = []
        for prev_str, curr_str in zip(prev_chunk, curr_chunk):
            delta_str = curr_str[len(prev_str):]
            delta_chunk.append(delta_str)
        delta_chunks.append(delta_chunk)
    
    return "delta", delta_chunks


def call_write(writer_mode, chunk_list, global_context, chunk_span, prompt_content, x_chunk_length, y_chunk_length, main_model, sub_model, max_thread_num, only_prompt):
    if ENABLE_ONLINE_DEMO:
        if max_thread_num > MAX_THREAD_NUM:
            raise Exception("在线Demo模型下，最大线程数不能超过" + str(MAX_THREAD_NUM) + "！")
    
    # 输入的chunk_list中每个chunk需要加上换行，除了最后一个chunk（因为是从页面中各个chunk传来的）
    chunk_list = [[e.strip() + ('\n' if e.strip() and rowi != len(chunk_list)-1 else '') for e in row] for rowi, row in enumerate(chunk_list)]

    prev_chunks = None
    def delta_wrapper(chunk_list, done=False, msg=None):
        # 返回的chunk_list中每个chunk需要去掉换行
        chunk_list = [[e.strip() for e in row] for row in chunk_list]

        nonlocal prev_chunks
        if prev_chunks is None:
            prev_chunks = chunk_list
            return {
                "done": done,
                "chunk_type": "init",
                "chunk_list": chunk_list,
                "msg": msg
            }
        else:
            chunk_type, new_chunks = get_delta_chunks(prev_chunks, chunk_list)
            prev_chunks = chunk_list
            return {
                "done": done,
                "chunk_type": chunk_type,
                "chunk_list": new_chunks,
                "msg": msg
            }
        
    novel_writer = load_novel_writer(writer_mode, chunk_list, global_context, x_chunk_length, y_chunk_length, main_model, sub_model, max_thread_num)
    

    # draft需要映射，所以进行初始划分
    if writer_mode == 'draft':
        target_chunk = novel_writer.get_chunk(pair_span=chunk_span)
        new_target_chunk = novel_writer.map_text_wo_llm(target_chunk)
        novel_writer.apply_chunks([target_chunk], [new_target_chunk])
        chunk_span = novel_writer.get_chunk_pair_span(new_target_chunk)
    
    init_novel_writer = load_novel_writer(writer_mode, list(novel_writer.xy_pairs), global_context, x_chunk_length, y_chunk_length, main_model, sub_model, max_thread_num)
    
    # TODO: writer.write 应该保证无论什么prompt，都能够同时适应y为空和y有值地情况
    # 换句话说，就是虽然可以单列出一个"新建正文"，但用扩写正文也能实现同样的效果。
    generator = novel_writer.write(prompt_content, pair_span=chunk_span) 
    
    prompt_outputs = []
    last_yield_time = time.time()  # Initialize the last yield time

    prompt_name = ''
    for kp_msg in generator:
        if isinstance(kp_msg, KeyPointMsg):
            # 如果要支持关键节点保存，需要计算一个编辑上的更改，然后在这里yield writer
            prompt_name = kp_msg.prompt_name
            continue
        else:
            chunk_list = kp_msg

        current_cost = 0
        currency_symbol = ''
        current_model = ''
        data_chunks = []
        prompt_outputs.clear()
        for e in chunk_list:
            if e is None: continue  # e为None说明该chunk还未处理
            output, chunk = e
            if output is None: continue # output为None说明该chunk未yield就return，说明未调用llm
            prompt_outputs.append(output)
            current_text = ""
            current_model = output['response_msgs'].model
            current_cost += output['response_msgs'].cost
            currency_symbol = output['response_msgs'].currency_symbol
            if 'plot2text' in output:
                current_text += f"正在建立映射关系..." + '\n'
            else:
                current_text = output['text']
            data_chunks.append((chunk.x_chunk, chunk.y_chunk, current_text))
            
        if only_prompt:
            yield {'prompts': [e['response_msgs'] for e in prompt_outputs]}
            return

        current_time = time.time()
        if current_time - last_yield_time >= 0.2:  # Check if 0.2 seconds have passed
            yield delta_wrapper(data_chunks, done=False, msg=f"正在 {prompt_name} （{len(prompt_outputs)} / {len(chunk_list)}）" + f" 模型：{current_model} 花费：{current_cost:.5f}{currency_symbol}" if current_model else '')
            last_yield_time = current_time  # Update the last yield time

    # 这里是计算出一个编辑上的更改，方便前端显示，后续diff功能将不由writer提供，因为这是为了显示的要求
    data_chunks = init_novel_writer.diff_to(novel_writer, pair_span=chunk_span)

    yield delta_wrapper(data_chunks, done=True, msg='创作完成!')


@app.route('/write', methods=['POST'])
def write():
    data = request.json                 
    writer_mode = data['writer_mode']
    chunk_list = data['chunk_list']
    chunk_span = data['chunk_span']
    prompt_content = data['prompt_content']
    x_chunk_length = data['x_chunk_length']
    y_chunk_length = data['y_chunk_length']
    main_model = data['main_model']
    sub_model = data['sub_model']
    global_context = data['global_context']
    only_prompt = data['only_prompt']
    
    # Update settings if provided
    if 'settings' in data:
        max_thread_num = data['settings']['MAX_THREAD_NUM']

    # Generate unique stream ID
    stream_id = str(time.time())
    active_streams[stream_id] = True
    
    def generate():
        try:
            # Send stream ID to client
            yield f"data: {json.dumps({'stream_id': stream_id})}\n\n"

            for result in call_write(writer_mode, list(chunk_list), global_context, chunk_span, prompt_content, x_chunk_length, y_chunk_length, main_model, sub_model, max_thread_num, only_prompt):
                if not active_streams.get(stream_id, False):
                    # Stream was stopped by client
                    print(f"Stream was stopped by client: {stream_id}")
                    return
                    
                yield f"data: {json.dumps(result)}\n\n"
        except Exception as e:
            error_msg = f"创作出错：\n{str(e)}"
            error_chunk_list = [[*e[:2], error_msg] for e in chunk_list[chunk_span[0]:chunk_span[1]]]
            
            error_data = {
                "done": True,
                "chunk_type": "init",
                "chunk_list": error_chunk_list
            }
            yield f"data: {json.dumps(error_data)}\n\n"
        finally:
            # Clean up stream tracking
            if stream_id in active_streams:
                del active_streams[stream_id]

    return Response(generate(), mimetype='text/event-stream')


@app.route('/summary', methods=['POST'])
def process_novel_text():
    data = request.json
    content = data['content']
    novel_name = data['novel_name']

    # Generate unique stream ID
    stream_id = str(time.time())
    active_streams[stream_id] = True

    def generate():
        try:
            yield f"data: {json.dumps({'stream_id': stream_id})}\n\n"

            main_model = get_model_config_from_provider_model(data['main_model'])
            sub_model = get_model_config_from_provider_model(data['sub_model'])
            max_novel_summary_length = data['settings']['MAX_NOVEL_SUMMARY_LENGTH']
            max_thread_num = data['settings']['MAX_THREAD_NUM']
            last_yield_time = 0
            for result in process_novel(content, novel_name, main_model, sub_model, max_novel_summary_length, max_thread_num):
                if not active_streams.get(stream_id, False):
                    # Stream was stopped by client
                    print(f"Stream was stopped by client: {stream_id}")
                    return
                    
                current_time = time.time()
                yield_value = f"data: {json.dumps(result)}\n\n"
                if current_time - last_yield_time >= 0.2:
                    last_yield_time = current_time
                    yield yield_value
            if current_time - last_yield_time < 0.2:
                # Save last yield to yaml file
                import yaml
                result_dict = json.loads(yield_value.replace('data: ', '').strip())
                with open('tmp.yaml', 'w', encoding='utf-8') as f:
                    yaml.dump(result_dict, f, allow_unicode=True)
                    
                yield yield_value   # Ensure last yield is returned
            
        except Exception as e:
            error_data = {
                "progress_msg": f"处理出错：{str(e)}",
            }
            yield f"data: {json.dumps(error_data)}\n\n"
        finally:
            # Clean up stream tracking
            if stream_id in active_streams:
                del active_streams[stream_id]

    return Response(generate(), mimetype='text/event-stream')

# Dictionary to track active streams
active_streams = {}

@app.route('/stop_stream', methods=['POST'])
def stop_stream():
    data = request.json
    stream_id = data.get('stream_id')
    if stream_id in active_streams:
        active_streams[stream_id] = False
    return jsonify({'success': True})

if __name__ == '__main__':
    app.run(host=BACKEND_HOST, port=BACKEND_PORT, debug=False)