tianyan/aistream-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

x

Browse Source
This commit is contained in:
Home
2025-12-01 23:32:34 +08:00
parent 67c44bc746
commit fee7149458
4 changed files with 410 additions and 287 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/WebSocketFrameHeader.py
+158 -87
View File
@@ -2,36 +2,41 @@ from enum import IntEnum
import struct
import json
import gzip
from typing import Optional, Union, Dict, Any, List
from typing import Optional, Union, Dict, Any, List, Tuple
# -------------------------- 协议常量定义 --------------------------
# -------------------------- 协议常量定义（与JS完全一致） --------------------------
class ProtocolConst:
PROTOCOL_VERSION = 0b0001 # 协议版本v1
HEADER_SIZE = 8 # 头部固定字节数（字节1~字节8）
MAX_BODY_SIZE = 1024 * 1024 * 10 # 最大包体大小（10MB，防止内存溢出）
STRING_ENCODING = "utf-8" # 字符串默认编码（统一编码格式，避免乱码）
PROTOCOL_VERSION = 0b0001 # 协议版本（4位，字节0低4位，0~15）
HEADER_SIZE = 8 # 头部固定字节数（字节0~7）
MAX_BODY_SIZE = 0xFFFFFF # 最大包体大小（24位，≈16MB，字节4~6存储）
STRING_ENCODING = "utf-8" # 字符串默认编码
# -------------------------- 枚举定义（业务约定） --------------------------
# -------------------------- 枚举定义（与JS码值完全对齐） --------------------------
class MessageType(IntEnum):
"""消息类型（4位，字节1高4位）"""
PING = 0b0000 # 心跳
AUDIO_DATA = 0b0001 # 纯音频数据(pcm)
TEXT_MESSAGE = 0b0010 # 纯文本消息
AUDIO_TEXT_MIX = 0b0011 # 音频+文本混合数据
"""消息类型（4位，字节0高4位，0~15）"""
PING = 0b0001 # 心跳（支持空包体）
AUDIO_DATA = 0b0010 # 纯音频数据(pcm)
TEXT_MESSAGE = 0b0011 # 纯文本消息
CONTROL_CMD = 0b0100 # 控制指令（暂停/继续等）
# 预留12种类型（0b0100 ~ 0b1111）
class SerializationType(IntEnum):
"""序列化方式（4位，字节2高4位）"""
RAW = 0b0000 # 原始二进制（音频等）
JSON = 0b0001 # JSON格式（文本/指令）
STRING = 0b0010 # 直接字符串（纯文本，UTF-8编码，无JSON包装）🔴新增类型
"""序列化方式（3位，字节1高3位，1~8）"""
RAW = 0b001 # 原始二进制（1）
JSON = 0b010 # JSON格式（2）
STRING = 0b011 # 直接字符串（3）
# 预留5种方式（0b100 ~ 0b111）
class CompressionType(IntEnum):
"""压缩方式（4位，字节2低4位）"""
NONE = 0b0000 # 无压缩
GZIP = 0b0001 # gzip压缩
"""压缩方式（3位，字节1中3位，1~8）"""
NONE = 0b001 # 无压缩（默认值1）
GZIP = 0b010 # gzip压缩（2）
# 预留6种方式（0b011 ~ 0b111）
class ControlCommand(IntEnum):
"""控制指令类型（配合MessageType.CONTROL_CMD使用）"""
@@ -40,35 +45,62 @@ class ControlCommand(IntEnum):
RESUME = 0b0011 # 继续
STOP = 0b0100 # 停止
# -------------------------- 协议工具类 --------------------------
# -------------------------- 协议工具类（与JS协议结构一致） --------------------------
class ProtocolCodec:
@staticmethod
def pack(
msg_type: MessageType,
body: Union[bytes, str, Dict[str, Any], List[Any]],
body: Union[bytes, str, Dict[str, Any], List[Any], None] = None,
sequence: int = 0,
serialization: Optional[SerializationType] = None,
compression: CompressionType = CompressionType.NONE
) -> bytes:
"""
封装协议包（头部 + 包体）
封装协议包（与JS pack 方法完全兼容）
:param msg_type: 消息类型
:param body: 包体数据（bytes/str/dict/list）
:param body: 包体数据（PING消息可传None/空，其他类型必填）
:param sequence: 消息顺序号（0~65535，默认0）
:param serialization: 序列化方式（None时自动推导）
:param compression: 压缩方式
:param compression: 压缩方式（默认无压缩）
:return: 完整协议包（bytes）
"""
# 1. 自动推导序列化方式（优化：纯文本消息自动用STRING，复杂结构用JSON）🔴修改推导逻辑
if serialization is None:
# 校验顺序号范围（0~65535）
if not isinstance(sequence, int) or not (0 <= sequence <= 0xFFFF):
raise ValueError(f"消息顺序号必须是0~65535的整数，当前传入：{sequence}")
# 特殊处理：PING消息支持空包体
if msg_type == MessageType.PING:
# PING消息强制RAW序列化
serialization = SerializationType.RAW
# 空包体转为空bytes
if body is None:
body = b""
if not isinstance(body, bytes):
raise TypeError(f"PING消息仅支持空包体或bytes类型，当前传入：{type(body)}")
else:
# 非PING消息包体不能为空
if body is None:
raise ValueError(f"非PING消息（{msg_type.name}）包体不能为空")
# 1. 自动推导序列化方式（非PING消息）
if serialization is None and msg_type != MessageType.PING:
if msg_type == MessageType.AUDIO_DATA:
serialization = SerializationType.RAW
elif msg_type == MessageType.TEXT_MESSAGE: # 纯文本消息→STRING
elif msg_type == MessageType.TEXT_MESSAGE:
serialization = SerializationType.STRING
elif msg_type in (MessageType.CONTROL_CMD, MessageType.AUDIO_TEXT_MIX): # 复杂结构→JSON
elif msg_type == MessageType.CONTROL_CMD:
serialization = SerializationType.JSON
else:
raise ValueError(f"不支持的消息类型：{msg_type}")
# 2. 序列化包体（新增STRING类型处理）🔴新增分支
# 2. 校验枚举值范围
if not (0b001 <= serialization.value <= 0b111):
raise ValueError(f"序列化方式必须在1~8（0b001~0b111）范围内，当前传入：{serialization.value}")
if not (0b001 <= compression.value <= 0b111):
raise ValueError(f"压缩方式必须在1~8（0b001~0b111）范围内，当前传入：{compression.value}")
# 3. 序列化包体
serialized_body: bytes
if serialization == SerializationType.RAW:
if not isinstance(body, bytes):
@@ -77,7 +109,7 @@ class ProtocolCodec:
elif serialization == SerializationType.STRING:
if not isinstance(body, str):
raise TypeError("STRING序列化要求body必须是str类型")
serialized_body = body.encode(ProtocolConst.STRING_ENCODING) # 直接UTF-8编码，无JSON包装
serialized_body = body.encode(ProtocolConst.STRING_ENCODING)
elif serialization == SerializationType.JSON:
if isinstance(body, str):
serialized_body = body.encode(ProtocolConst.STRING_ENCODING)
@@ -88,7 +120,7 @@ class ProtocolCodec:
else:
raise ValueError(f"不支持的序列化方式：{serialization}")
# 3. 压缩包体（不变）
# 4. 压缩包体
compressed_body: bytes
if compression == CompressionType.GZIP:
compressed_body = gzip.compress(serialized_body)
@@ -97,50 +129,81 @@ class ProtocolCodec:
else:
raise ValueError(f"不支持的压缩方式：{compression}")
# 4. 校验包体大小（不变）
# 5. 校验包体大小（24位最大支持0xFFFFFF字节）
body_len = len(compressed_body)
if body_len > ProtocolConst.MAX_BODY_SIZE:
raise OverflowError(f"包体过大（{body_len}字节），最大支持{ProtocolConst.MAX_BODY_SIZE}字节")
raise OverflowError(
f"包体过大（{body_len}字节），最大支持{ProtocolConst.MAX_BODY_SIZE}字节（≈16MB）"
)
# 5. 构造头部（不变）
byte1 = (msg_type.value << 4) | 0x00 # 消息类型(4位) + 保留位1(4位)
byte2 = (serialization.value << 4) | (compression.value & 0x0F) # 序列化(4位) + 压缩(4位)
byte3_4 = struct.pack(">H", ProtocolConst.PROTOCOL_VERSION) # 协议版本(16位)
byte5_8 = struct.pack(">I", body_len) # 包体长度(32位)
header = bytes([byte1, byte2]) + byte3_4 + byte5_8
# 6. 构造头部（与JS头部结构完全一致）
# 字节0：消息类型(高4位) + 协议版本(低4位)
byte0 = ((msg_type.value & 0x0F) << 4) | (ProtocolConst.PROTOCOL_VERSION & 0x0F)
# 字节1：序列化方式(高3位) + 压缩方式(中3位) + 保留位(低2位)
byte1 = ((serialization.value & 0x07) << 5) | ((compression.value & 0x07) << 2) | 0x00
# 字节2~3：消息顺序号（16位大端序）
byte2_3 = struct.pack(">H", sequence)
# 字节4~6：包体长度（24位大端序）
byte4 = (body_len >> 16) & 0xFF
byte5 = (body_len >> 8) & 0xFF
byte6 = body_len & 0xFF
# 字节7：保留位（固定0x00）
byte7 = 0x00
# 拼接头部
header = (
bytes([byte0, byte1]) +
byte2_3 +
bytes([byte4, byte5, byte6, byte7])
)
assert len(header) == ProtocolConst.HEADER_SIZE, f"头部长度错误：实际{len(header)}字节，预期{ProtocolConst.HEADER_SIZE}字节"
return header + compressed_body
@staticmethod
def unpack(packet: bytes) -> tuple[MessageType, SerializationType, CompressionType, Any]:
def unpack(packet: bytes) -> Tuple[MessageType, SerializationType, CompressionType, int, Any]:
"""
解析协议包
解析协议包（与JS unpack 方法完全兼容）
:param packet: 完整协议包（头部 + 包体）
:return: (消息类型, 序列化方式, 压缩方式, 原始包体数据)
:return: (消息类型, 序列化方式, 压缩方式, 顺序号, 原始包体数据)
"""
# 1. 校验包长度（不变）
# 1. 校验包长度
if len(packet) < ProtocolConst.HEADER_SIZE:
raise ValueError(f"包长度过短（{len(packet)}字节），至少需要{ProtocolConst.HEADER_SIZE}字节头部")
raise ValueError(
f"包长度过短（{len(packet)}字节），至少需要{ProtocolConst.HEADER_SIZE}字节头部"
)
# 2. 解析头部（不变）
# 2. 解析头部
header = packet[:ProtocolConst.HEADER_SIZE]
body = packet[ProtocolConst.HEADER_SIZE:]
byte1 = header[0]
byte2 = header[1]
msg_type = MessageType((byte1 >> 4) & 0x0F)
serialization = SerializationType((byte2 >> 4) & 0x0F)
compression = CompressionType(byte2 & 0x0F)
version = struct.unpack(">H", header[2:4])[0]
body_len = struct.unpack(">I", header[4:8])[0]
# 字节0：消息类型(高4位) + 协议版本(低4位)
byte0 = header[0]
msg_type = MessageType((byte0 >> 4) & 0x0F)
version = byte0 & 0x0F
# 字节1：序列化方式(高3位) + 压缩方式(中3位) + 保留位(低2位)
byte1 = header[1]
serialization = SerializationType((byte1 >> 5) & 0x07)
compression = CompressionType((byte1 >> 2) & 0x07)
# 字节2~3：消息顺序号（16位大端序）
sequence = struct.unpack(">H", header[2:4])[0]
# 字节4~6：包体长度（24位大端序），字节7：保留位（忽略）
body_len = (header[4] << 16) | (header[5] << 8) | header[6]
# 校验版本和包体长度
if version != ProtocolConst.PROTOCOL_VERSION:
raise ValueError(f"协议版本不匹配：收到v{version}，支持v{ProtocolConst.PROTOCOL_VERSION}")
raise ValueError(
f"协议版本不匹配：收到v{version}（0b{version:04b}），支持v{ProtocolConst.PROTOCOL_VERSION}（0b{ProtocolConst.PROTOCOL_VERSION:04b}）"
)
if len(body) != body_len:
raise ValueError(f"包体长度不匹配：头部声明{body_len}字节，实际{len(body)}字节")
raise ValueError(
f"包体长度不匹配：头部声明{body_len}字节，实际{len(body)}字节"
)
# 3. 解压包体（不变）
# 3. 解压包体
decompressed_body: bytes
if compression == CompressionType.GZIP:
try:
@@ -152,13 +215,15 @@ class ProtocolCodec:
else:
raise ValueError(f"不支持的压缩方式：{compression}")
# 4. 反序列化包体（新增STRING类型处理）🔴新增分支
# 4. 反序列化包体（空包体返回None）
original_body: Any
if serialization == SerializationType.RAW:
if len(decompressed_body) == 0:
original_body = None
elif serialization == SerializationType.RAW:
original_body = decompressed_body
elif serialization == SerializationType.STRING:
try:
original_body = decompressed_body.decode(ProtocolConst.STRING_ENCODING) # 直接UTF-8解码，无JSON解析
original_body = decompressed_body.decode(ProtocolConst.STRING_ENCODING)
except UnicodeDecodeError:
raise ValueError(f"STRING反序列化失败：{ProtocolConst.STRING_ENCODING}解码错误")
elif serialization == SerializationType.JSON:
@@ -166,58 +231,64 @@ class ProtocolCodec:
original_body = json.loads(decompressed_body.decode(ProtocolConst.STRING_ENCODING))
except UnicodeDecodeError:
raise ValueError(f"JSON反序列化失败：{ProtocolConst.STRING_ENCODING}解码错误")
except json.JSONDecodeError:
raise ValueError("JSON反序列化失败：格式错误")
except json.JSONDecodeError as e:
raise ValueError(f"JSON反序列化失败：格式错误（{str(e)}）")
else:
raise ValueError(f"不支持的序列化方式：{serialization}")
return msg_type, serialization, compression, original_body
return msg_type, serialization, compression, sequence, original_body
# -------------------------- 使用示例（新增纯字符串消息测试） --------------------------
# -------------------------- 使用示例（验证与JS兼容性） --------------------------
if __name__ == "__main__":
# 示例1：纯文本消息（自动用STRING序列化，无JSON包装）🔴测试新增类型
text_body = "你好，这是纯字符串消息（无JSON）"
# 示例1：PING消息（空包体，默认顺序号0）
ping_packet = ProtocolCodec.pack(MessageType.PING)
print(f"PING消息包长度：{len(ping_packet)}字节（仅头部）")
msg_type1, ser1, comp1, seq1, body1 = ProtocolCodec.unpack(ping_packet)
print(f"PING解析结果：类型={msg_type1.name}，序列化={ser1.name}，压缩={comp1.name}，顺序号={seq1}，包体={body1}\n")
# 示例2：纯文本消息（STRING序列化，指定顺序号）
text_body = "Python与JS协议兼容测试（纯字符串）"
text_packet = ProtocolCodec.pack(
msg_type=MessageType.TEXT_MESSAGE,
body=text_body,
sequence=1001,
compression=CompressionType.NONE
)
print(f"纯文本消息包长度：{len(text_packet)}字节")
msg_type1, ser1, comp1, body1 = ProtocolCodec.unpack(text_packet)
print(f"解析结果：类型={msg_type1.name}，序列化={ser1.name}，压缩={comp1.name}，内容={body1}\n")
print(f"文本消息包长度：{len(text_packet)}字节")
msg_type2, ser2, comp2, seq2, body2 = ProtocolCodec.unpack(text_packet)
print(f"文本解析结果：类型={msg_type2.name}，序列化={ser2.name}，顺序号={seq2}，内容={body2}\n")
# 示例2：控制指令（JSON序列化，复杂结构）
control_body = {"cmd": ControlCommand.HEARTBEAT.value, "timestamp": 1699999999}
# 示例3：控制指令（JSON序列化）
control_body = {"cmd": ControlCommand.PAUSE.value, "reason": "用户主动暂停"}
control_packet = ProtocolCodec.pack(
msg_type=MessageType.CONTROL_CMD,
body=control_body,
compression=CompressionType.NONE
sequence=1002
)
print(f"控制指令包长度：{len(control_packet)}字节")
msg_type2, ser2, comp2, body2 = ProtocolCodec.unpack(control_packet)
print(f"解析结果：类型={msg_type2.name}，序列化={ser2.name}，压缩={comp2.name}，内容={body2}\n")
msg_type3, ser3, comp3, seq3, body3 = ProtocolCodec.unpack(control_packet)
print(f"控制指令解析结果：类型={msg_type3.name}，序列化={ser3.name}，顺序号={seq3}，内容={body3}\n")
# 示例3：音频数据（RAW序列化）
# 示例4：音频数据（RAW序列化）
audio_body = b"\x00\x01\x02\x03\x04\x05" * 100 # 模拟PCM数据
audio_packet = ProtocolCodec.pack(
msg_type=MessageType.AUDIO_DATA,
body=audio_body,
compression=CompressionType.NONE
sequence=1003
)
print(f"音频数据包长度：{len(audio_packet)}字节")
msg_type3, ser3, comp3, body3 = ProtocolCodec.unpack(audio_packet)
print(f"解析结果：类型={msg_type3.name}，序列化={ser3.name}，压缩={comp3.name}，数据长度={len(body3)}字节\n")
msg_type4, ser4, comp4, seq4, body4 = ProtocolCodec.unpack(audio_packet)
print(f"音频解析结果：类型={msg_type4.name}，序列化={ser4.name}，顺序号={seq4}，数据长度={len(body4)}字节\n")
# 示例4：混合数据（JSON序列化）
mix_body = {
"audio_data": list(audio_body[:10]),
"text_content": "混合数据仍用JSON"
}
mix_packet = ProtocolCodec.pack(
msg_type=MessageType.AUDIO_TEXT_MIX,
body=mix_body
# 示例5：GZIP压缩测试（需JS端启用GZIP解压）
long_text_body = "这是一段很长的文本，用于测试GZIP压缩效果" * 100
gzip_packet = ProtocolCodec.pack(
msg_type=MessageType.TEXT_MESSAGE,
body=long_text_body,
sequence=1004,
compression=CompressionType.GZIP
)
print(f"混合数据包长度：{len(mix_packet)}字节")
msg_type4, ser4, comp4, body4 = ProtocolCodec.unpack(mix_packet)
print(f"解析结果：类型={msg_type4.name}，序列化={ser4.name}，压缩={comp4.name}，内容={body4}")
print(f"GZIP压缩后包长度：{len(gzip_packet)}字节（原始文本长度：{len(long_text_body.encode())}字节）")
msg_type5, ser5, comp5, seq5, body5 = ProtocolCodec.unpack(gzip_packet)
print(f"GZIP解析结果：类型={msg_type5.name}，压缩={comp5.name}，内容前50字：{body5[:50]}...")