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