#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 流式数据包处理器 处理流式protobuf数据包，支持实时解析和WebSocket推送。 """ import asyncio import json import base64 from typing import AsyncGenerator, List, Dict, Any, Optional from datetime import datetime from .logging import logger from .protobuf_utils import protobuf_to_dict class StreamProcessor: """流式数据包处理器""" def __init__(self, websocket_manager=None): self.websocket_manager = websocket_manager self.active_streams: Dict[str, StreamSession] = {} async def create_stream_session(self, stream_id: str, message_type: str = "warp.multi_agent.v1.Response") -> 'StreamSession': """创建流式会话""" session = StreamSession(stream_id, message_type, self.websocket_manager) self.active_streams[stream_id] = session logger.info(f"创建流式会话: {stream_id}, 消息类型: {message_type}") return session async def get_stream_session(self, stream_id: str) -> Optional['StreamSession']: """获取流式会话""" return self.active_streams.get(stream_id) async def close_stream_session(self, stream_id: str): """关闭流式会话""" if stream_id in self.active_streams: session = self.active_streams[stream_id] await session.close() del self.active_streams[stream_id] logger.info(f"关闭流式会话: {stream_id}") async def process_stream_chunk(self, stream_id: str, chunk_data: bytes) -> Dict[str, Any]: """处理流式数据块""" session = await self.get_stream_session(stream_id) if not session: raise ValueError(f"流式会话不存在: {stream_id}") return await session.process_chunk(chunk_data) async def finalize_stream(self, stream_id: str) -> Dict[str, Any]: """完成流式处理""" session = await self.get_stream_session(stream_id) if not session: raise ValueError(f"流式会话不存在: {stream_id}") result = await session.finalize() await self.close_stream_session(stream_id) return result class StreamSession: """流式会话""" def __init__(self, session_id: str, message_type: str, websocket_manager=None): self.session_id = session_id self.message_type = message_type self.websocket_manager = websocket_manager self.chunks: List[bytes] = [] self.chunk_count = 0 self.total_size = 0 self.start_time = datetime.now() self.parsed_chunks: List[Dict] = [] self.complete_message: Optional[Dict] = None async def process_chunk(self, chunk_data: bytes) -> Dict[str, Any]: """处理单个数据块""" self.chunk_count += 1 self.total_size += len(chunk_data) self.chunks.append(chunk_data) logger.debug(f"流式会话 {self.session_id}: 处理数据块 {self.chunk_count}, 大小 {len(chunk_data)} 字节") chunk_result = { "chunk_index": self.chunk_count - 1, "size": len(chunk_data), "timestamp": datetime.now().isoformat() } try: chunk_json = protobuf_to_dict(chunk_data, self.message_type) chunk_result["json_data"] = chunk_json chunk_result["parsed_successfully"] = True self.parsed_chunks.append(chunk_json) if self.websocket_manager: await self.websocket_manager.broadcast({ "event": "stream_chunk_parsed", "stream_id": self.session_id, "chunk": chunk_result }) except Exception as e: chunk_result["error"] = str(e) chunk_result["parsed_successfully"] = False logger.warning(f"数据块解析失败: {e}") if self.websocket_manager: await self.websocket_manager.broadcast({ "event": "stream_chunk_error", "stream_id": self.session_id, "chunk": chunk_result }) return chunk_result async def finalize(self) -> Dict[str, Any]: """完成流式处理，尝试拼接完整消息""" duration = (datetime.now() - self.start_time).total_seconds() logger.info(f"流式会话 {self.session_id} 完成: {self.chunk_count} 块, 总大小 {self.total_size} 字节, 耗时 {duration:.2f}s") result = { "session_id": self.session_id, "chunk_count": self.chunk_count, "total_size": self.total_size, "duration_seconds": duration, "chunks": [] } for i, chunk in enumerate(self.chunks): chunk_info = { "index": i, "size": len(chunk), "hex_preview": chunk[:32].hex() if len(chunk) >= 32 else chunk.hex() } if i < len(self.parsed_chunks): chunk_info["parsed_data"] = self.parsed_chunks[i] result["chunks"].append(chunk_info) try: complete_data = b''.join(self.chunks) complete_json = protobuf_to_dict(complete_data, self.message_type) result["complete_message"] = { "size": len(complete_data), "json_data": complete_json, "assembly_successful": True } self.complete_message = complete_json logger.info(f"流式消息拼接成功: {len(complete_data)} 字节") except Exception as e: result["complete_message"] = { "error": str(e), "assembly_successful": False } logger.warning(f"流式消息拼接失败: {e}") if self.websocket_manager: await self.websocket_manager.broadcast({ "event": "stream_completed", "stream_id": self.session_id, "result": result }) return result async def close(self): """关闭会话""" self.chunks.clear() self.parsed_chunks.clear() self.complete_message = None logger.debug(f"流式会话 {self.session_id} 已关闭") class StreamPacketAnalyzer: """流式数据包分析器""" @staticmethod def analyze_chunk_patterns(chunks: List[bytes]) -> Dict[str, Any]: if not chunks: return {"error": "无数据块"} analysis = { "total_chunks": len(chunks), "size_distribution": {}, "size_stats": {}, "pattern_analysis": {} } sizes = [len(chunk) for chunk in chunks] analysis["size_stats"] = { "min": min(sizes), "max": max(sizes), "avg": sum(sizes) / len(sizes), "total": sum(sizes) } size_ranges = [(0, 100), (100, 500), (500, 1000), (1000, 5000), (5000, float('inf'))] for start, end in size_ranges: range_name = f"{start}-{end if end != float('inf') else '∞'}" count = sum(1 for size in sizes if start <= size < end) analysis["size_distribution"][range_name] = count if len(chunks) >= 2: first_bytes = [chunk[:4].hex() if len(chunk) >= 4 else chunk.hex() for chunk in chunks[:5]] analysis["pattern_analysis"]["first_bytes_samples"] = first_bytes if chunks: common_prefix_len = 0 first_chunk = chunks[0] for i in range(min(len(first_chunk), 10)): if all(len(chunk) > i and chunk[i] == first_chunk[i] for chunk in chunks[1:]): common_prefix_len = i + 1 else: break if common_prefix_len > 0: analysis["pattern_analysis"]["common_prefix_length"] = common_prefix_len analysis["pattern_analysis"]["common_prefix_hex"] = first_chunk[:common_prefix_len].hex() return analysis @staticmethod def extract_streaming_deltas(parsed_chunks: List[Dict]) -> List[Dict]: if not parsed_chunks: return [] deltas = [] previous_content = "" for i, chunk in enumerate(parsed_chunks): delta = { "chunk_index": i, "timestamp": datetime.now().isoformat() } current_content = StreamPacketAnalyzer._extract_text_content(chunk) if current_content and current_content != previous_content: if previous_content and current_content.startswith(previous_content): delta["content_delta"] = current_content[len(previous_content):] delta["delta_type"] = "append" else: delta["content_delta"] = current_content delta["delta_type"] = "replace" delta["total_content_length"] = len(current_content) previous_content = current_content else: delta["content_delta"] = "" delta["delta_type"] = "no_change" if i > 0: delta["field_changes"] = StreamPacketAnalyzer._compare_dicts(parsed_chunks[i-1], chunk) deltas.append(delta) return deltas @staticmethod def _extract_text_content(data: Dict) -> str: text_paths = [ ["content"], ["text"], ["message"], ["agent_output", "text"], ["choices", 0, "delta", "content"], ["choices", 0, "message", "content"] ] for path in text_paths: try: current = data for key in path: if isinstance(current, dict) and key in current: current = current[key] elif isinstance(current, list) and isinstance(key, int) and 0 <= key < len(current): current = current[key] else: break else: if isinstance(current, str): return current except Exception: continue return "" @staticmethod def _compare_dicts(dict1: Dict, dict2: Dict, prefix: str = "") -> List[str]: changes = [] all_keys = set(dict1.keys()) | set(dict2.keys()) for key in all_keys: current_path = f"{prefix}.{key}" if prefix else key if key not in dict1: changes.append(f"添加: {current_path}") elif key not in dict2: changes.append(f"删除: {current_path}") elif dict1[key] != dict2[key]: if isinstance(dict1[key], dict) and isinstance(dict2[key], dict): changes.extend(StreamPacketAnalyzer._compare_dicts(dict1[key], dict2[key], current_path)) else: changes.append(f"修改: {current_path}") return changes[:10] _global_processor: Optional[StreamProcessor] = None def get_stream_processor() -> StreamProcessor: global _global_processor if _global_processor is None: _global_processor = StreamProcessor() return _global_processor def set_websocket_manager(manager): processor = get_stream_processor() processor.websocket_manager = manager