""" Anti-Truncation Module - Ensures complete streaming output 保持一个流式请求内完整输出的反截断模块 """ import io import json import re from typing import Any, AsyncGenerator, Dict, List, Tuple from fastapi.responses import StreamingResponse from log import log # 反截断配置 DONE_MARKER = "[done]" 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 = 3, enable_prefill_mode: bool = False, ): self.original_request_func = original_request_func self.base_payload = payload.copy() self.max_attempts = max_attempts self.enable_prefill_mode = enable_prefill_mode # 使用 StringIO 避免字符串拼接的内存问题 self.collected_content = io.StringIO() self.current_attempt = 0 def _get_collected_text(self) -> str: """获取收集的文本内容""" return self.collected_content.getvalue() def _append_content(self, content: str): """追加内容到收集器""" if content: self.collected_content.write(content) def _clear_content(self): """清空收集的内容，释放内存""" self.collected_content.close() self.collected_content = io.StringIO() 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_buffer = io.StringIO() # 使用 StringIO 缓存当前轮次的chunk found_done_marker = False async for line in response.body_iterator: if not line: yield line continue # 处理上游生成器 yield 出 Response 对象的情况（错误响应） from fastapi import Response as FastAPIResponse if isinstance(line, FastAPIResponse): log.error(f"Anti-truncation: Received Response object from stream (status={line.status_code}), treating as error") error_chunk = { "error": { "message": line.body.decode('utf-8', errors='ignore') if hasattr(line, 'body') and line.body else "Upstream error", "type": "api_error", "code": line.status_code, } } yield f"data: {json.dumps(error_chunk)}\n\n".encode() yield b"data: [DONE]\n\n" return # 处理 bytes 类型的流式数据 if isinstance(line, bytes): # 解码 bytes 为字符串 line_str = line.decode('utf-8', errors='ignore').strip() else: line_str = str(line).strip() # 跳过空行 if not line_str: yield line continue # 处理 SSE 格式的数据行 if line_str.startswith("data: "): payload_str = line_str[6:] # 去掉 "data: " 前缀 # 检查是否是 [DONE] 标记 if payload_str.strip() == "[DONE]": if found_done_marker: log.info("Anti-truncation: Found [done] marker, output complete") yield line # 清理内存 chunk_buffer.close() self._clear_content() return else: log.warning("Anti-truncation: Stream ended without [done] marker") # 不发送[DONE]，准备继续 break # 尝试解析 JSON 数据 try: data = json.loads(payload_str) content = self._extract_content_from_chunk(data) log.debug(f"Anti-truncation: Extracted content: {repr(content[:100] if content else '')}") if content: chunk_buffer.write(content) # 检查是否包含done标记 has_marker = self._check_done_marker_in_chunk_content(content) log.debug(f"Anti-truncation: Check done marker result: {has_marker}, DONE_MARKER='{DONE_MARKER}'") if has_marker: found_done_marker = True log.debug(f"Anti-truncation: Found [done] marker in chunk, content: {content[:200]}") # 清理行中的[done]标记后再发送 cleaned_line = self._remove_done_marker_from_line(line, line_str, data) yield cleaned_line except (json.JSONDecodeError, ValueError): # 无法解析的行，直接传递 yield line continue else: # 非 data: 开头的行，直接传递 yield line # 更新收集的内容 - 使用 StringIO 高效处理 chunk_text = chunk_buffer.getvalue() if chunk_text: self._append_content(chunk_text) chunk_buffer.close() log.debug(f"Anti-truncation: After processing stream, found_done_marker={found_done_marker}") # 如果找到了done标记，结束 if found_done_marker: # 立即清理内容释放内存 self._clear_content() yield b"data: [DONE]\n\n" return # 只有在单个chunk中没有找到done标记时，才检查累积内容（防止done标记跨chunk出现） if not found_done_marker: accumulated_text = self._get_collected_text() if self._check_done_marker_in_text(accumulated_text): log.info("Anti-truncation: Found [done] marker in accumulated content") # 立即清理内容释放内存 self._clear_content() yield b"data: [DONE]\n\n" return # 如果没找到done标记且不是最后一次尝试，准备续传 if self.current_attempt < self.max_attempts: accumulated_text = self._get_collected_text() total_length = len(accumulated_text) log.info( f"Anti-truncation: No [done] marker found in output (length: {total_length}), preparing continuation (attempt {self.current_attempt + 1})" ) if total_length > 100: log.debug( f"Anti-truncation: Current collected content ends with: ...{accumulated_text[-100:]}" ) # 在下一次循环中会继续 continue else: # 最后一次尝试，直接结束 log.warning("Anti-truncation: Max attempts reached, ending stream") # 立即清理内容释放内存 self._clear_content() 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._clear_content() 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() # 如果有收集到的内容，添加到对话中 accumulated_text = self._get_collected_text() if accumulated_text: new_contents.append({"role": "model", "parts": [{"text": accumulated_text}]}) # 预填充模式：直接用拼接内容作为末尾 model 预填充，不再增加 user 续写指令 if self.enable_prefill_mode: log.debug("Anti-truncation: Using prefill continuation mode (no user continuation prompt)") request_data["contents"] = new_contents continuation_payload["request"] = request_data return continuation_payload # 构建具体的续写指令，包含前面的内容摘要 content_summary = "" if accumulated_text: 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 = "" # 先尝试解包 response 字段（Gemini API 格式） if "response" in data: data = data["response"] # 处理 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 流式格式（choices/delta） elif "choices" in data: for choice in data["choices"]: if "delta" in choice and "content" in choice["delta"]: delta_content = choice["delta"]["content"] if delta_content: content += delta_content return content async def _handle_non_streaming_response(self, response) -> bytes: """处理非流式响应 - 使用循环代替递归避免栈溢出""" # 使用循环代替递归 while True: try: # 特殊处理：如果返回的是StreamingResponse，需要读取其body_iterator if isinstance(response, StreamingResponse): log.error("Anti-truncation: Received StreamingResponse in non-streaming handler - this should not happen") # 尝试读取流式响应的内容 chunks = [] async for chunk in response.body_iterator: chunks.append(chunk) content = b"".join(chunks).decode() if chunks else "" # 提取响应内容 elif 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: log.error(f"Anti-truncation: Unknown response type: {type(response)}") content = str(response) # 验证内容不为空 if not content or not content.strip(): log.error("Anti-truncation: Received empty response content") return json.dumps( { "error": { "message": "Empty response from server", "type": "api_error", "code": 500, } } ).encode() # 尝试解析 JSON try: response_data = json.loads(content) except json.JSONDecodeError as json_err: log.error(f"Anti-truncation: Failed to parse JSON response: {json_err}, content: {content[:200]}") # 如果不是 JSON，直接返回原始内容 return content.encode() if isinstance(content, str) else content # 检查是否包含done标记 text_content = self._extract_content_from_response(response_data) has_done_marker = self._check_done_marker_in_text(text_content) if has_done_marker or self.current_attempt >= self.max_attempts: # 找到done标记或达到最大尝试次数，返回结果 return content.encode() if isinstance(content, str) else content # 需要继续，收集内容并构建下一个请求 if text_content: self._append_content(text_content) log.info("Anti-truncation: Non-streaming response needs continuation") # 增加尝试次数 self.current_attempt += 1 # 构建续传payload并发送下一个请求 next_payload = self._build_current_payload() response = await self.original_request_func(next_payload) # 继续循环处理下一个响应 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 = "" # 先尝试解包 response 字段（Gemini API 格式） if "response" in data: data = data["response"] # 处理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_line(self, line: bytes, line_str: str, data: Dict[str, Any]) -> bytes: """从行中移除[done]标记""" try: # 首先检查是否真的包含[done]标记 if "[done]" not in line_str.lower(): return line # 没有[done]标记，直接返回原始行 log.info(f"Anti-truncation: Attempting to remove [done] marker from line") log.debug(f"Anti-truncation: Original line (first 200 chars): {line_str[:200]}") # 编译正则表达式，匹配[done]标记（忽略大小写，包括可能的空白字符） done_pattern = re.compile(r"\s*\[done\]\s*", re.IGNORECASE) # 检查是否有 response 包裹层 has_response_wrapper = "response" in data log.debug(f"Anti-truncation: has_response_wrapper={has_response_wrapper}, data keys={list(data.keys())}") if has_response_wrapper: # 需要保留外层的 response 字段 inner_data = data["response"] else: inner_data = data log.debug(f"Anti-truncation: inner_data keys={list(inner_data.keys())}") log.debug(f"Anti-truncation: inner_data keys={list(inner_data.keys())}") # 处理Gemini格式 if "candidates" in inner_data: log.info(f"Anti-truncation: Processing Gemini format to remove [done] marker") modified_inner = inner_data.copy() modified_inner["candidates"] = [] for i, candidate in enumerate(inner_data["candidates"]): modified_candidate = candidate.copy() # 只在最后一个candidate中清理[done]标记 is_last_candidate = i == len(inner_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() original_text = part["text"] # 只在最后一个candidate中清理[done]标记 if is_last_candidate: modified_part["text"] = done_pattern.sub("", part["text"]) if "[done]" in original_text.lower(): log.debug(f"Anti-truncation: Removed [done] from text: '{original_text[:100]}' -> '{modified_part['text'][:100]}'") modified_parts.append(modified_part) else: modified_parts.append(part) modified_content["parts"] = modified_parts modified_candidate["content"] = modified_content modified_inner["candidates"].append(modified_candidate) # 如果有 response 包裹层，需要重新包装 if has_response_wrapper: modified_data = data.copy() modified_data["response"] = modified_inner else: modified_data = modified_inner # 重新编码为行格式 - SSE格式需要两个换行符 json_str = json.dumps(modified_data, separators=(",", ":"), ensure_ascii=False) result = f"data: {json_str}\n\n".encode("utf-8") log.debug(f"Anti-truncation: Modified line (first 200 chars): {result.decode('utf-8', errors='ignore')[:200]}") return result # 处理OpenAI格式 elif "choices" in inner_data: modified_inner = inner_data.copy() modified_inner["choices"] = [] for choice in inner_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_inner["choices"].append(modified_choice) # 如果有 response 包裹层，需要重新包装 if has_response_wrapper: modified_data = data.copy() modified_data["response"] = modified_inner else: modified_data = modified_inner # 重新编码为行格式 - SSE格式需要两个换行符 json_str = json.dumps(modified_data, separators=(",", ":"), ensure_ascii=False) return f"data: {json_str}\n\n".encode("utf-8") # 如果没有找到支持的格式，返回原始行 return line except Exception as e: log.warning(f"Failed to remove [done] marker from line: {str(e)}") return line async def apply_anti_truncation_to_stream( request_func, payload: Dict[str, Any], max_attempts: int = 3, enable_prefill_mode: bool = False, ) -> StreamingResponse: """ 对流式请求应用反截断处理 Args: request_func: 原始请求函数 payload: 请求payload max_attempts: 最大续传尝试次数 enable_prefill_mode: 是否启用预填充模式。启用后续传请求不再添加 user 续写指令， 而是将已收集内容作为末尾 model 内容进行预填充 Returns: 处理后的StreamingResponse """ # 首先对payload应用反截断指令 anti_truncation_payload = apply_anti_truncation(payload) # 创建反截断处理器 processor = AntiTruncationStreamProcessor( lambda p: request_func(p), anti_truncation_payload, max_attempts, enable_prefill_mode, ) # 返回包装后的流式响应 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)