|
|
|
|
| """
|
| 流式数据包处理器
|
|
|
| 处理流式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 |
