core/writer.py

import re
import numpy as np
import bisect
from dataclasses import asdict, dataclass

from llm_api import ModelConfig
from prompts.对齐剧情和正文 import prompt as match_plot_and_text
from prompts.审阅.prompt import main as prompt_review
from core.writer_utils import split_text_into_chunks, detect_max_edit_span, run_yield_func
from core.writer_utils import KeyPointMsg
from core.diff_utils import get_chunk_changes


class Chunk(dict):
    def __init__(self, chunk_pairs: tuple[tuple[str, str, str]], source_slice: tuple[int, int], text_slice: tuple[int, int]):
        super().__init__()
        self['chunk_pairs'] = tuple(chunk_pairs)
        
        if isinstance(source_slice, slice):
            source_slice = (source_slice.start, source_slice.stop)
        self['source_slice'] = source_slice

        if isinstance(text_slice, slice):
            text_slice = (text_slice.start, text_slice.stop)
        assert text_slice[1] is None or text_slice[1] < 0, 'text_slice end must be None or negative'
        self['text_slice'] = text_slice

    def edit(self, x_chunk=None, y_chunk=None, text_pairs=None):
        if x_chunk is not None:
            text_pairs = [(x_chunk, self.y_chunk), ]
        elif y_chunk is not None:
            text_pairs = [(self.x_chunk, y_chunk), ]
        else:
            text_pairs = text_pairs

        chunk_pairs = list(self['chunk_pairs'])
        chunk_pairs[self.text_slice] = list(text_pairs)

        return Chunk(chunk_pairs=tuple(chunk_pairs), source_slice=self.source_slice, text_slice=self.text_slice)
    
    @property
    def source_slice(self) -> slice:
        return slice(*self['source_slice'])

    @property
    def chunk_pairs(self) -> tuple[tuple[str, str]]:
        return self['chunk_pairs']
    
    @property
    def text_slice(self) -> slice:
        return slice(*self['text_slice'])
    
    @property
    def text_source_slice(self) -> slice:
        source_start = self.source_slice.start + self.text_slice.start
        source_stop = self.source_slice.stop + (self.text_slice.stop or 0)
        return slice(source_start, source_stop)
    
    @property
    def text_pairs(self) -> tuple[tuple[str, str]]:
        return self.chunk_pairs[self.text_slice]
    
    @property
    def x_chunk(self) -> str:
        return ''.join(pair[0] for pair in self.text_pairs)
    
    @property
    def y_chunk(self) -> str:
        return ''.join(pair[1] for pair in self.text_pairs)
    
    @property
    def x_chunk_len(self) -> int:
        return sum(len(pair[0]) for pair in self.text_pairs)
    
    @property
    def y_chunk_len(self) -> int:
        return sum(len(pair[1]) for pair in self.text_pairs)
    
    @property
    def x_chunk_context(self) -> str:
        return ''.join(pair[0] for pair in self.chunk_pairs)
    
    @property
    def y_chunk_context(self) -> str:
        return ''.join(pair[1] for pair in self.chunk_pairs)
    
    @property
    def x_chunk_context_len(self) -> int:
        return sum(len(pair[0]) for pair in self.chunk_pairs)
    
    @property
    def y_chunk_context_len(self) -> int:
        return sum(len(pair[1]) for pair in self.chunk_pairs)
    
    
class Writer:
    def __init__(self, xy_pairs, global_context=None, model:ModelConfig=None, sub_model:ModelConfig=None, x_chunk_length=1000, y_chunk_length=1000, max_thread_num=5):
        self.xy_pairs = xy_pairs
        self.global_context = global_context or {}

        self.model = model
        self.sub_model = sub_model

        self.x_chunk_length = x_chunk_length
        self.y_chunk_length = y_chunk_length

        # x_chunk_length是指一次prompt调用时输入的x长度（由batch_map函数控制）, 此参数会影响到映射到y的扩写率（即：LLM的输出窗口长度/x_chunk_length）
        # 同时，x_chunk_length会影响到map的chunk大小，map的pair大小主要由x_chunk_length决定（具体来说，由update_map函数控制，为x_chunk_length//2)
        # y_chunk_length对pair大小的影响较少（因为映射是一对多）

        self.max_thread_num = max_thread_num    # 使得可以单独控制某个chunk变量的线程数，这在同时运行多个Writer变量时有用
    
    @property
    def x(self):    # TODO: 考虑x经常访问的情况
        return ''.join(pair[0] for pair in self.xy_pairs)

    @property
    def y(self):
        return ''.join(pair[1] for pair in self.xy_pairs)
    
    @property
    def x_len(self):
        return sum(len(pair[0]) for pair in self.xy_pairs)

    @property
    def y_len(self):
        return sum(len(pair[1]) for pair in self.xy_pairs)

    def get_model(self):
        return self.model

    def get_sub_model(self):
        return self.sub_model
    
    def count_span_length(self, span):
        pairs = self.xy_pairs[span[0]:span[1]]
        return sum(len(pair[0]) for pair in pairs), sum(len(pair[1]) for pair in pairs)

    def align_span(self, x_span=None, y_span=None):
        if x_span is None and y_span is None:
            raise ValueError("Either x_span or y_span must be provided")
        
        if x_span is not None and y_span is not None:
            raise ValueError("Only one of x_span or y_span should be provided")
        
        is_x = x_span is not None
        z_span = x_span if is_x else y_span
        cumsum_z = np.cumsum([0] + [len(pair[0 if is_x else 1]) for pair in self.xy_pairs]).tolist()
        
        l, r = z_span
        start_chunk = bisect.bisect_right(cumsum_z, l) - 1
        end_chunk = bisect.bisect_left(cumsum_z, r)
        
        aligned_l = cumsum_z[start_chunk]
        aligned_r = cumsum_z[end_chunk]
        
        aligned_span = (aligned_l, aligned_r)
        pair_span = (start_chunk, end_chunk)
        
        # Add assertions to verify the correctness of the output
        assert aligned_l <= l < aligned_r, "aligned_span does not properly contain the start of the input span"
        assert aligned_l < r <= aligned_r, "aligned_span does not properly contain the end of the input span"
        assert 0 <= start_chunk < end_chunk <= len(self.xy_pairs), "pair_span is out of bounds"
        assert sum(len(pair[0 if is_x else 1]) for pair in self.xy_pairs[start_chunk:end_chunk]) == aligned_r - aligned_l, "aligned_span and pair_span do not match"

        return aligned_span, pair_span
    
    def get_chunk(self, pair_span=None, x_span=None, y_span=None, context_length=0, smooth=True):
        if sum(x is not None for x in [pair_span, x_span, y_span]) != 1:
            raise ValueError("Exactly one of pair_span, x_span, or y_span must be provided")
        
        assert pair_span is None or (pair_span[0] >= 0 and pair_span[1] <= len(self.xy_pairs)), "pair_span is out of bounds"

        is_x = x_span is not None
        is_pair = pair_span is not None

        if is_pair:
            context_pair_span = (
                max(0, pair_span[0] - context_length),
                min(len(self.xy_pairs), pair_span[1] + context_length)
            )
        else:
            assert smooth, "smooth must be True"
            span = x_span if is_x else y_span
            if smooth:
                span, pair_span = self.align_span(x_span=span if is_x else None, y_span=span if not is_x else None)

            context_span = (
                max(0, span[0] - context_length),
                min(self.x_len if is_x else self.y_len, span[1] + context_length)
            )

            context_span, context_pair_span = self.align_span(x_span=context_span if is_x else None, y_span=context_span if not is_x else None)

        chunk_pairs = self.xy_pairs[context_pair_span[0]:context_pair_span[1]]
        source_slice = context_pair_span
        text_slice = (pair_span[0] - context_pair_span[0], pair_span[1] - context_pair_span[1])
        assert text_slice[1] <= 0, "text_slice end must be negative"
        text_slice = (text_slice[0], None if text_slice[1] == 0 else text_slice[1])

        return Chunk(
            chunk_pairs=chunk_pairs,
            source_slice=source_slice,
            text_slice=text_slice
        )
    
    def get_chunk_pair_span(self, chunk: Chunk):
        pair_start, pair_end = chunk.text_source_slice.start, chunk.text_source_slice.stop
        merged_x_chunk = ''.join(p[0] for p in self.xy_pairs[pair_start:pair_end])
        merged_y_chunk = ''.join(p[1] for p in self.xy_pairs[pair_start:pair_end])
        if merged_x_chunk == chunk.x_chunk and merged_y_chunk == chunk.y_chunk:
            return pair_start, pair_end

        pair_start, pair_end = 0, len(self.xy_pairs)
        x_chunk, y_chunk = chunk.x_chunk, chunk.y_chunk
        for i, (x, y) in enumerate(self.xy_pairs):
            if x_chunk[:50].startswith(x[:50]) and y_chunk[:50].startswith(y[:50]):
                pair_start = i
                break

        for i in range(pair_start, len(self.xy_pairs)):
            x, y = self.xy_pairs[i]
            if x_chunk[-50:].endswith(x[-50:]) and y_chunk[-50:].endswith(y[-50:]):
                pair_end = i + 1
                break

        # Verify the pair_span
        merged_x_chunk = ''.join(p[0] for p in self.xy_pairs[pair_start:pair_end])
        merged_y_chunk = ''.join(p[1] for p in self.xy_pairs[pair_start:pair_end])
        assert x_chunk == merged_x_chunk and y_chunk == merged_y_chunk, "Chunk mismatch"

        return (pair_start, pair_end)
    
    def apply_chunks(self, chunks: list[Chunk], new_chunks: list[Chunk]):
        occupied_pair_span = [False] * len(self.xy_pairs)
        pair_span_list = [self.get_chunk_pair_span(e) for e in chunks]
        for pair_span in pair_span_list:
            assert not any(occupied_pair_span[pair_span[0]:pair_span[1]]), "Chunk overlap"
            occupied_pair_span[pair_span[0]:pair_span[1]] = [True] * (pair_span[1] - pair_span[0])
        # TODO: 这里可以验证occupied_pair_span是否全被占据
        new_pairs_list = [e.text_pairs for e in new_chunks]

        sorted_spans_with_new_pairs = sorted(
            zip(pair_span_list, new_pairs_list),
            key=lambda x: x[0][0],
            reverse=True
        )

        for (start, end), new_pairs in sorted_spans_with_new_pairs:
            self.xy_pairs[start:end] = new_pairs

    def get_chunks(self, pair_span=None, chunk_length_ratio=1, context_length_ratio=1, offset_ratio=0):
        pair_span = pair_span or (0, len(self.xy_pairs))
        chunk_length = self.x_chunk_length * chunk_length_ratio, self.y_chunk_length * chunk_length_ratio
        context_length = self.x_chunk_length//2 * context_length_ratio, self.y_chunk_length//2 * context_length_ratio
        
        if 0 < offset_ratio < 1:
            offset_ratio = int(chunk_length[0] * offset_ratio), int(chunk_length[1] * offset_ratio)

        # Generate chunks
        chunks = []
        start = pair_span[0]
        cstart = self.count_span_length((0, start))  # char_start
        max_cend = self.count_span_length((0, pair_span[1]))  # char_end
        while start < pair_span[1]:
            if offset_ratio != 0:
                cend = cstart[0] + offset_ratio[0], cstart[1] + offset_ratio[1]
                offset_ratio = 0
            else:
                cend = cstart[0] + int(chunk_length[0] * 0.8), cstart[1] + int(chunk_length[1] * 0.8) # 八二原则，偷个懒，不求最优划分
            cend = min(cend[0], max_cend[0]), min(cend[1], max_cend[1])

            # 选择非零长度的span来获取chunk
            x_len, y_len = cend[0] - cstart[0], cend[1] - cstart[1]
            if x_len > 0:
                chunk1 = self.get_chunk(x_span=(cstart[0], cend[0]), context_length=context_length[0])
            if y_len > 0:
                chunk2 = self.get_chunk(y_span=(cstart[1], cend[1]), context_length=context_length[1])
            
            if x_len > 0 and y_len == 0:
                chunk = chunk1
            elif x_len == 0 and y_len > 0:
                chunk = chunk2
            elif x_len > 0 and y_len > 0:
                # 选其中source_slice更小的chunk
                chunk = chunk1 if chunk1.source_slice.stop - chunk1.source_slice.start < chunk2.source_slice.stop - chunk2.source_slice.start else chunk2
            else:
                raise ValueError("Both x_span and y_span have zero length")
             
            # assert chunk.x_chunk_context_len <= self.x_chunk_length * 2 and chunk.y_chunk_context_len <= self.y_chunk_length * 2, \
            #     "无法获取到一个足够短的区块，请调整区块长度或窗口长度！"

            chunks.append(chunk)
            start = chunk.text_source_slice.stop
            cstart = self.count_span_length((0, start))

        return chunks

    # TODO: batch_yield 可以考虑输入生成器，而不是函数及参数 
    def batch_yield(self, generators, chunks, prompt_name=None):
        # TODO: 后续考虑只输出new_chunks, 不必重复输出chunks

        # Process all pairs with the prompt and yield intermediate results
        results = [None] * len(generators)
        yields = [None] * len(generators)
        finished = [False] * len(generators)
        first_iter_flag = True
        while True:
            co_num = 0
            for i, gen in enumerate(generators):
                if finished[i]:
                    continue

                try:
                    co_num += 1
                    yield_value = next(gen)
                    yields[i] = (yield_value, chunks[i])    # TODO: yield 带上chunk是为了配合前端
                except StopIteration as e:
                    results[i] = e.value
                    finished[i] = True
                    if yields[i] is None: yields[i] = (None, chunks[i])
                
                if co_num >= self.max_thread_num:
                        break
            
            if all(finished):
                break

            if first_iter_flag and prompt_name is not None:
                yield (kp_msg := KeyPointMsg(prompt_name=prompt_name))
                first_iter_flag = False

            yield yields  # 如果是yield的值，那必定为tuple

        if not first_iter_flag and prompt_name is not None:
            yield kp_msg.set_finished()

        return results

    # 临时函数，用于配合前端，返回一个更改，对self施加该更改可以变为cur
    def diff_to(self, cur, pair_span=None):
        if pair_span is None:
            pair_span = (0, len(self.xy_pairs))
        
        if self.count_span_length(pair_span)[0] == 0:
            # 2.1版本中，章节和剧情的创作不参考x
            pair_span2 = (0 + pair_span[0], len(cur.xy_pairs) - (len(self.xy_pairs) - pair_span[1]))
            y_list = [e[1] for e in self.xy_pairs[pair_span[0]:pair_span[1]]] 
            y2_list =[e[1] for e in cur.xy_pairs[pair_span2[0]:pair_span2[1]]]
            
            y_list += ['',] * max(len(y2_list) - len(y_list), 0)
            y2_list += ['',] * max(len(y_list) - len(y2_list), 0)

            data_chunks = [('', y, y2) for y, y2 in zip(y_list, y2_list)]

            return data_chunks

        pre_pointer = 0, 1
        cur_pointer = 0, 1

        cum_sum_pre = np.cumsum([0] + [len(pair[0]) for pair in self.xy_pairs])
        cum_sum_cur = np.cumsum([0] + [len(pair[0]) for pair in cur.xy_pairs])

        apply_chunks = []

        while pre_pointer[1] <= len(self.xy_pairs) and cur_pointer[1] <= len(cur.xy_pairs):
            if cum_sum_pre[pre_pointer[1]] - cum_sum_pre[pre_pointer[0]] == cum_sum_cur[cur_pointer[1]] - cum_sum_cur[cur_pointer[0]]:
                chunk = self.get_chunk(pair_span=pre_pointer)
                value = "".join(pair[1] for pair in cur.xy_pairs[cur_pointer[0]:cur_pointer[1]])
                apply_chunks.append((chunk, 'y_chunk', value))

                pre_pointer = pre_pointer[1], pre_pointer[1] + 1
                cur_pointer = cur_pointer[1], cur_pointer[1] + 1
            elif cum_sum_pre[pre_pointer[1]] - cum_sum_pre[pre_pointer[0]] < cum_sum_cur[cur_pointer[1]] - cum_sum_cur[cur_pointer[0]]:
                pre_pointer = pre_pointer[0], pre_pointer[1] + 1
            else:
                cur_pointer = cur_pointer[0], cur_pointer[1] + 1
        
        assert pre_pointer[1] == len(self.xy_pairs) + 1 and cur_pointer[1] == len(cur.xy_pairs) + 1

        filtered_apply_chunks = []
        for e in apply_chunks:
            text_source_slice = e[0].text_source_slice
            if text_source_slice.start >= pair_span[0] and text_source_slice.stop <= pair_span[1]:
                filtered_apply_chunks.append(e)

        data_chunks = []
        for chunk, key, value in filtered_apply_chunks:
            data_chunks.append((chunk.x_chunk, chunk.y_chunk, value))

        return data_chunks

    # 临时函数，用于配合前端
    def apply_chunk(self, chunk:Chunk, key, value):
        if not isinstance(chunk, Chunk):
            chunk = Chunk(**chunk)
        new_chunk = chunk.edit(**{key: value})
        self.apply_chunks([chunk], [new_chunk])
    
    def write_text(self, chunk:Chunk, prompt_main, user_prompt_text, input_keys=None, model=None):
        chunk2prompt_key = {
            'x_chunk': 'x',
            'y_chunk': 'y',
            'x_chunk_context': 'context_x',
            'y_chunk_context': 'context_y'
        }
        
    
        if input_keys is not None:
            prompt_kwargs = {k: getattr(chunk, k) for k in input_keys}
            assert all(prompt_kwargs.values()), "Missing required context keys"
        else:
            prompt_kwargs = {k: getattr(chunk, k) for k in chunk2prompt_key.keys()}
        
        prompt_kwargs = {chunk2prompt_key.get(k, k): v for k, v in prompt_kwargs.items()}

        prompt_kwargs.update(self.global_context)   # prompt_kwargs会把所有的信息都带上，至于要用哪些由prompt决定
        
        result = yield from prompt_main(
            model=model or self.get_model(),
            user_prompt=user_prompt_text,
            **prompt_kwargs
        )

        # 为了在V2.2版本兼容summary_prompt, 后续text_key这种设计会舍弃
        update_dict = {}
        if 'text_key' in result:
            update_dict[result['text_key']] = result['text']
        else:
            update_dict['y_chunk'] = result['text']

        return chunk.edit(**update_dict)
    
    # 目前review(审阅)的评分机制暂未实装
    def review_text(self, chunk:Chunk, prompt_name, model=None):
        result = yield from prompt_review(
            model=model or self.get_model(),
            prompt_name=prompt_name,
            y=chunk.y_chunk
        )

        return result['text']

    def map_text_wo_llm(self, chunk:Chunk):
        # 该函数尝试不用LLM进行映射，目标是保证chunk.pairs中每个pair的长度合适，如果长了，进行划分，如果无法划分，报错
        new_xy_pairs = []
        for x, y in chunk.text_pairs:
            if x.strip() and not y.strip():
                x_pairs = split_text_into_chunks(x, self.x_chunk_length, min_chunk_n=1, min_chunk_size=5)
                new_xy_pairs.extend([(x_pair, y) for x_pair in x_pairs])
            elif not x.strip() and y.strip():
                y_pairs = split_text_into_chunks(y, self.y_chunk_length, min_chunk_n=1, min_chunk_size=5)
                new_xy_pairs.extend([(x, y_pair) for y_pair in y_pairs])
            else:
                if len(x) > self.x_chunk_length or len(y) > self.y_chunk_length:
                    raise ValueError("窗口太小或段落太长!考虑选择更大的窗口长度或手动分段。")
                new_xy_pairs.append((x, y))
        
        return chunk.edit(text_pairs=new_xy_pairs)

    def map_text(self, chunk:Chunk):
        # TODO: map会检查映射的内容是否大致匹配，是否有错误映射到context的情况

        if chunk.x_chunk.strip():
            x_pairs = split_text_into_chunks(chunk.x_chunk, self.x_chunk_length, min_chunk_n=1, min_chunk_size=5, max_chunk_n=20)
            assert len(x_pairs) >= len(chunk.text_pairs), "未知错误！合并所有区块后再分区块，结果更少？"
            if len(x_pairs) == len(chunk.text_pairs):
                return chunk, True, ''
        else:
            # 这说明y的创作是不参照x的，而是参照global_context
            y_pairs = split_text_into_chunks(chunk.y_chunk, self.y_chunk_length, min_chunk_n=1, min_chunk_size=5, max_chunk_n=20)
            new_xy_pairs = [('', y) for y in y_pairs]
            return chunk.edit(text_pairs=new_xy_pairs), True, ''

        try:
            y_pairs = split_text_into_chunks(chunk.y_chunk, self.y_chunk_length, min_chunk_n=len(x_pairs), min_chunk_size=5, max_chunk_n=20)
        except Exception as e:
            # 如果y_chunk不能找到更多的区块划分，干脆让x_chunk划分更少的区块
            y_pairs = split_text_into_chunks(chunk.y_chunk, self.y_chunk_length, min_chunk_n=1, min_chunk_size=5, max_chunk_n=20)
            x_pairs = split_text_into_chunks(chunk.x_chunk, self.x_chunk_length, min_chunk_n=1, min_chunk_size=5, max_chunk_n=int(0.8 * len(y_pairs)))
            
            # TODO: 这是因为目前映射Prompt的设计需要x数量小于y，后续会对Prompt进行改进

        try:
            gen = match_plot_and_text.main(
                model=self.get_sub_model(),
                plot_chunks=x_pairs,
                text_chunks=y_pairs
                )
            while True:
                yield next(gen)
        except StopIteration as e:
            output = e.value
        
        x2y = output['plot2text']
        new_xy_pairs = []
        for xi_list, yi_list in x2y:
            xl, xr = xi_list[0], xi_list[-1]
            new_xy_pairs.append(("".join(x_pairs[xl:xr+1]), "".join(y_pairs[i] for i in yi_list)))

        new_chunk = chunk.edit(text_pairs=new_xy_pairs)
        return new_chunk, True, ''
    
    def batch_map_text(self, chunks):
        results = yield from self.batch_yield(
            [self.map_text(e) for e in chunks], chunks, prompt_name='映射文本')
        return results
    
    def batch_write_apply_text(self, chunks, prompt_main, user_prompt_text):
        new_chunks = yield from self.batch_yield(
            [self.write_text(e, prompt_main, user_prompt_text) for e in chunks], 
            chunks, prompt_name='创作文本')
        
        results = yield from self.batch_map_text(new_chunks)
        new_chunks2 = [e[0] for e in results]

        self.apply_chunks(chunks, new_chunks2)

    def batch_review_write_apply_text(self, chunks, write_prompt_main, review_prompt_name):
        reviews = yield from self.batch_yield(
            [self.review_text(e, review_prompt_name) for e in chunks], 
            chunks, prompt_name='审阅文本')
        
        rewrite_instrustion = "\n\n根据审阅意见，重新创作，如果审阅意见表示无需改动，则保持原样输出。"

        new_chunks = yield from self.batch_yield(
            [self.write_text(chunk, write_prompt_main, review + rewrite_instrustion) for chunk, review in zip(chunks, reviews)], 
            chunks, prompt_name='创作文本')
        
        results = yield from self.batch_map_text(new_chunks)
        new_chunks2 = [e[0] for e in results]

        self.apply_chunks(chunks, new_chunks2)