src/anti_truncation.py

"""
Anti-Truncation Module - Ensures complete streaming output
保持一个流式请求内完整输出的反截断模块
"""
import json
import re
from typing import Dict, Any, AsyncGenerator, List, Tuple

from fastapi.responses import StreamingResponse

from log import log

# 反截断配置
DONE_MARKER = "[done]"
MAX_CONTINUATION_ATTEMPTS = 3
CONTINUATION_PROMPT = f"""请从刚才被截断的地方继续输出剩余的所有内容。

重要提醒：
1. 不要重复前面已经输出的内容
2. 直接继续输出，无需任何前言或解释
3. 当你完整完成所有内容输出后，必须在最后一行单独输出：{DONE_MARKER}
4. {DONE_MARKER} 标记表示你的回答已经完全结束，这是必需的结束标记

现在请继续输出："""

# 正则替换配置
REGEX_REPLACEMENTS: List[Tuple[str, str, str]] = [
    (
        "age_pattern",  # 替换规则名称
        r"(?:[1-9]|1[0-8])岁(?:的)?|(?:十一|十二|十三|十四|十五|十六|十七|十八|十|一|二|三|四|五|六|七|八|九)岁(?:的)?",  # 正则模式
        ""  # 替换文本
    ),
    # 可在此处添加更多替换规则
    # ("rule_name", r"pattern", "replacement"),
]

def apply_regex_replacements(text: str) -> str:
    """
    对文本应用正则替换规则
    
    Args:
        text: 要处理的文本
        
    Returns:
        处理后的文本
    """
    if not text:
        return text
    
    processed_text = text
    replacement_count = 0
    
    for rule_name, pattern, replacement in REGEX_REPLACEMENTS:
        try:
            # 编译正则表达式，使用IGNORECASE标志
            regex = re.compile(pattern, re.IGNORECASE)
            
            # 执行替换
            new_text, count = regex.subn(replacement, processed_text)
            
            if count > 0:
                log.debug(f"Regex replacement '{rule_name}': {count} matches replaced")
                processed_text = new_text
                replacement_count += count
                
        except re.error as e:
            log.error(f"Invalid regex pattern in rule '{rule_name}': {e}")
            continue
    
    if replacement_count > 0:
        log.info(f"Applied {replacement_count} regex replacements to text")
    
    return processed_text

def apply_regex_replacements_to_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
    """
    对请求payload中的文本内容应用正则替换
    
    Args:
        payload: 请求payload
        
    Returns:
        应用替换后的payload
    """
    if not REGEX_REPLACEMENTS:
        return payload
    
    modified_payload = payload.copy()
    request_data = modified_payload.get("request", {})
    
    # 处理contents中的文本
    contents = request_data.get("contents", [])
    if contents:
        new_contents = []
        for content in contents:
            if isinstance(content, dict):
                new_content = content.copy()
                parts = new_content.get("parts", [])
                if parts:
                    new_parts = []
                    for part in parts:
                        if isinstance(part, dict) and "text" in part:
                            new_part = part.copy()
                            new_part["text"] = apply_regex_replacements(part["text"])
                            new_parts.append(new_part)
                        else:
                            new_parts.append(part)
                    new_content["parts"] = new_parts
                new_contents.append(new_content)
            else:
                new_contents.append(content)
        
        request_data["contents"] = new_contents
        modified_payload["request"] = request_data
        log.debug("Applied regex replacements to request contents")
    
    return modified_payload

def apply_anti_truncation(payload: Dict[str, Any]) -> Dict[str, Any]:
    """
    对请求payload应用反截断处理和正则替换
    在systemInstruction中添加提醒，要求模型在结束时输出DONE_MARKER标记
    
    Args:
        payload: 原始请求payload
        
    Returns:
        添加了反截断指令并应用了正则替换的payload
    """
    # 首先应用正则替换
    modified_payload = apply_regex_replacements_to_payload(payload)
    request_data = modified_payload.get("request", {})
    
    # 获取或创建systemInstruction
    system_instruction = request_data.get("systemInstruction", {})
    if not system_instruction:
        system_instruction = {"parts": []}
    elif "parts" not in system_instruction:
        system_instruction["parts"] = []
    
    # 添加反截断指令
    anti_truncation_instruction = {
        "text": f"""严格执行以下输出结束规则：

1. 当你完成完整回答时，必须在输出的最后单独一行输出：{DONE_MARKER}
2. {DONE_MARKER} 标记表示你的回答已经完全结束，这是必需的结束标记
3. 只有输出了 {DONE_MARKER} 标记，系统才认为你的回答是完整的
4. 如果你的回答被截断，系统会要求你继续输出剩余内容
5. 无论回答长短，都必须以 {DONE_MARKER} 标记结束

示例格式：
```
你的回答内容...
更多回答内容...
{DONE_MARKER}
```

注意：{DONE_MARKER} 必须单独占一行，前面不要有任何其他字符。

这个规则对于确保输出完整性极其重要，请严格遵守。"""
    }
    
    # 检查是否已经包含反截断指令
    has_done_instruction = any(
        part.get("text", "").find(DONE_MARKER) != -1 
        for part in system_instruction["parts"] 
        if isinstance(part, dict)
    )
    
    if not has_done_instruction:
        system_instruction["parts"].append(anti_truncation_instruction)
        request_data["systemInstruction"] = system_instruction
        modified_payload["request"] = request_data
        
        log.debug("Applied anti-truncation instruction to request")
    
    return modified_payload

class AntiTruncationStreamProcessor:
    """反截断流式处理器"""
    
    def __init__(self, 
                 original_request_func,
                 payload: Dict[str, Any],
                 max_attempts: int = MAX_CONTINUATION_ATTEMPTS):
        self.original_request_func = original_request_func
        self.base_payload = payload.copy()
        self.max_attempts = max_attempts
        self.collected_content = []  # 使用列表避免字符串重复拼接
        self.current_attempt = 0
        
    async def process_stream(self) -> AsyncGenerator[bytes, None]:
        """处理流式响应，检测并处理截断"""
        
        while self.current_attempt < self.max_attempts:
            self.current_attempt += 1
            
            # 构建当前请求payload
            current_payload = self._build_current_payload()
            
            log.debug(f"Anti-truncation attempt {self.current_attempt}/{self.max_attempts}")
            
            # 发送请求
            try:
                response = await self.original_request_func(current_payload)
                
                if not isinstance(response, StreamingResponse):
                    # 非流式响应，直接处理
                    yield await self._handle_non_streaming_response(response)
                    return
                
                # 处理流式响应
                chunk_content = ""
                found_done_marker = False
                
                async for chunk in response.body_iterator:
                    if not chunk:
                        yield chunk
                        continue
                    
                    # 处理不同数据类型的startswith问题
                    if isinstance(chunk, bytes):
                        if not chunk.startswith(b'data: '):
                            yield chunk
                            continue
                        payload_data = chunk[len(b'data: '):]
                    else:
                        chunk_str = str(chunk)
                        if not chunk_str.startswith('data: '):
                            yield chunk
                            continue
                        payload_data = chunk_str[len('data: '):].encode()
                    
                    # 解析chunk内容
                    
                    if payload_data.strip() == b'[DONE]':
                        # 检查是否找到了done标记
                        if found_done_marker:
                            log.info("Anti-truncation: Found [done] marker, output complete")
                            yield chunk
                            return
                        else:
                            log.warning("Anti-truncation: Stream ended without [done] marker")
                            # 不发送[DONE]，准备继续
                            break
                    
                    try:
                        data = json.loads(payload_data.decode())
                        content = self._extract_content_from_chunk(data)
                        
                        if content:
                            chunk_content += content
                            
                            # 检查是否包含done标记
                            if self._check_done_marker_in_chunk_content(content):
                                found_done_marker = True
                                log.info("Anti-truncation: Found [done] marker in chunk")
                        
                        # 清理chunk中的[done]标记后再发送
                        cleaned_chunk = self._remove_done_marker_from_chunk(chunk, data)
                        yield cleaned_chunk
                        
                    except (json.JSONDecodeError, UnicodeDecodeError):
                        yield chunk
                        continue
                
                # 更新收集的内容 - 使用列表避免字符串重复拼接
                if chunk_content:
                    self.collected_content.append(chunk_content)
                
                # 如果找到了done标记，结束
                if found_done_marker:
                    # 立即清理内容释放内存
                    self.collected_content.clear()
                    yield b'data: [DONE]\n\n'
                    return
                
                # 只有在单个chunk中没有找到done标记时，才检查累积内容（防止done标记跨chunk出现）
                if not found_done_marker:
                    accumulated_text = ''.join(self.collected_content) if self.collected_content else ""
                    if self._check_done_marker_in_text(accumulated_text):
                        log.info("Anti-truncation: Found [done] marker in accumulated content")
                        # 立即清理内容释放内存
                        self.collected_content.clear()
                        yield b'data: [DONE]\n\n'
                        return
                
                # 如果没找到done标记且不是最后一次尝试，准备续传
                if self.current_attempt < self.max_attempts:
                    total_length = sum(len(chunk) for chunk in self.collected_content) if self.collected_content else 0
                    log.info(f"Anti-truncation: No [done] marker found in output (length: {total_length}), preparing continuation (attempt {self.current_attempt + 1})")
                    if self.collected_content and total_length > 100:
                        last_chunk = self.collected_content[-1] if self.collected_content else ""
                        log.debug(f"Anti-truncation: Current collected content ends with: {'...' + last_chunk[-100:]}")
                    # 在下一次循环中会继续
                    continue
                else:
                    # 最后一次尝试，直接结束
                    log.warning("Anti-truncation: Max attempts reached, ending stream")
                    # 立即清理内容释放内存
                    self.collected_content.clear()
                    yield b'data: [DONE]\n\n'
                    return
                
            except Exception as e:
                log.error(f"Anti-truncation error in attempt {self.current_attempt}: {str(e)}")
                if self.current_attempt >= self.max_attempts:
                    # 发送错误chunk
                    error_chunk = {
                        "error": {
                            "message": f"Anti-truncation failed: {str(e)}",
                            "type": "api_error",
                            "code": 500
                        }
                    }
                    yield f"data: {json.dumps(error_chunk)}\n\n".encode()
                    yield b'data: [DONE]\n\n'
                    return
                # 否则继续下一次尝试
                
        # 如果所有尝试都失败了
        log.error("Anti-truncation: All attempts failed")
        # 确保清理内容释放内存
        self.collected_content.clear()
        yield b'data: [DONE]\n\n'
    
    def _build_current_payload(self) -> Dict[str, Any]:
        """构建当前请求的payload"""
        if self.current_attempt == 1:
            # 第一次请求，使用原始payload（已经包含反截断指令）
            return self.base_payload
        
        # 后续请求，添加续传指令
        continuation_payload = self.base_payload.copy()
        request_data = continuation_payload.get("request", {})
        
        # 获取原始对话内容
        contents = request_data.get("contents", [])
        new_contents = contents.copy()
        
        # 如果有收集到的内容，添加到对话中
        if self.collected_content:
            # 拼接收集的内容并添加模型的回复
            accumulated_text = ''.join(self.collected_content)
            new_contents.append({
                "role": "model",
                "parts": [{"text": accumulated_text}]
            })
        
        # 构建具体的续写指令，包含前面的内容摘要
        content_summary = ""
        if self.collected_content:
            accumulated_text = ''.join(self.collected_content)
            if len(accumulated_text) > 200:
                content_summary = f"\n\n前面你已经输出了约 {len(accumulated_text)} 个字符的内容，结尾是：\n\"...{accumulated_text[-100:]}\""
            else:
                content_summary = f"\n\n前面你已经输出的内容是：\n\"{accumulated_text}\""
        
        detailed_continuation_prompt = f"""{CONTINUATION_PROMPT}{content_summary}"""
        
        # 添加继续指令
        continuation_message = {
            "role": "user", 
            "parts": [{"text": detailed_continuation_prompt}]
        }
        new_contents.append(continuation_message)
        
        request_data["contents"] = new_contents
        continuation_payload["request"] = request_data
        
        return continuation_payload
    
    def _extract_content_from_chunk(self, data: Dict[str, Any]) -> str:
        """从chunk数据中提取文本内容"""
        content = ""
        
        # 处理Gemini格式
        if "candidates" in data:
            for candidate in data["candidates"]:
                if "content" in candidate:
                    parts = candidate["content"].get("parts", [])
                    for part in parts:
                        if "text" in part:
                            content += part["text"]
        
        # 处理OpenAI格式
        elif "choices" in data:
            for choice in data["choices"]:
                if "delta" in choice and "content" in choice["delta"]:
                    content += choice["delta"]["content"]
                elif "message" in choice and "content" in choice["message"]:
                    content += choice["message"]["content"]
        
        return content
    
    async def _handle_non_streaming_response(self, response) -> bytes:
        """处理非流式响应"""
        try:
            if hasattr(response, 'body'):
                content = response.body.decode() if isinstance(response.body, bytes) else response.body
            elif hasattr(response, 'content'):
                content = response.content.decode() if isinstance(response.content, bytes) else response.content
            else:
                content = str(response)
            
            response_data = json.loads(content)
            
            # 检查是否包含done标记
            text_content = self._extract_content_from_response(response_data)
            has_done_marker = self._check_done_marker_in_text(text_content)
            
            if not has_done_marker and self.current_attempt < self.max_attempts:
                log.info("Anti-truncation: Non-streaming response needs continuation")
                if text_content:
                    self.collected_content.append(text_content)
                # 递归处理续传
                return await self._handle_non_streaming_response(
                    await self.original_request_func(self._build_current_payload())
                )
            
            return content.encode()
            
        except Exception as e:
            log.error(f"Anti-truncation non-streaming error: {str(e)}")
            return json.dumps({
                "error": {
                    "message": f"Anti-truncation failed: {str(e)}",
                    "type": "api_error",
                    "code": 500
                }
            }).encode()
    
    def _check_done_marker_in_text(self, text: str) -> bool:
        """检测文本中是否包含DONE_MARKER（只检测指定标记）"""
        if not text:
            return False

        # 只要文本中出现DONE_MARKER即可
        return DONE_MARKER in text
    
    def _check_done_marker_in_chunk_content(self, content: str) -> bool:
        """检查单个chunk内容中是否包含done标记"""
        return self._check_done_marker_in_text(content)
    
    def _extract_content_from_response(self, data: Dict[str, Any]) -> str:
        """从响应数据中提取文本内容"""
        content = ""
        
        # 处理Gemini格式
        if "candidates" in data:
            for candidate in data["candidates"]:
                if "content" in candidate:
                    parts = candidate["content"].get("parts", [])
                    for part in parts:
                        if "text" in part:
                            content += part["text"]
        
        # 处理OpenAI格式
        elif "choices" in data:
            for choice in data["choices"]:
                if "message" in choice and "content" in choice["message"]:
                    content += choice["message"]["content"]
        
        return content
    
    def _remove_done_marker_from_chunk(self, chunk: bytes, data: Dict[str, Any]) -> bytes:
        """使用正则表达式从chunk中移除[done]标记"""
        try:
            # 首先检查是否真的包含[done]标记，如果没有则直接返回原始chunk
            chunk_text = chunk.decode('utf-8', errors='ignore') if isinstance(chunk, bytes) else str(chunk)
            if '[done]' not in chunk_text.lower():
                return chunk  # 没有[done]标记，直接返回原始chunk
            
            # 编译正则表达式，匹配[done]标记（忽略大小写，包括可能的空白字符）
            done_pattern = re.compile(r'\s*\[done\]\s*', re.IGNORECASE)
            
            # 处理Gemini格式
            if "candidates" in data:
                modified_data = data.copy()
                modified_data["candidates"] = []
                
                for i, candidate in enumerate(data["candidates"]):
                    modified_candidate = candidate.copy()
                    # 只在最后一个candidate中清理[done]标记
                    is_last_candidate = (i == len(data["candidates"]) - 1)
                    
                    if "content" in candidate:
                        modified_content = candidate["content"].copy()
                        if "parts" in modified_content:
                            modified_parts = []
                            for part in modified_content["parts"]:
                                if "text" in part and isinstance(part["text"], str):
                                    modified_part = part.copy()
                                    # 只在最后一个candidate中清理[done]标记
                                    if is_last_candidate:
                                        modified_part["text"] = done_pattern.sub('', part["text"])
                                    modified_parts.append(modified_part)
                                else:
                                    modified_parts.append(part)
                            modified_content["parts"] = modified_parts
                        modified_candidate["content"] = modified_content
                    modified_data["candidates"].append(modified_candidate)
                
                # 重新编码为chunk格式，保持原始的换行符
                if isinstance(chunk, bytes):
                    prefix = b'data: '
                    suffix = b'\n\n'  # 确保有正确的换行符
                    json_data = json.dumps(modified_data, separators=(',',':'), ensure_ascii=False).encode('utf-8')
                    return prefix + json_data + suffix
                else:
                    return f"data: {json.dumps(modified_data, separators=(',',':'), ensure_ascii=False)}\n\n"
            
            # 处理OpenAI格式
            elif "choices" in data:
                modified_data = data.copy()
                modified_data["choices"] = []
                
                for choice in data["choices"]:
                    modified_choice = choice.copy()
                    if "delta" in choice and "content" in choice["delta"]:
                        modified_delta = choice["delta"].copy()
                        modified_delta["content"] = done_pattern.sub('', choice["delta"]["content"])
                        modified_choice["delta"] = modified_delta
                    elif "message" in choice and "content" in choice["message"]:
                        modified_message = choice["message"].copy()
                        modified_message["content"] = done_pattern.sub('', choice["message"]["content"])
                        modified_choice["message"] = modified_message
                    modified_data["choices"].append(modified_choice)
                
                # 重新编码为chunk格式，保持原始的换行符
                if isinstance(chunk, bytes):
                    prefix = b'data: '
                    suffix = b'\n\n'  # 确保有正确的换行符
                    json_data = json.dumps(modified_data, separators=(',',':'), ensure_ascii=False).encode('utf-8')
                    return prefix + json_data + suffix
                else:
                    return f"data: {json.dumps(modified_data, separators=(',',':'), ensure_ascii=False)}\n\n"
            
            # 如果没有找到支持的格式，返回原始chunk
            return chunk
            
        except Exception as e:
            log.warning(f"Failed to remove [done] marker from chunk: {str(e)}")
            return chunk

async def apply_anti_truncation_to_stream(
    request_func,
    payload: Dict[str, Any],
    max_attempts: int = MAX_CONTINUATION_ATTEMPTS
) -> StreamingResponse:
    """
    对流式请求应用反截断处理
    
    Args:
        request_func: 原始请求函数
        payload: 请求payload
        max_attempts: 最大续传尝试次数
        
    Returns:
        处理后的StreamingResponse
    """
    
    # 首先对payload应用反截断指令
    anti_truncation_payload = apply_anti_truncation(payload)
    
    # 创建反截断处理器
    processor = AntiTruncationStreamProcessor(
        lambda p: request_func(p),
        anti_truncation_payload,
        max_attempts
    )
    
    # 返回包装后的流式响应
    return StreamingResponse(
        processor.process_stream(),
        media_type="text/event-stream"
    )

def is_anti_truncation_enabled(request_data: Dict[str, Any]) -> bool:
    """
    检查请求是否启用了反截断功能
    
    Args:
        request_data: 请求数据
        
    Returns:
        是否启用反截断
    """
    return request_data.get("enable_anti_truncation", False)