6 Commits
Author SHA1 Message Date
zimonianhua 223ead1292 refactor(docs): 调整 README 结构,将问题反馈移至文末
Validation And Formatting / HACS (push) Has been cancelled
Validation And Formatting / Hassfest (push) Has been cancelled
2026-05-01 22:46:04 +08:00
zimonianhua c629ad306b docs: 在 README 中添加问题反馈链接
Validation And Formatting / HACS (push) Has been cancelled
Validation And Formatting / Hassfest (push) Has been cancelled
2026-05-01 22:42:27 +08:00
zimonianhua 10ccc01c6d perf(asr): 优化 SSL 上下文创建,避免阻塞事件循环
- 将 SSL 上下文初始化从 __init__ 延迟到 __aenter__
- 使用 asyncio.to_thread 在线程池中执行 SSL 上下文创建
- 避免在同步构造函数中执行潜在的阻塞操作
2026-05-01 22:40:17 +08:00
zimonianhua 2ffc20b14c refactor(asr): 重构 ASR 模块架构,优化代码组织
Validation And Formatting / HACS (push) Has been cancelled
Validation And Formatting / Hassfest (push) Has been cancelled
- 提取 CredentialManager 凭据管理类到独立模块 (credential.py)
- 提取 ASR 客户端实现到独立模块 (client.py)
- 提取响应解析器到独立模块
- 提取协议构建器到独立模块
- 提取数据模型定义到独立模块
- 新增配置选项热更新支持,无需重启即可生效
- 简化配置验证逻辑,延迟凭据初始化至首次转录
- 更新 manifest 文档链接和 iot_class 配置
2026-05-01 22:32:19 +08:00
zimonianhua 6f28d70b0a feat(ner): 新增 NER 命名实体识别功能
Validation And Formatting / HACS (push) Has been cancelled
Validation And Formatting / Hassfest (push) Has been cancelled
- 新增 NER 接口支持 (ner.py)
- 新增 SAMI 认证模块 (sami.py)
- 新增 Wave 加密协议客户端 (wave_client.py)
- 扩展 ASR 响应结构支持详细识别结果
- 扩展设备凭据支持 sami_token 和 wave_session
- 重构 AudioEncoder 为同步实现,优化异步文件操作
- 移除不安全的 SSL 证书验证禁用
2026-05-01 22:06:11 +08:00
zimonianhua a8c13fd908 修正 HACS 安装步骤中的格式问题
Validation And Formatting / HACS (push) Has been cancelled
Validation And Formatting / Hassfest (push) Has been cancelled
2026-02-06 09:26:02 +08:00
18 changed files with 1607 additions and 874 deletions
+6 -2
View File
@@ -2,7 +2,6 @@
豆包语音识别 Home Assistant 集成插件,基于 [doubaoime-asr](https://github.com/starccy/doubaoime-asr) 开发。 豆包语音识别 Home Assistant 集成插件,基于 [doubaoime-asr](https://github.com/starccy/doubaoime-asr) 开发。
## 功能特性 ## 功能特性
- ✅ 完全兼容 Home Assistant STT 组件规范 - ✅ 完全兼容 Home Assistant STT 组件规范
@@ -30,7 +29,7 @@
### 方法 2: HACS 安装(推荐) ### 方法 2: HACS 安装(推荐)
1. 确保你已经在 Home Assistant 中安装并配置了 HACSHome Assistant Community Store)。 1. 确保你已经在 Home Assistant 中安装并配置了 HACSHome Assistant Community Store)。
2. 点击 HACS 界面左上角的菜单,选择 自定义仓库(Custom repositories),在弹出的窗口中输入仓库地址 https://github.com/xyzmos/hass_stt_doubao,类别选择 集成,然后点击 添加。之后再搜索 "Doubao STT" 并安装。 2. 点击 HACS 界面左上角的菜单,选择 自定义仓库(Custom repositories),在弹出的窗口中输入仓库地址`https://github.com/xyzmos/hass_stt_doubao`,类别选择`集成`,然后点击 `添加`。之后再搜索 "Doubao STT" 并安装。
3. 重启 Home Assistant 3. 重启 Home Assistant
## 配置 ## 配置
@@ -140,6 +139,11 @@ custom_components/hass_stt_doubao/
- 不保证未来的可用性和稳定性 - 不保证未来的可用性和稳定性
- 服务端协议可能随时变更导致功能失效 - 服务端协议可能随时变更导致功能失效
## 问题反馈
[HomeAssistant 豆包输入法 STT 语音识别插件
反馈](https://bbs.nextrt.com/d/3-homeassistant-dou-bao-shu-ru-fa-stt-yu-yin-shi-bie-cha-jian)
## 许可证 ## 许可证
MIT License MIT License
+27 -13
View File
@@ -25,21 +25,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up Doubao Speech-to-Text from a config entry.""" """Set up Doubao Speech-to-Text from a config entry."""
hass.data.setdefault(DOMAIN, {}) hass.data.setdefault(DOMAIN, {})
# 获取配置 _update_entry_data(hass, entry)
credential_path = entry.data.get(CONF_CREDENTIAL_PATH, DEFAULT_CREDENTIAL_PATH)
enable_punctuation = entry.data.get(CONF_ENABLE_PUNCTUATION, DEFAULT_ENABLE_PUNCTUATION)
# 将相对路径转换为绝对路径 entry.async_on_unload(
if not Path(credential_path).is_absolute(): entry.add_update_listener(_async_update_entry_listener)
credential_path = hass.config.path(credential_path) )
# 存储配置供 STT 实体使用
hass.data[DOMAIN][entry.entry_id] = {
CONF_CREDENTIAL_PATH: credential_path,
CONF_ENABLE_PUNCTUATION: enable_punctuation,
}
# 加载 STT 平台
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
_LOGGER.info("Doubao STT 集成已设置完成") _LOGGER.info("Doubao STT 集成已设置完成")
@@ -47,6 +38,29 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return True return True
def _update_entry_data(hass: HomeAssistant, entry: ConfigEntry) -> None:
"""Store config from entry.options (falling back to entry.data)."""
merged = {**entry.data, **entry.options}
credential_path = merged.get(CONF_CREDENTIAL_PATH, DEFAULT_CREDENTIAL_PATH)
enable_punctuation = merged.get(CONF_ENABLE_PUNCTUATION, DEFAULT_ENABLE_PUNCTUATION)
if not Path(credential_path).is_absolute():
credential_path = hass.config.path(credential_path)
hass.data[DOMAIN][entry.entry_id] = {
CONF_CREDENTIAL_PATH: credential_path,
CONF_ENABLE_PUNCTUATION: enable_punctuation,
}
async def _async_update_entry_listener(
hass: HomeAssistant, entry: ConfigEntry,
) -> None:
"""Handle options update."""
_update_entry_data(hass, entry)
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry.""" """Unload a config entry."""
if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
@@ -26,38 +26,27 @@ _LOGGER = logging.getLogger(__name__)
async def validate_input(hass: HomeAssistant, data: dict[str, Any]) -> dict[str, Any]: async def validate_input(hass: HomeAssistant, data: dict[str, Any]) -> dict[str, Any]:
"""Validate the user input allows us to connect. """Validate the user input allows us to connect.
Data has the keys from STEP_USER_DATA_SCHEMA with values provided by the user. Only checks that the credential file exists and is valid JSON if provided.
Does NOT register a device or perform network I/O.
Actual credential initialization is deferred to the first transcribe call.
""" """
from .doubaoime_asr import ASRConfig, DoubaoASR, ASRError import json as json_module
credential_path = data.get(CONF_CREDENTIAL_PATH, DEFAULT_CREDENTIAL_PATH) credential_path = data.get(CONF_CREDENTIAL_PATH, DEFAULT_CREDENTIAL_PATH)
enable_punctuation = data.get(CONF_ENABLE_PUNCTUATION, DEFAULT_ENABLE_PUNCTUATION) enable_punctuation = data.get(CONF_ENABLE_PUNCTUATION, DEFAULT_ENABLE_PUNCTUATION)
# 将相对路径转换为 Home Assistant 配置目录下的绝对路径
if not Path(credential_path).is_absolute(): if not Path(credential_path).is_absolute():
credential_path = hass.config.path(credential_path) credential_path = hass.config.path(credential_path)
# 尝试初始化配置并验证凭据 cred_path = Path(credential_path)
try: if cred_path.exists():
config = ASRConfig( try:
credential_path=credential_path, cred_data = json_module.loads(cred_path.read_text(encoding="utf-8"))
enable_punctuation=enable_punctuation, if not isinstance(cred_data, dict):
) raise CannotConnect("凭据文件格式不正确")
# 确保凭据已初始化(会自动注册设备如果需要) except (json_module.JSONDecodeError, OSError) as err:
await config.async_ensure_credentials() raise CannotConnect(f"凭据文件读取失败: {err}") from err
# 简单验证:检查是否成功获取了 device_id 和 token
if not config.device_id or not config.token:
raise CannotConnect("无法获取设备凭据")
except ASRError as err:
_LOGGER.error("验证 Doubao STT 配置失败: %s", err)
raise CannotConnect(str(err)) from err
except Exception as err:
_LOGGER.exception("验证 Doubao STT 配置时发生未知错误")
raise CannotConnect(str(err)) from err
# 返回用户可读的标题信息
return { return {
"title": "Doubao STT", "title": "Doubao STT",
"credential_path": credential_path, "credential_path": credential_path,
@@ -128,7 +117,6 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
errors: dict[str, str] = {} errors: dict[str, str] = {}
if user_input is not None: if user_input is not None:
# 验证新的配置
try: try:
await validate_input(self.hass, user_input) await validate_input(self.hass, user_input)
except CannotConnect: except CannotConnect:
@@ -137,6 +125,9 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
_LOGGER.exception("Unexpected exception") _LOGGER.exception("Unexpected exception")
errors["base"] = "unknown" errors["base"] = "unknown"
else: else:
self.hass.config_entries.async_update_entry(
self.config_entry, data={**self.config_entry.data, **user_input}
)
return self.async_create_entry(title="", data=user_input) return self.async_create_entry(title="", data=user_input)
# 获取当前配置值 # 获取当前配置值
@@ -1,23 +1,41 @@
from .asr import ( from .client import (
DoubaoASR, DoubaoASR,
ASRResponse,
ASRError,
ResponseType,
AudioChunk,
transcribe, transcribe,
transcribe_stream, transcribe_stream,
transcribe_realtime, transcribe_realtime,
) )
from .models import (
ASRResponse,
ASRResult,
ASRAlternative,
ASRWord,
ASRExtra,
OIDecodingInfo,
ASRError,
ResponseType,
AudioChunk,
)
from .config import ASRConfig from .config import ASRConfig
from .ner import NerResponse, NerResult, NerWord, ner, async_ner
__all__ = [ __all__ = [
"DoubaoASR", "DoubaoASR",
"ASRConfig", "ASRConfig",
"ASRResponse", "ASRResponse",
"ASRResult",
"ASRAlternative",
"ASRWord",
"ASRExtra",
"OIDecodingInfo",
"ASRError", "ASRError",
"ResponseType", "ResponseType",
"AudioChunk", "AudioChunk",
"transcribe", "transcribe",
"transcribe_stream", "transcribe_stream",
"transcribe_realtime", "transcribe_realtime",
"NerResponse",
"NerResult",
"NerWord",
"ner",
"async_ner",
] ]
@@ -1,656 +1,33 @@
from __future__ import annotations from .client import (
DoubaoASR,
import asyncio transcribe,
import contextlib transcribe_stream,
from dataclasses import dataclass transcribe_realtime,
from enum import Enum, auto )
import json from .models import (
from pathlib import Path ASRResponse,
import ssl ASRResult,
import time ASRAlternative,
from typing import AsyncIterator, Callable, List, Optional, Union ASRWord,
import uuid ASRExtra,
from pydantic import BaseModel, Field OIDecodingInfo,
import websockets ASRError,
from websockets import ClientConnection ResponseType,
AudioChunk,
from .config import ASRConfig, SessionConfig )
from .audio import AudioEncoder
from .asr_pb2 import AsrRequest, AsrResponse as AsrResponsePb, FrameState __all__ = [
"DoubaoASR",
# PCM 音频数据的类型别名 "ASRResponse",
AudioChunk = bytes "ASRResult",
"ASRAlternative",
"ASRWord",
class ResponseType(Enum): "ASRExtra",
""" "OIDecodingInfo",
ASR 响应类型 "ASRError",
""" "ResponseType",
TASK_STARTED = auto() "AudioChunk",
SESSION_STARTED = auto() "transcribe",
SESSION_FINISHED = auto() "transcribe_stream",
VAD_START = auto() "transcribe_realtime",
INTERIM_RESULT = auto() ]
FINAL_RESULT = auto()
HEARTBEAT = auto()
ERROR = auto()
UNKNOWN = auto()
@dataclass
class ASRResponse:
"""
ASR 响应
"""
type: ResponseType
text: str = ""
is_final: bool = False
vad_start: bool = False
vad_finished: bool = False
packet_number: int = -1
error_msg: str = ""
raw_json: Optional[dict] = None
class ASRError(Exception):
"""
ASR 错误
"""
def __init__(self, message: str, response: Optional[ASRResponse] = None) -> None:
super().__init__(message)
self.response = response
class _SessionState(BaseModel):
"""
ASR 会话状态
"""
request_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
final_text: str = ""
is_finished: bool = False
error: Optional[ASRResponse] = None
class DoubaoASR:
"""
豆包输入法 ASR 客户端
"""
def __init__(self, config: Optional[ASRConfig] = None):
self.config = config
self._encoder = AudioEncoder(self.config)
self._ssl_context: Optional[ssl.SSLContext] = None
async def __aenter__(self) -> DoubaoASR:
return self
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
pass
async def _ensure_ssl_context(self):
"""
确保 SSL 上下文已创建(在 executor 中运行避免阻塞)
"""
if self._ssl_context is None:
loop = asyncio.get_event_loop()
self._ssl_context = await loop.run_in_executor(None, ssl.create_default_context)
async def transcribe(self, audio: Union[str, Path, bytes], *, realtime = False, on_interim: Callable[[str], None] = None) -> str:
"""
非流式语音识别
:param audio: 音频文件路径或 PCM 字节数据
:param on_interim: 可选的中间结果回调
:return: 最终识别文本
"""
final_text = ""
async for response in self.transcribe_stream(audio, realtime=realtime):
if response.type == ResponseType.INTERIM_RESULT and on_interim:
on_interim(response.text)
elif response.type == ResponseType.FINAL_RESULT:
final_text = response.text
elif response.type == ResponseType.ERROR:
raise ASRError(response.error_msg, response)
return final_text
async def transcribe_stream(self, audio: Union[str, Path, bytes], *, realtime: bool = False) -> AsyncIterator[ASRResponse]:
"""
流式语音识别(完整音频)
:param audio: 音频文件路径或 PCM 字节数据
:param realtime: 是否按实时速度发送
:return: ASR 响应流,包括中间结果和最终结果
"""
if isinstance(audio, (str, Path)):
pcm_data = self._encoder.convert_audio_to_pcm(
audio, self.config.sample_rate, self.config.channels,
)
else:
pcm_data = audio
opus_frames = await self._encoder.pcm_to_opus_frames(pcm_data)
state = _SessionState()
# 预创建 SSL 上下文避免阻塞
await self._ensure_ssl_context()
# 获取 WebSocket URL
ws_url = await self.config.get_ws_url()
try:
async with websockets.connect(
ws_url,
additional_headers=self.config.headers,
open_timeout=self.config.connect_timeout,
ssl=self._ssl_context,
) as ws:
# 初始化会话
async for resp in self._initialize_session(ws, state):
yield resp
# 响应队列
response_queue: asyncio.Queue[Optional[ASRResponse]] = asyncio.Queue()
# 启动发送和接收任务
send_task = asyncio.create_task(
self._send_audio(ws, opus_frames, state, realtime)
)
recv_task = asyncio.create_task(
self._receive_responses(ws, state, response_queue)
)
try:
# 从队列中获取服务器响应
while True:
try:
resp = await asyncio.wait_for(
response_queue.get(),
timeout=self.config.recv_timeout,
)
if resp is None: # 结束标记
break
# 心跳包仅用于重置超时,不转发给用户
if resp.type == ResponseType.HEARTBEAT:
continue
yield resp
if resp.type == ResponseType.ERROR:
break
except asyncio.TimeoutError:
break
await send_task
finally:
send_task.cancel()
recv_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await send_task
with contextlib.suppress(asyncio.CancelledError):
await recv_task
except websockets.exceptions.WebSocketException as e:
raise ASRError(f"WebSocket 错误: {e}") from e
async def transcribe_realtime(
self,
audio_source: AsyncIterator[AudioChunk],
) -> AsyncIterator[ASRResponse]:
"""
实时流式语音识别(支持麦克风等持续音频源)
:param audio_source: PCM 音频数据的异步迭代器
- 每个 chunk 应为 16-bit PCM 数据
- 采样率和声道数应与 config 中配置一致
- 迭代器结束时会自动发送 FinishSession
:return: ASR 响应流
"""
state = _SessionState()
# 预创建 SSL 上下文避免阻塞
await self._ensure_ssl_context()
# 获取 WebSocket URL
ws_url = await self.config.get_ws_url()
try:
async with websockets.connect(
ws_url,
additional_headers=self.config.headers,
open_timeout=self.config.connect_timeout,
ssl=self._ssl_context,
) as ws:
# 初始化会话
async for resp in self._initialize_session(ws, state):
yield resp
# 响应队列
response_queue: asyncio.Queue[Optional[ASRResponse]] = asyncio.Queue()
# 启动发送和接收任务
send_task = asyncio.create_task(
self._send_audio_realtime(ws, audio_source, state)
)
recv_task = asyncio.create_task(
self._receive_responses(ws, state, response_queue)
)
try:
# 实时模式不使用超时,依靠 WebSocket 层检测断开
while True:
resp = await response_queue.get()
if resp is None:
break
if resp.type == ResponseType.HEARTBEAT:
continue
yield resp
if resp.type == ResponseType.ERROR:
break
await send_task
finally:
send_task.cancel()
recv_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await send_task
with contextlib.suppress(asyncio.CancelledError):
await recv_task
except websockets.exceptions.WebSocketException as e:
raise ASRError(f"WebSocket 错误: {e}") from e
async def _send_audio_realtime(
self,
ws: ClientConnection,
audio_source: AsyncIterator[AudioChunk],
state: _SessionState,
):
"""
从异步迭代器读取 PCM 数据并实时发送
"""
# 预先获取编码器,避免在循环中触发阻塞导入
encoder = await self._encoder.get_encoder()
timestamp_ms = int(time.time() * 1000)
frame_index = 0
pcm_buffer = b""
samples_per_frame = (
self.config.sample_rate * self.config.frame_duration_ms // 1000
)
bytes_per_frame = samples_per_frame * 2 # 16-bit
async for chunk in audio_source:
if state.is_finished:
break
pcm_buffer += chunk
# 当缓冲区有足够数据时,编码并发送
while len(pcm_buffer) >= bytes_per_frame:
pcm_frame = pcm_buffer[:bytes_per_frame]
pcm_buffer = pcm_buffer[bytes_per_frame:]
# 编码为 Opus
opus_frame = encoder.encode(pcm_frame, samples_per_frame)
# 确定帧状态(实时模式下不知道最后一帧,使用 FIRST/MIDDLE
if frame_index == 0:
frame_state = FrameState.FRAME_STATE_FIRST
else:
frame_state = FrameState.FRAME_STATE_MIDDLE
msg = _build_asr_request(
opus_frame,
state.request_id,
frame_state,
timestamp_ms + frame_index * self.config.frame_duration_ms,
)
await ws.send(msg)
frame_index += 1
# 迭代器结束,处理剩余数据
if pcm_buffer and not state.is_finished:
# 补零到完整帧
if len(pcm_buffer) < bytes_per_frame:
pcm_buffer += b"\x00" * (bytes_per_frame - len(pcm_buffer))
opus_frame = encoder.encode(pcm_buffer, samples_per_frame)
msg = _build_asr_request(
opus_frame,
state.request_id,
FrameState.FRAME_STATE_LAST,
timestamp_ms + frame_index * self.config.frame_duration_ms,
)
await ws.send(msg)
elif frame_index > 0 and not state.is_finished:
# 没有剩余数据,但需要发送一个 LAST 帧标记
# 发送一个空的 LAST 帧(静音)
silent_frame = b"\x00" * bytes_per_frame
opus_frame = encoder.encode(silent_frame, samples_per_frame)
msg = _build_asr_request(
opus_frame,
state.request_id,
FrameState.FRAME_STATE_LAST,
timestamp_ms + frame_index * self.config.frame_duration_ms,
)
await ws.send(msg)
# FinishSession
if not state.is_finished:
token = await self.config.get_token()
await ws.send(_build_finish_session(state.request_id, token))
async def _initialize_session(self, ws: ClientConnection, state: _SessionState) -> AsyncIterator[ASRResponse]:
"""
初始化 ASR 会话
"""
token = await self.config.get_token()
# StartTask
await ws.send(_build_start_task(state.request_id, token))
resp = await ws.recv()
parsed = _parse_response(resp)
if parsed.type == ResponseType.ERROR:
raise ASRError(f'StartTask 失败:{parsed.error_msg}', parsed)
yield parsed
# StartSession
session_config = await self.config.get_session_config()
await ws.send(
_build_start_session(state.request_id, token, session_config)
)
resp = await ws.recv()
parsed = _parse_response(resp)
if parsed.type == ResponseType.ERROR:
raise ASRError(f'StartSession 失败:{parsed.error_msg}', parsed)
yield parsed
async def _send_audio(
self,
ws: ClientConnection,
opus_frames: List[bytes],
state: _SessionState,
realtime: bool,
):
"""
发送音频帧
"""
timestamp_ms = int(time.time() * 1000)
frame_interval = self.config.frame_duration_ms / 1000.0
for i, opus_frame in enumerate(opus_frames):
if state.is_finished:
break
if i == 0:
frame_state = FrameState.FRAME_STATE_FIRST
elif i == len(opus_frames) - 1:
frame_state = FrameState.FRAME_STATE_LAST
else:
frame_state = FrameState.FRAME_STATE_MIDDLE
msg = _build_asr_request(
opus_frame,
state.request_id,
frame_state,
timestamp_ms + i * self.config.frame_duration_ms,
)
await ws.send(msg)
if realtime:
await asyncio.sleep(frame_interval)
# FinishSession
token = await self.config.get_token()
await ws.send(_build_finish_session(state.request_id, token))
async def _receive_responses(
self,
ws: ClientConnection,
state: _SessionState,
queue: asyncio.Queue[Optional[ASRResponse]],
):
"""
接受响应并放入队列
"""
try:
while not state.is_finished:
response = await ws.recv()
resp = _parse_response(response)
if resp.type == ResponseType.ERROR:
state.error = resp
state.is_finished = True
await queue.put(resp)
break
elif resp.type == ResponseType.HEARTBEAT:
# 心跳包也放入队列,用于重置超时计时器
await queue.put(resp)
elif resp.type == ResponseType.SESSION_FINISHED:
state.is_finished = True
await queue.put(resp)
break
elif resp.type == ResponseType.FINAL_RESULT:
state.final_text = resp.text
await queue.put(resp)
else:
await queue.put(resp)
except websockets.exceptions.ConnectionClosed:
state.is_finished = True
finally:
# 结束信号
await queue.put(None)
def _build_start_task(request_id: str, token: str) -> bytes:
"""构建 StartTask 消息 pb 数据"""
request = AsrRequest()
request.token = token
request.service_name = "ASR"
request.method_name = "StartTask"
request.request_id = request_id
return request.SerializeToString()
def _build_start_session(request_id: str, token: str, config: SessionConfig) -> bytes:
"""构建 StartSession 消息 pb 数据"""
request = AsrRequest()
request.token = token
request.service_name = "ASR"
request.method_name = "StartSession"
request.request_id = request_id
request.payload = config.model_dump_json()
return request.SerializeToString()
def _build_finish_session(request_id: str, token: str) -> bytes:
"""构建 FinishSession 消息 pb 数据"""
request = AsrRequest()
request.token = token
request.service_name = "ASR"
request.method_name = "FinishSession"
request.request_id = request_id
return request.SerializeToString()
def _build_asr_request(
audio_data: bytes,
request_id: str,
frame_state: FrameState,
timestamp_ms: int,
) -> bytes:
request = AsrRequest()
metadata = json.dumps({"extra": {}, "timestamp_ms": timestamp_ms})
request.service_name = "ASR"
request.method_name = "TaskRequest"
request.payload = metadata
request.audio_data = audio_data
request.request_id = request_id
request.frame_state = frame_state
return request.SerializeToString()
def _parse_response(data: bytes) -> ASRResponse:
"""解析 ASR 响应 (使用 protobuf)"""
pb = AsrResponsePb()
pb.ParseFromString(data)
message_type = pb.message_type
result_json = pb.result_json # 字段 7: 识别结果 JSON
status_message = pb.status_message # 字段 6: 状态消息
# 根据 message_type 判断响应类型
if message_type == "TaskStarted":
return ASRResponse(type=ResponseType.TASK_STARTED)
if message_type == "SessionStarted":
return ASRResponse(type=ResponseType.SESSION_STARTED)
if message_type == "SessionFinished":
return ASRResponse(type=ResponseType.SESSION_FINISHED)
if message_type in ("TaskFailed", "SessionFailed"):
return ASRResponse(type=ResponseType.ERROR, error_msg=status_message)
# 识别结果在 result_json 字段(字段 7
if not result_json:
return ASRResponse(type=ResponseType.UNKNOWN)
try:
json_data = json.loads(result_json)
except json.JSONDecodeError:
return ASRResponse(type=ResponseType.UNKNOWN)
results = json_data.get("results")
extra = json_data.get("extra", {})
# 无 results,可能是心跳包
if results is None:
return ASRResponse(
type=ResponseType.HEARTBEAT,
packet_number=extra.get("packet_number", -1),
raw_json=json_data,
)
# VAD 开始
if extra.get("vad_start"):
return ASRResponse(type=ResponseType.VAD_START, vad_start=True, raw_json=json_data)
# 解析识别结果
text = ""
is_interim = True
vad_finished = False
nonstream_result = False
for r in results:
if r.get("text"):
text = r.get("text")
if r.get("is_interim") is False:
is_interim = False
if r.get("is_vad_finished"):
vad_finished = True
if r.get("extra", {}).get("nonstream_result"):
nonstream_result = True
# 最终结果
if nonstream_result or (not is_interim and vad_finished):
return ASRResponse(
type=ResponseType.FINAL_RESULT,
text=text,
is_final=True,
vad_finished=vad_finished,
raw_json=json_data,
)
# 中间结果
return ASRResponse(
type=ResponseType.INTERIM_RESULT,
text=text,
is_final=False,
raw_json=json_data,
)
# =============
# 便捷函数
# =============
async def transcribe(
audio: str | Path | bytes,
*,
config: ASRConfig | None = None,
on_interim: Callable[[str], None] | None = None,
realtime: bool = False,
) -> str:
"""
便捷函数:非流式语音识别
Args:
audio: 音频文件路径或 PCM 字节数据
config: ASR 配置(可选)
on_interim: 中间结果回调(可选)
realtime: 是否模拟实时语音输入
- True: 按音频实际时长发送,每帧间插入延迟,模拟实时的语音输入
- False(默认): 尽快发送所有帧,会更快拿到结果(不知道会不会被风控)
Returns:
最终识别文本
"""
async with DoubaoASR(config) as asr:
return await asr.transcribe(audio, on_interim=on_interim, realtime=realtime)
async def transcribe_stream(
audio: str | Path | bytes,
*,
config: ASRConfig | None = None,
realtime: bool = False,
) -> AsyncIterator[ASRResponse]:
"""
便捷函数:流式语音识别(完整音频)
Args:
audio: 音频文件路径或 PCM 字节数据
config: ASR 配置(可选)
realtime: 是否模拟实时语音输入
- True: 按音频实际时长发送,每帧间插入延迟,模拟实时的语音输入
- False(默认): 尽快发送所有帧,会更快拿到结果(不知道会不会被风控)
Yields:
ASRResponse 对象
"""
async with DoubaoASR(config) as asr:
async for resp in asr.transcribe_stream(audio, realtime=realtime):
yield resp
async def transcribe_realtime(
audio_source: AsyncIterator[AudioChunk],
*,
config: ASRConfig | None = None,
) -> AsyncIterator[ASRResponse]:
"""
便捷函数:实时流式语音识别(支持麦克风等持续音频源)
Args:
audio_source: PCM 音频数据的异步迭代器
- 每个 chunk 应为 16-bit PCM 数据
- 采样率和声道数应与 config 中配置一致
config: ASR 配置(可选)
Yields:
ASRResponse 对象
"""
async with DoubaoASR(config) as asr:
async for resp in asr.transcribe_realtime(audio_source):
yield resp
@@ -1,13 +1,11 @@
from typing import Optional, List, Union, TYPE_CHECKING from typing import Optional, List, Union
from pathlib import Path from pathlib import Path
import asyncio
import opuslib
import miniaudio import miniaudio
from .config import ASRConfig from .config import ASRConfig
if TYPE_CHECKING:
import opuslib
class AudioEncoder: class AudioEncoder:
@@ -16,36 +14,19 @@ class AudioEncoder:
""" """
def __init__(self, config: ASRConfig) -> None: def __init__(self, config: ASRConfig) -> None:
self.config = config self.config = config
self._encoder: Optional["opuslib.Encoder"] = None self._encoder: Optional[opuslib.Encoder] = None
self._opuslib_module = None
async def _ensure_opuslib(self): @property
"""在 executor 中导入 opuslib,避免阻塞事件循环""" def encoder(self) -> opuslib.Encoder:
if self._opuslib_module is None:
loop = asyncio.get_event_loop()
self._opuslib_module = await loop.run_in_executor(None, self._import_opuslib)
@staticmethod
def _import_opuslib():
"""在独立线程中导入 opuslib"""
import opuslib
return opuslib
async def get_encoder(self) -> "opuslib.Encoder":
"""异步获取编码器"""
if self._encoder is None: if self._encoder is None:
await self._ensure_opuslib() self._encoder = opuslib.Encoder(
self._encoder = self._opuslib_module.Encoder(
self.config.sample_rate, self.config.sample_rate,
self.config.channels, self.config.channels,
self._opuslib_module.APPLICATION_AUDIO, opuslib.APPLICATION_AUDIO,
) )
return self._encoder return self._encoder
async def pcm_to_opus_frames(self, pcm_data: bytes) -> List[bytes]: def pcm_to_opus_frames(self, pcm_data: bytes) -> List[bytes]:
"""将 PCM 数据转换为 Opus 帧(异步)"""
encoder = await self.get_encoder()
samples_per_frame = ( samples_per_frame = (
self.config.sample_rate * self.config.frame_duration_ms // 1000 self.config.sample_rate * self.config.frame_duration_ms // 1000
) )
@@ -57,7 +38,7 @@ class AudioEncoder:
if len(chunk) < bytes_per_frame: if len(chunk) < bytes_per_frame:
chunk = chunk + b"\x00" * (bytes_per_frame - len(chunk)) chunk = chunk + b"\x00" * (bytes_per_frame - len(chunk))
opus_frame = encoder.encode(chunk, samples_per_frame) opus_frame = self.encoder.encode(chunk, samples_per_frame)
frames.append(opus_frame) frames.append(opus_frame)
return frames return frames
@@ -0,0 +1,340 @@
from __future__ import annotations
import asyncio
import logging
from pathlib import Path
import ssl
import time
from typing import AsyncIterator, Callable, List, Optional, Union
import websockets
from websockets import ClientConnection
from .config import ASRConfig
from .audio import AudioEncoder
from .models import (
ASRResponse,
ASRError,
AudioChunk,
ResponseType,
_SessionState,
)
from .parser import parse_response as _parse_response
from .protocol import (
build_start_task as _build_start_task,
build_start_session as _build_start_session,
build_finish_session as _build_finish_session,
build_asr_request as _build_asr_request,
)
from .asr_pb2 import FrameState
_LOGGER = logging.getLogger(__name__)
class DoubaoASR:
def __init__(self, config: Optional[ASRConfig] = None):
self.config = config
self._encoder = AudioEncoder(self.config)
self._ssl_context: Optional[ssl.SSLContext] = None
async def __aenter__(self) -> DoubaoASR:
self._ssl_context = await asyncio.to_thread(ssl.create_default_context)
return self
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
pass
async def _run_ws_session(
self,
send_coroutine,
recv_timeout: Optional[float] = None,
) -> AsyncIterator[ASRResponse]:
state = _SessionState()
ws_url = await self.config.async_ws_url()
try:
async with websockets.connect(
ws_url,
additional_headers=self.config.headers,
open_timeout=self.config.connect_timeout,
ssl=self._ssl_context,
) as ws:
async for resp in self._initialize_session(ws, state):
yield resp
response_queue: asyncio.Queue[Optional[ASRResponse]] = asyncio.Queue()
async def _send():
await send_coroutine(ws, state)
async def _recv():
await self._receive_responses(ws, state, response_queue)
send_task = asyncio.create_task(_send())
recv_task = asyncio.create_task(_recv())
try:
while True:
try:
if recv_timeout is not None:
resp = await asyncio.wait_for(
response_queue.get(), timeout=recv_timeout,
)
else:
resp = await response_queue.get()
if resp is None:
break
if resp.type == ResponseType.HEARTBEAT:
continue
yield resp
if resp.type == ResponseType.ERROR:
break
except asyncio.TimeoutError:
break
finally:
results = await asyncio.gather(
send_task, recv_task, return_exceptions=True,
)
for r in results:
if isinstance(r, BaseException) and not isinstance(r, asyncio.CancelledError):
_LOGGER.debug("WS task exception: %s", r)
except websockets.exceptions.WebSocketException as e:
raise ASRError(f"WebSocket 错误: {e}") from e
async def transcribe(self, audio: Union[str, Path, bytes], *, realtime=False, on_interim: Callable[[str], None] = None) -> str:
final_text = ""
async for response in self.transcribe_stream(audio, realtime=realtime):
if response.type == ResponseType.INTERIM_RESULT and on_interim:
on_interim(response.text)
elif response.type == ResponseType.FINAL_RESULT:
final_text = response.text
elif response.type == ResponseType.ERROR:
raise ASRError(response.error_msg, response)
return final_text
async def transcribe_stream(self, audio: Union[str, Path, bytes], *, realtime: bool = False) -> AsyncIterator[ASRResponse]:
if isinstance(audio, (str, Path)):
pcm_data = self._encoder.convert_audio_to_pcm(
audio, self.config.sample_rate, self.config.channels,
)
else:
pcm_data = audio
opus_frames = self._encoder.pcm_to_opus_frames(pcm_data)
async def _send(ws, state):
await self._send_audio(ws, opus_frames, state, realtime)
async for resp in self._run_ws_session(
_send, recv_timeout=self.config.recv_timeout,
):
yield resp
async def transcribe_realtime(
self,
audio_source: AsyncIterator[AudioChunk],
) -> AsyncIterator[ASRResponse]:
async def _send(ws, state):
await self._send_audio_realtime(ws, audio_source, state)
async for resp in self._run_ws_session(_send, recv_timeout=None):
yield resp
async def _send_audio_realtime(
self,
ws: ClientConnection,
audio_source: AsyncIterator[AudioChunk],
state: _SessionState,
):
timestamp_ms = int(time.time() * 1000)
frame_index = 0
pcm_buffer = b""
samples_per_frame = (
self.config.sample_rate * self.config.frame_duration_ms // 1000
)
bytes_per_frame = samples_per_frame * 2
async for chunk in audio_source:
if state.is_finished:
break
pcm_buffer += chunk
while len(pcm_buffer) >= bytes_per_frame:
pcm_frame = pcm_buffer[:bytes_per_frame]
pcm_buffer = pcm_buffer[bytes_per_frame:]
opus_frame = self._encoder.encoder.encode(pcm_frame, samples_per_frame)
if frame_index == 0:
frame_state = FrameState.FRAME_STATE_FIRST
else:
frame_state = FrameState.FRAME_STATE_MIDDLE
msg = _build_asr_request(
opus_frame,
state.request_id,
frame_state,
timestamp_ms + frame_index * self.config.frame_duration_ms,
)
await ws.send(msg)
frame_index += 1
if pcm_buffer and not state.is_finished:
if len(pcm_buffer) < bytes_per_frame:
pcm_buffer += b"\x00" * (bytes_per_frame - len(pcm_buffer))
opus_frame = self._encoder.encoder.encode(pcm_buffer, samples_per_frame)
msg = _build_asr_request(
opus_frame,
state.request_id,
FrameState.FRAME_STATE_LAST,
timestamp_ms + frame_index * self.config.frame_duration_ms,
)
await ws.send(msg)
elif frame_index > 0 and not state.is_finished:
silent_frame = b"\x00" * bytes_per_frame
opus_frame = self._encoder.encoder.encode(silent_frame, samples_per_frame)
msg = _build_asr_request(
opus_frame,
state.request_id,
FrameState.FRAME_STATE_LAST,
timestamp_ms + frame_index * self.config.frame_duration_ms,
)
await ws.send(msg)
if not state.is_finished:
token = await self.config.async_get_token()
await ws.send(_build_finish_session(state.request_id, token))
async def _initialize_session(self, ws: ClientConnection, state: _SessionState) -> AsyncIterator[ASRResponse]:
token = await self.config.async_get_token()
await ws.send(_build_start_task(state.request_id, token))
resp = await ws.recv()
parsed = _parse_response(resp)
if parsed.type == ResponseType.ERROR:
raise ASRError(f'StartTask 失败:{parsed.error_msg}', parsed)
yield parsed
session_cfg = await self.config.async_session_config()
await ws.send(
_build_start_session(state.request_id, token, session_cfg)
)
resp = await ws.recv()
parsed = _parse_response(resp)
if parsed.type == ResponseType.ERROR:
raise ASRError(f'StartSession 失败:{parsed.error_msg}', parsed)
yield parsed
async def _send_audio(
self,
ws: ClientConnection,
opus_frames: List[bytes],
state: _SessionState,
realtime: bool,
):
timestamp_ms = int(time.time() * 1000)
frame_interval = self.config.frame_duration_ms / 1000.0
for i, opus_frame in enumerate(opus_frames):
if state.is_finished:
break
if i == 0:
frame_state = FrameState.FRAME_STATE_FIRST
elif i == len(opus_frames) - 1:
frame_state = FrameState.FRAME_STATE_LAST
else:
frame_state = FrameState.FRAME_STATE_MIDDLE
msg = _build_asr_request(
opus_frame,
state.request_id,
frame_state,
timestamp_ms + i * self.config.frame_duration_ms,
)
await ws.send(msg)
if realtime:
await asyncio.sleep(frame_interval)
token = await self.config.async_get_token()
await ws.send(_build_finish_session(state.request_id, token))
async def _receive_responses(
self,
ws: ClientConnection,
state: _SessionState,
queue: asyncio.Queue[Optional[ASRResponse]],
):
try:
while not state.is_finished:
response = await ws.recv()
resp = _parse_response(response)
if resp.type == ResponseType.ERROR:
state.error = resp
state.is_finished = True
await queue.put(resp)
break
elif resp.type == ResponseType.HEARTBEAT:
await queue.put(resp)
elif resp.type == ResponseType.SESSION_FINISHED:
state.is_finished = True
await queue.put(resp)
break
elif resp.type == ResponseType.FINAL_RESULT:
state.final_text = resp.text
await queue.put(resp)
else:
await queue.put(resp)
except websockets.exceptions.ConnectionClosed as exc:
state.is_finished = True
await queue.put(ASRResponse(
type=ResponseType.ERROR,
error_msg=f"WebSocket 连接已关闭: {exc}",
))
finally:
await queue.put(None)
async def transcribe(
audio: str | Path | bytes,
*,
config: ASRConfig | None = None,
on_interim: Callable[[str], None] | None = None,
realtime: bool = False,
) -> str:
async with DoubaoASR(config) as asr:
return await asr.transcribe(audio, on_interim=on_interim, realtime=realtime)
async def transcribe_stream(
audio: str | Path | bytes,
*,
config: ASRConfig | None = None,
realtime: bool = False,
) -> AsyncIterator[ASRResponse]:
async with DoubaoASR(config) as asr:
async for resp in asr.transcribe_stream(audio, realtime=realtime):
yield resp
async def transcribe_realtime(
audio_source: AsyncIterator[AudioChunk],
*,
config: ASRConfig | None = None,
) -> AsyncIterator[ASRResponse]:
async with DoubaoASR(config) as asr:
async for resp in asr.transcribe_realtime(audio_source):
yield resp
@@ -1,13 +1,13 @@
from dataclasses import dataclass, field
import asyncio import asyncio
import json from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import Optional, Union from typing import Optional, Union
from pydantic import BaseModel from pydantic import BaseModel
import aiofiles
from .constants import WEBSOCKET_URL, USER_AGENT, AID from .constants import WEBSOCKET_URL, USER_AGENT, AID
from .device import DeviceCredentials, register_device, get_asr_token from .credential import CredentialManager
from .device import DeviceCredentials
class _AudioInfo(BaseModel): class _AudioInfo(BaseModel):
@@ -26,9 +26,6 @@ class _SessionExtraConfig(BaseModel):
class SessionConfig(BaseModel): class SessionConfig(BaseModel):
"""
ASR 任务开始前需要初始化 Session 的配置
"""
audio_info: _AudioInfo audio_info: _AudioInfo
enable_punctuation: bool enable_punctuation: bool
enable_speech_rejection: bool enable_speech_rejection: bool
@@ -37,162 +34,61 @@ class SessionConfig(BaseModel):
@dataclass @dataclass
class ASRConfig: class ASRConfig:
"""
ASR 配置
如果不提供 device_id 和 token,将自动注册设备并获取 token。
示例:
# 自动获取凭据(首次使用时会注册设备,不持久化)
config = ASRConfig()
# 使用已有凭据
config = ASRConfig(device_id="1234567890123456", token="MyToken123")
# 使用凭据文件(推荐,首次注册后自动缓存)
config = ASRConfig(credential_path="~/.config/doubao-asr/credentials.json")
# 凭据文件 + 覆盖部分参数
config = ASRConfig(
credential_path="~/.config/doubao-asr/credentials.json",
token="NewToken", # 覆盖文件中的 token
)
"""
url: str = WEBSOCKET_URL url: str = WEBSOCKET_URL
aid: str = AID aid: str = AID
user_agent: str = USER_AGENT user_agent: str = USER_AGENT
device_id: Optional[str] = None # 空则自动获取 device_id: Optional[str] = None
token: Optional[str] = None # 空则自动获取 token: Optional[str] = None
credential_path: Union[str, Path, None] = None credential_path: Union[str, Path, None] = None
"""
凭据文件路径
"""
# 这些都是客户端给的默认值,挺通用的。其实我也没尝试过改了服务器会不会认
# 音频配置
sample_rate: int = 16000 sample_rate: int = 16000
channels: int = 1 channels: int = 1
frame_duration_ms: int = 20 frame_duration_ms: int = 20
# 会话配置
enable_punctuation: bool = True enable_punctuation: bool = True
enable_speech_rejection: bool = False enable_speech_rejection: bool = False
enable_asr_twopass: bool = True enable_asr_twopass: bool = True
enable_asr_threepass: bool = True enable_asr_threepass: bool = True
# 这里是输入法当前作用在哪个应用上
# 可能服务器会根据当前所使用的程序调整不同的语音识别策略??
# 这里用 Chrome 浏览器,算是比较通用的了吧?
app_name: str = "com.android.chrome" app_name: str = "com.android.chrome"
# 连接配置
connect_timeout: float = 10.0 connect_timeout: float = 10.0
recv_timeout: float = 10.0 recv_timeout: float = 10.0
# 内部状态 _credential_manager: Optional[CredentialManager] = field(default=None, repr=False)
_credentials: Optional[DeviceCredentials] = field(default=None, repr=None) _wave_client: Optional[object] = field(default=None, repr=False)
_initialized: bool = field(default=False, repr=False)
async def _load_credentials_from_file(self) -> Optional[DeviceCredentials]: @property
""" def _cred_mgr(self) -> CredentialManager:
从缓存文件中加载凭据信息(异步) if self._credential_manager is None:
""" self._credential_manager = CredentialManager(
if self.credential_path is None: credential_path=self.credential_path,
return None device_id=self.device_id,
token=self.token,
)
return self._credential_manager
path = Path(self.credential_path).expanduser() def _sync_from_mgr(self) -> None:
if not path.exists(): self.device_id = self._cred_mgr.device_id
return None self.token = self._cred_mgr.token
try:
async with aiofiles.open(path, 'r', encoding='utf-8') as f:
content = await f.read()
data = json.loads(content)
return DeviceCredentials(**data)
except (json.JSONDecodeError, OSError): def ensure_credentials(self):
return None self._cred_mgr.ensure_sync()
self._sync_from_mgr()
async def _save_credentials_to_file(self, creds: DeviceCredentials):
"""
保存凭据至缓存文件(异步)
"""
if self.credential_path is None:
return
path = Path(self.credential_path).expanduser()
# 使用 run_in_executor 来执行同步的 mkdir 操作
loop = asyncio.get_event_loop()
await loop.run_in_executor(None, lambda: path.parent.mkdir(parents=True, exist_ok=True))
async with aiofiles.open(path, 'w', encoding='utf-8') as f:
content = json.dumps(creds.model_dump(), indent=2, ensure_ascii=False)
await f.write(content)
async def async_ensure_credentials(self): async def async_ensure_credentials(self):
""" await self._cred_mgr.ensure_async()
确保凭据已初始化(异步) self._sync_from_mgr()
优先级: @property
1. 直接传入的 device_id/token 参数(最高优先级) def ws_url(self) -> str:
2. credential_path 文件中的值 self.ensure_credentials()
3. 自动注册获取(最低优先级) return f'{self.url}?aid={self.aid}&device_id={self.device_id}'
如果指定了 credential_path 且文件不存在,会注册设备并保存到文件。 async def async_ws_url(self) -> str:
"""
if self._initialized:
return
# 保存直接通过参数传入的凭据,用于进行覆盖
user_device_id = self.device_id
user_token = self.token
# 尝试从文件中加载凭据
file_creds = await self._load_credentials_from_file()
if file_creds:
self._credentials = file_creds
# 使用文件中的值作为默认
if self.device_id is None:
self.device_id = file_creds.device_id
if self.token is None:
self.token = file_creds.token
# 如果 device_id 仍为 None, 则注册设备
need_save = False
if self.device_id is None:
# 在 executor 中运行同步的 register_device
loop = asyncio.get_event_loop()
self._credentials = await loop.run_in_executor(None, register_device)
self.device_id = self._credentials.device_id
need_save = True
# 如果 token 仍为 None, 则获取 token
if self.token is None:
cdid = self._credentials.cdid if self._credentials else None
# 在 executor 中运行同步的 get_asr_token
loop = asyncio.get_event_loop()
self.token = await loop.run_in_executor(None, get_asr_token, self.device_id, cdid)
# 如果指定了 credential_path 且有新注册的凭据,则保存至文件
if self.credential_path and need_save and self._credentials:
self._credentials.token = self.token
await self._save_credentials_to_file(self._credentials)
# 覆盖用户传入的参数
if user_device_id is not None:
self.device_id = user_device_id
if user_token is not None:
self.token = user_token
self._initialized = True
async def get_ws_url(self) -> str:
"""获取 WebSocket URL(异步)"""
await self.async_ensure_credentials() await self.async_ensure_credentials()
return f'{self.url}?aid={self.aid}&device_id={self.device_id}' return f'{self.url}?aid={self.aid}&device_id={self.device_id}'
@property @property
def headers(self) -> dict[str, str]: def headers(self) -> dict[str, str]:
return { return {
@@ -201,9 +97,15 @@ class ASRConfig:
"x-custom-keepalive": "true" "x-custom-keepalive": "true"
} }
async def get_session_config(self) -> SessionConfig: def session_config(self) -> SessionConfig:
"""获取会话配置(异步)""" self.ensure_credentials()
return self._build_session_config()
async def async_session_config(self) -> SessionConfig:
await self.async_ensure_credentials() await self.async_ensure_credentials()
return self._build_session_config()
def _build_session_config(self) -> SessionConfig:
audio_info = _AudioInfo( audio_info = _AudioInfo(
channel=self.channels, channel=self.channels,
format="speech_opus", format="speech_opus",
@@ -224,8 +126,71 @@ class ASRConfig:
enable_speech_rejection=self.enable_speech_rejection, enable_speech_rejection=self.enable_speech_rejection,
extra=extra, extra=extra,
) )
async def get_token(self) -> str: def get_token(self) -> str:
"""获取 token(异步)""" self.ensure_credentials()
return self.token
async def async_get_token(self) -> str:
await self.async_ensure_credentials() await self.async_ensure_credentials()
return self.token return self.token
def _on_wave_session_update(self, session) -> None:
mgr = self._cred_mgr
if mgr.credentials:
mgr.credentials.wave_session = session.to_dict()
mgr.save_credentials_sync(mgr.credentials)
def get_wave_client(self):
from .wave_client import WaveClient, WaveSession
self.ensure_credentials()
if self._wave_client is None:
cached_session = None
creds = self._cred_mgr.credentials
if creds and creds.wave_session:
try:
session = WaveSession.from_dict(creds.wave_session)
if not session.is_expired():
cached_session = session
except (KeyError, ValueError):
pass
self._wave_client = WaveClient(
self.device_id,
self.aid,
session=cached_session,
on_session_update=self._on_wave_session_update,
)
return self._wave_client
async def async_get_wave_client(self):
from .wave_client import WaveClient, WaveSession
await self.async_ensure_credentials()
if self._wave_client is None:
cached_session = None
creds = self._cred_mgr.credentials
if creds and creds.wave_session:
try:
session = WaveSession.from_dict(creds.wave_session)
if not session.is_expired():
cached_session = session
except (KeyError, ValueError):
pass
self._wave_client = WaveClient(
self.device_id,
self.aid,
session=cached_session,
on_session_update=self._on_wave_session_update,
)
return self._wave_client
def get_sami_token(self) -> str:
self.ensure_credentials()
return self._cred_mgr.get_sami_token_sync()
async def async_get_sami_token(self) -> str:
await self.async_ensure_credentials()
return await self._cred_mgr.get_sami_token_async()
@@ -7,9 +7,24 @@ SETTINGS_URL = "https://is.snssdk.com/service/settings/v3/"
# ASR WebSocket URL # ASR WebSocket URL
WEBSOCKET_URL = "wss://frontier-audio-ime-ws.doubao.com/ocean/api/v1/ws" WEBSOCKET_URL = "wss://frontier-audio-ime-ws.doubao.com/ocean/api/v1/ws"
# 需要通过该接口获取 SAMI TOKEN 来访问 NER 接口
SAMI_CONFIG_URL = "https://ime.oceancloudapi.com/api/v1/user/get_config"
# 加密协议握手 URL
HANDSHAKE_URL = "https://keyhub.zijieapi.com/handshake"
# NER 接口 URL
NER_URL = "https://speech.bytedance.com/api/v3/context/ime/ner"
# 豆包输入法的 APP ID # 豆包输入法的 APP ID
AID = 401734 AID = 401734
# SAMI 语音相关的服务 APP KEY (写死在客户端的)
SAMI_APP_KEY = "SYlxZr6LnvBaIVmF"
# HKDF info 字符串 (写死在客户端的)
HKDF_INFO = b"4e30514609050cd3"
# 应用配置 (豆包输入法) # 应用配置 (豆包输入法)
APP_CONFIG = { APP_CONFIG = {
"aid": AID, "aid": AID,
@@ -0,0 +1,196 @@
from __future__ import annotations
import asyncio
import json
from pathlib import Path
from typing import Optional
from .device import DeviceCredentials, register_device, get_asr_token
from .sami import get_sami_token
def _jwt_is_expired(token: str, margin: int = 60) -> bool:
import base64
import time
try:
payload_b64 = token.split(".")[1]
padding = 4 - len(payload_b64) % 4
if padding != 4:
payload_b64 += "=" * padding
payload = json.loads(base64.urlsafe_b64decode(payload_b64))
exp = payload.get("exp")
if exp is None:
return False
return time.time() >= exp - margin
except (IndexError, ValueError, json.JSONDecodeError):
return False
class CredentialManager:
"""Manages device credentials: loading, saving, registration, and token refresh."""
def __init__(
self,
credential_path: Optional[str | Path] = None,
device_id: Optional[str] = None,
token: Optional[str] = None,
) -> None:
self.credential_path = credential_path
self.device_id = device_id
self.token = token
self._credentials: Optional[DeviceCredentials] = None
self._initialized = False
def _load_sync(self) -> Optional[DeviceCredentials]:
if self.credential_path is None:
return None
path = Path(self.credential_path).expanduser()
if not path.exists():
return None
try:
with open(path, "r", encoding="utf-8") as f:
data = json.loads(f.read())
return DeviceCredentials(**data)
except (json.JSONDecodeError, OSError):
return None
def _save_sync(self, creds: DeviceCredentials) -> None:
if self.credential_path is None:
return
path = Path(self.credential_path).expanduser()
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
json.dump(creds.model_dump(), f, indent=2, ensure_ascii=False)
async def _load_async(self) -> Optional[DeviceCredentials]:
return await asyncio.to_thread(self._load_sync)
async def _save_async(self, creds: DeviceCredentials) -> None:
await asyncio.to_thread(self._save_sync, creds)
@property
def credentials(self) -> Optional[DeviceCredentials]:
return self._credentials
def ensure(self, *, is_async: bool = False) -> None:
"""Core credential initialization logic.
When *is_async* is True the caller is expected to have already
arranged for I/O to happen off the event loop (the sync helper
methods are called from ``asyncio.to_thread`` by the async
wrappers). For simplicity the sync/async split is handled by
``ensure_sync`` / ``ensure_async`` below.
"""
raise NotImplementedError # pragma: no cover dispatched below
def ensure_sync(self) -> None:
if self._initialized:
return
user_device_id = self.device_id
user_token = self.token
file_creds = self._load_sync()
if file_creds:
self._credentials = file_creds
if self.device_id is None:
self.device_id = file_creds.device_id
if self.token is None:
self.token = file_creds.token
need_save = False
if self.device_id is None:
self._credentials = register_device()
self.device_id = self._credentials.device_id
need_save = True
if self.token is None:
cdid = self._credentials.cdid if self._credentials else None
self.token = get_asr_token(self.device_id, cdid)
if self.credential_path and need_save and self._credentials:
self._credentials.token = self.token
self._save_sync(self._credentials)
if user_device_id is not None:
self.device_id = user_device_id
if user_token is not None:
self.token = user_token
self._initialized = True
async def ensure_async(self) -> None:
if self._initialized:
return
user_device_id = self.device_id
user_token = self.token
file_creds = await self._load_async()
if file_creds:
self._credentials = file_creds
if self.device_id is None:
self.device_id = file_creds.device_id
if self.token is None:
self.token = file_creds.token
need_save = False
if self.device_id is None:
self._credentials = await asyncio.to_thread(register_device)
self.device_id = self._credentials.device_id
need_save = True
if self.token is None:
cdid = self._credentials.cdid if self._credentials else None
self.token = await asyncio.to_thread(get_asr_token, self.device_id, cdid)
if self.credential_path and need_save and self._credentials:
self._credentials.token = self.token
await self._save_async(self._credentials)
if user_device_id is not None:
self.device_id = user_device_id
if user_token is not None:
self.token = user_token
self._initialized = True
def save_credentials_sync(self, creds: DeviceCredentials) -> None:
self._credentials = creds
self._save_sync(creds)
async def save_credentials_async(self, creds: DeviceCredentials) -> None:
self._credentials = creds
await self._save_async(creds)
def get_sami_token_sync(self) -> str:
self.ensure_sync()
if (self._credentials
and self._credentials.sami_token
and not _jwt_is_expired(self._credentials.sami_token)):
return self._credentials.sami_token
cdid = self._credentials.cdid if self._credentials else None
sami_token = get_sami_token(cdid)
if self._credentials:
self._credentials.sami_token = sami_token
self._save_sync(self._credentials)
return sami_token
async def get_sami_token_async(self) -> str:
await self.ensure_async()
if (self._credentials
and self._credentials.sami_token
and not _jwt_is_expired(self._credentials.sami_token)):
return self._credentials.sami_token
cdid = self._credentials.cdid if self._credentials else None
sami_token = await asyncio.to_thread(get_sami_token, cdid)
if self._credentials:
self._credentials.sami_token = sami_token
await self._save_async(self._credentials)
return sami_token
@@ -29,6 +29,14 @@ class DeviceCredentials(BaseModel):
""" """
用于 ASR 的 token 用于 ASR 的 token
""" """
sami_token: Optional[str] = None
"""
用于 NER 等 SAMI 服务的 token
"""
wave_session: Optional[dict] = None
"""
Wave 加密会话缓存(序列化后的 WaveSession
"""
class DeviceRegisterHeaderField(BaseModel): class DeviceRegisterHeaderField(BaseModel):
@@ -273,7 +281,6 @@ def register_device() -> DeviceCredentials:
params=params.model_dump(), params=params.model_dump(),
json=body.model_dump(), json=body.model_dump(),
headers=headers, headers=headers,
verify=False, # 禁用SSL证书验证以兼容某些环境
) )
response.raise_for_status() response.raise_for_status()
@@ -311,7 +318,6 @@ def get_asr_token(device_id: str, cdid: str) -> str:
params=params, params=params,
data=body_str, data=body_str,
headers=headers, headers=headers,
verify=False, # 禁用SSL证书验证以兼容某些环境
) )
response.raise_for_status() response.raise_for_status()
@@ -0,0 +1,91 @@
from __future__ import annotations
from enum import Enum, auto
from typing import List, Optional
from pydantic import BaseModel, Field
import uuid
AudioChunk = bytes
class ResponseType(Enum):
TASK_STARTED = auto()
SESSION_STARTED = auto()
SESSION_FINISHED = auto()
VAD_START = auto()
INTERIM_RESULT = auto()
FINAL_RESULT = auto()
HEARTBEAT = auto()
ERROR = auto()
UNKNOWN = auto()
class ASRWord(BaseModel):
word: str
start_time: float
end_time: float
class OIDecodingInfo(BaseModel):
oi_former_word_num: int = 0
oi_latter_word_num: int = 0
oi_words: Optional[List] = None
class ASRAlternative(BaseModel):
text: str
start_time: float
end_time: float
words: List[ASRWord] = []
semantic_related_to_prev: Optional[bool] = None
oi_decoding_info: Optional[OIDecodingInfo] = None
class ASRResult(BaseModel):
text: str
start_time: float
end_time: float
confidence: float = 0.0
alternatives: List[ASRAlternative] = []
is_interim: bool = True
is_vad_finished: bool = False
index: int = 0
class ASRExtra(BaseModel):
audio_duration: Optional[int] = None
model_avg_rtf: Optional[float] = None
model_send_first_response: Optional[int] = None
speech_adaptation_version: Optional[str] = None
model_total_process_time: Optional[int] = None
packet_number: Optional[int] = None
vad_start: Optional[bool] = None
req_payload: Optional[dict] = None
class ASRResponse(BaseModel):
type: ResponseType
text: str = ""
is_final: bool = False
vad_start: bool = False
vad_finished: bool = False
packet_number: int = -1
error_msg: str = ""
raw_json: Optional[dict] = None
results: List[ASRResult] = []
extra: Optional[ASRExtra] = None
class ASRError(Exception):
def __init__(self, message: str, response: Optional[ASRResponse] = None) -> None:
super().__init__(message)
self.response = response
class _SessionState(BaseModel):
request_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
final_text: str = ""
is_finished: bool = False
error: Optional[ASRResponse] = None
@@ -0,0 +1,126 @@
from __future__ import annotations
import base64
import json
import uuid
from typing import TYPE_CHECKING, List
import requests
from pydantic import BaseModel, Field
from .wave_client import WaveClient
from .constants import AID, APP_CONFIG, NER_URL, SAMI_APP_KEY
if TYPE_CHECKING:
from .config import ASRConfig
class NerUserInfo(BaseModel):
uid: str = "0"
did: str
app_name: str
app_version: str
sdk_version: str = ""
platform: str = "android"
experience_improve: bool = False
@classmethod
def new(cls, did: str, app_name: str):
app_version = APP_CONFIG.get("version_name", "")
return cls(did=did, app_name=app_name, app_version=app_version)
class NerRequest(BaseModel):
"""
ner 接口请求体
"""
user: NerUserInfo
text: str
additions: dict = Field(default_factory=dict)
@classmethod
def new(cls, text: str, did: str, app_name: str = "", addiction: dict = None):
return cls(user=NerUserInfo.new(did, app_name), text=text, additions=addiction or {})
class NerWord(BaseModel):
freq: int
word: str
class NerResult(BaseModel):
text: str
words: List[NerWord]
class NerResponse(BaseModel):
"""
ner 接口响应体
"""
results: List[NerResult]
def get_ner_results(wave_client: WaveClient, sami_token: str, text: str, did: str, app_name: str = "") -> NerResponse:
"""
调用 ner 接口获取结果
"""
request = NerRequest.new(text, did, app_name)
headers = {
'app_version': APP_CONFIG.get('version_name', ''),
'app_id': str(AID),
'os_type': 'android',
'x-api-resource-id': 'asr.user.context',
'x-api-app-key': SAMI_APP_KEY,
'x-api-token': sami_token,
'x-api-request-id': str(uuid.uuid4()),
}
req_data = request.model_dump_json().encode()
payload, headers = wave_client.prepare_request(req_data, headers)
response = requests.post(NER_URL, data=payload, headers=headers)
resp_headers = response.headers
nonce = base64.b64decode(resp_headers.get('x-tt-e-p'))
decoded = wave_client.decrypt(response.content, nonce=nonce)
return NerResponse(**json.loads(decoded.decode()))
def ner(config: ASRConfig, text: str, app_name: str = "") -> NerResponse:
"""
通过 ASRConfig 调用 NER 接口(同步便捷函数)
自动管理 WaveClient、sami_token 和 device_id。
:param config: ASR 配置(会自动初始化凭据)
:param text: 需要进行 NER 的文本
:param app_name: 应用名称(可选),可能会根据不同应用的使用场景适配不同的识别策略?"
:return: NER 响应
"""
config.ensure_credentials()
wave_client = config.get_wave_client()
sami_token = config.get_sami_token()
return get_ner_results(wave_client, sami_token, text, config.device_id, app_name)
async def async_ner(config: ASRConfig, text: str, app_name: str = "") -> NerResponse:
"""
通过 ASRConfig 调用 NER 接口(异步便捷函数,不阻塞事件循环)
自动管理 WaveClient、sami_token 和 device_id。
:param config: ASR 配置(会自动初始化凭据)
:param text: 需要进行 NER 的文本
:param app_name: 应用名称(可选)
:return: NER 响应
"""
import asyncio
await config.async_ensure_credentials()
wave_client = await config.async_get_wave_client()
sami_token = await config.async_get_sami_token()
return await asyncio.to_thread(
get_ner_results, wave_client, sami_token, text, config.device_id, app_name
)
@@ -0,0 +1,161 @@
from __future__ import annotations
import json
from typing import Optional
from .models import (
ASRWord,
OIDecodingInfo,
ASRAlternative,
ASRResult,
ASRExtra,
ASRResponse,
ResponseType,
)
from .asr_pb2 import AsrResponse as AsrResponsePb
def parse_word(data: dict) -> ASRWord:
return ASRWord(
word=data.get("word", ""),
start_time=data.get("start_time", 0.0),
end_time=data.get("end_time", 0.0),
)
def parse_oi_decoding_info(data: Optional[dict]) -> Optional[OIDecodingInfo]:
if data is None:
return None
return OIDecodingInfo(
oi_former_word_num=data.get("oi_former_word_num", 0),
oi_latter_word_num=data.get("oi_latter_word_num", 0),
oi_words=data.get("oi_words"),
)
def parse_alternative(data: dict) -> ASRAlternative:
words = [parse_word(w) for w in data.get("words", [])]
return ASRAlternative(
text=data.get("text", ""),
start_time=data.get("start_time", 0.0),
end_time=data.get("end_time", 0.0),
words=words,
semantic_related_to_prev=data.get("semantic_related_to_prev"),
oi_decoding_info=parse_oi_decoding_info(data.get("oi_decoding_info")),
)
def parse_result(data: dict) -> ASRResult:
alternatives = [parse_alternative(a) for a in data.get("alternatives", [])]
return ASRResult(
text=data.get("text", ""),
start_time=data.get("start_time", 0.0),
end_time=data.get("end_time", 0.0),
confidence=data.get("confidence", 0.0),
alternatives=alternatives,
is_interim=data.get("is_interim", True),
is_vad_finished=data.get("is_vad_finished", False),
index=data.get("index", 0),
)
def parse_extra(data: dict) -> ASRExtra:
return ASRExtra(
audio_duration=data.get("audio_duration"),
model_avg_rtf=data.get("model_avg_rtf"),
model_send_first_response=data.get("model_send_first_response"),
speech_adaptation_version=data.get("speech_adaptation_version"),
model_total_process_time=data.get("model_total_process_time"),
packet_number=data.get("packet_number"),
vad_start=data.get("vad_start"),
req_payload=data.get("req_payload"),
)
def parse_response(data: bytes) -> ASRResponse:
pb = AsrResponsePb()
pb.ParseFromString(data)
message_type = pb.message_type
result_json = pb.result_json
status_message = pb.status_message
if message_type == "TaskStarted":
return ASRResponse(type=ResponseType.TASK_STARTED)
if message_type == "SessionStarted":
return ASRResponse(type=ResponseType.SESSION_STARTED)
if message_type == "SessionFinished":
return ASRResponse(type=ResponseType.SESSION_FINISHED)
if message_type in ("TaskFailed", "SessionFailed"):
return ASRResponse(type=ResponseType.ERROR, error_msg=status_message)
if not result_json:
return ASRResponse(type=ResponseType.UNKNOWN)
try:
json_data = json.loads(result_json)
except json.JSONDecodeError:
return ASRResponse(type=ResponseType.UNKNOWN)
results_raw = json_data.get("results")
extra_raw = json_data.get("extra", {})
parsed_extra = parse_extra(extra_raw)
if results_raw is None:
return ASRResponse(
type=ResponseType.HEARTBEAT,
packet_number=extra_raw.get("packet_number", -1),
raw_json=json_data,
extra=parsed_extra,
)
parsed_results = [parse_result(r) for r in results_raw]
if extra_raw.get("vad_start"):
return ASRResponse(
type=ResponseType.VAD_START,
vad_start=True,
raw_json=json_data,
results=parsed_results,
extra=parsed_extra,
)
text = ""
is_interim = True
vad_finished = False
for r in parsed_results:
if r.text:
text = r.text
if not r.is_interim:
is_interim = False
if r.is_vad_finished:
vad_finished = True
nonstream_result = any(
r.get("extra", {}).get("nonstream_result") for r in results_raw
)
if nonstream_result or (not is_interim and vad_finished):
return ASRResponse(
type=ResponseType.FINAL_RESULT,
text=text,
is_final=True,
vad_finished=vad_finished,
raw_json=json_data,
results=parsed_results,
extra=parsed_extra,
)
return ASRResponse(
type=ResponseType.INTERIM_RESULT,
text=text,
is_final=False,
raw_json=json_data,
results=parsed_results,
extra=parsed_extra,
)
@@ -0,0 +1,52 @@
from __future__ import annotations
import json
from .config import SessionConfig
from .asr_pb2 import AsrRequest, FrameState
def build_start_task(request_id: str, token: str) -> bytes:
request = AsrRequest()
request.token = token
request.service_name = "ASR"
request.method_name = "StartTask"
request.request_id = request_id
return request.SerializeToString()
def build_start_session(request_id: str, token: str, config: SessionConfig) -> bytes:
request = AsrRequest()
request.token = token
request.service_name = "ASR"
request.method_name = "StartSession"
request.request_id = request_id
request.payload = config.model_dump_json()
return request.SerializeToString()
def build_finish_session(request_id: str, token: str) -> bytes:
request = AsrRequest()
request.token = token
request.service_name = "ASR"
request.method_name = "FinishSession"
request.request_id = request_id
return request.SerializeToString()
def build_asr_request(
audio_data: bytes,
request_id: str,
frame_state: FrameState,
timestamp_ms: int,
) -> bytes:
request = AsrRequest()
metadata = json.dumps({"extra": {}, "timestamp_ms": timestamp_ms})
request.service_name = "ASR"
request.method_name = "TaskRequest"
request.payload = metadata
request.audio_data = audio_data
request.request_id = request_id
request.frame_state = frame_state
return request.SerializeToString()
@@ -0,0 +1,137 @@
"""
SAMI 服务相关
"""
import hashlib
import time
import uuid
from typing import Optional
import requests
from pydantic import BaseModel, ConfigDict, Field
from .constants import SAMI_CONFIG_URL, SAMI_APP_KEY, USER_AGENT, APP_CONFIG, DEFAULT_DEVICE_CONFIG
class _SamiConfigParams(BaseModel):
"""
SAMI 配置接口的 URL Params
"""
model_config = ConfigDict(populate_by_name=True, serialize_by_alias=True)
device_platform: str = "android"
os: str = "android"
ssmix: str = "a"
rticket: str = Field(default_factory=lambda: str(int(time.time() * 1000)), serialization_alias="_rticket")
cdid: str
channel: str
aid: str
app_name: str
version_code: str
version_name: str
manifest_version_code: str
update_version_code: str
resolution: str
dpi: str
device_type: str
device_brand: str
language: str
os_api: str
os_version: str
ac: str = "wifi"
use_olympus_account: str = Field(default="1", serialization_alias="use-olympus-account")
@classmethod
def default(cls, cdid: str) -> "_SamiConfigParams":
"""
使用默认配置构建 SAMI 配置 Params
"""
app_config = {
**{k: APP_CONFIG[k] for k in ("channel", "app_name", "version_name")},
"aid": str(APP_CONFIG["aid"]),
"version_code": str(APP_CONFIG["version_code"]),
"manifest_version_code": str(APP_CONFIG["manifest_version_code"]),
"update_version_code": str(APP_CONFIG["update_version_code"]),
}
device_keys = ("device_platform", "os", "resolution", "dpi", "device_type",
"device_brand", "language", "os_api", "os_version")
device_config = {k: DEFAULT_DEVICE_CONFIG[k] for k in device_keys}
return cls(cdid=cdid, **app_config, **device_config)
class _SamiConfigRequest(BaseModel):
"""
SAMI 配置接口的请求体
"""
sami_app_key: str = SAMI_APP_KEY
class _SamiConfigData(BaseModel):
"""SAMI 配置数据"""
sami_token: str
class _SamiConfigResponse(BaseModel):
"""
SAMI 配置接口的响应
"""
code: int
msg: str
data: _SamiConfigData
@property
def sami_token(self) -> str:
return self.data.sami_token
def get_sami_config(cdid: str) -> requests.Response:
"""
获取 SAMI 配置 (包含 token)
Args:
cdid: 客户端设备 ID
Returns:
响应对象
"""
params = _SamiConfigParams.default(cdid)
body = _SamiConfigRequest()
body_json = body.model_dump_json()
x_ss_stub = hashlib.md5(body_json.encode()).hexdigest().upper()
headers = {
"User-Agent": USER_AGENT,
"Content-Type": "application/json",
"app_version": APP_CONFIG["version_name"],
"app_id": str(APP_CONFIG["aid"]),
"os_type": "Android",
"x-ss-stub": x_ss_stub,
}
response = requests.post(
SAMI_CONFIG_URL,
params=params.model_dump(by_alias=True),
data=body_json,
headers=headers,
)
return response
def get_sami_token(cdid: Optional[str] = None) -> str:
"""
获取 SAMI token
:param cdid: 客户端设备 ID,如果为 None 则自动生成
:return: SAMI token 字符串
"""
if cdid is None:
cdid = str(uuid.uuid4())
response = get_sami_config(cdid)
response.raise_for_status()
data = _SamiConfigResponse(**response.json())
return data.sami_token
@@ -0,0 +1,259 @@
"""
ByteDance Wave 加密协议客户端
"""
import base64
import hashlib
import secrets
import time
from typing import Callable, Optional, Union
import requests
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from pydantic import BaseModel
from .constants import HANDSHAKE_URL, HKDF_INFO, USER_AGENT
class _KeyShare(BaseModel):
"""
密钥交换信息
"""
curve: str
pubkey: str
class _HandshakeRequest(BaseModel):
"""
握手请求
"""
version: int = 2
random: str
app_id: str
did: str
key_shares: list[_KeyShare]
cipher_suites: list[int] = [4097] # ChaCha20
class _HandshakeResponse(BaseModel):
"""
握手响应
"""
version: int
random: str
key_share: _KeyShare
cipher_suite: int
cert: str
ticket: str
ticket_exp: int
ticket_long: str
ticket_long_exp: int
class WaveSession(BaseModel):
"""
Wave 加密会话
"""
ticket: str
ticket_long: str
encryption_key: bytes
client_random: bytes
server_random: bytes
shared_key: bytes
ticket_exp: int
ticket_long_exp: int
expires_at: float
class Config:
arbitrary_types_allowed = True
def is_expired(self) -> bool:
"""检查会话是否已过期"""
return time.time() >= self.expires_at
def to_dict(self) -> dict:
"""序列化为可 JSON 存储的 dictbytes 字段用 base64 编码)"""
return {
"ticket": self.ticket,
"ticket_long": self.ticket_long,
"encryption_key": base64.b64encode(self.encryption_key).decode(),
"client_random": base64.b64encode(self.client_random).decode(),
"server_random": base64.b64encode(self.server_random).decode(),
"shared_key": base64.b64encode(self.shared_key).decode(),
"ticket_exp": self.ticket_exp,
"ticket_long_exp": self.ticket_long_exp,
"expires_at": self.expires_at,
}
@classmethod
def from_dict(cls, data: dict) -> "WaveSession":
"""从 dict 反序列化(base64 解码 bytes 字段)"""
return cls(
ticket=data["ticket"],
ticket_long=data["ticket_long"],
encryption_key=base64.b64decode(data["encryption_key"]),
client_random=base64.b64decode(data["client_random"]),
server_random=base64.b64decode(data["server_random"]),
shared_key=base64.b64decode(data["shared_key"]),
ticket_exp=data["ticket_exp"],
ticket_long_exp=data["ticket_long_exp"],
expires_at=data["expires_at"],
)
class WaveClient:
"""
Wave 协议客户端
"""
def __init__(
self,
device_id: str,
app_id: Union[str, int],
session: Optional[WaveSession] = None,
on_session_update: Optional[Callable[[WaveSession], None]] = None,
):
self.device_id = device_id
self.app_id = str(app_id)
self.session = session
self._on_session_update = on_session_update
@staticmethod
def _chacha20_crypt(key: bytes, nonce: bytes, data: bytes) -> bytes:
"""ChaCha20 加密/解密"""
if len(nonce) == 12:
nonce_16 = b'\x00\x00\x00\x00' + nonce
else:
nonce_16 = nonce
cipher = Cipher(algorithms.ChaCha20(key, nonce_16), mode=None, backend=default_backend())
encryptor = cipher.encryptor()
return encryptor.update(data) + encryptor.finalize()
@staticmethod
def _derive_key(shared_key: bytes, salt: bytes, info: bytes) -> bytes:
"""HKDF 密钥派生"""
return HKDF(
algorithm=hashes.SHA256(),
length=32,
salt=salt,
info=info,
backend=default_backend()
).derive(shared_key)
def handshake(self) -> bool:
"""
执行 Wave 握手,建立加密会话
:return: 握手是否成功
"""
# 生成 ECDH 密钥对
private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())
client_random = secrets.token_bytes(32)
pubkey_bytes = private_key.public_key().public_bytes(
encoding=serialization.Encoding.X962,
format=serialization.PublicFormat.UncompressedPoint
)
request = _HandshakeRequest(
random=base64.b64encode(client_random).decode(),
app_id=self.app_id,
did=self.device_id,
key_shares=[_KeyShare(curve="secp256r1", pubkey=base64.b64encode(pubkey_bytes).decode())],
)
request_json = request.model_dump_json(by_alias=True)
# ECDSA 签名
signature = private_key.sign(request_json.encode(), ec.ECDSA(hashes.SHA256()))
headers = {
"Content-Type": "application/json",
"x-tt-s-sign": base64.b64encode(signature).decode(),
"User-Agent": USER_AGENT,
}
response = requests.post(HANDSHAKE_URL, data=request_json, headers=headers)
if response.status_code != 200:
return False
resp = _HandshakeResponse(**response.json())
# 计算共享密钥
server_pubkey = base64.b64decode(resp.key_share.pubkey)
server_public_key = ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), server_pubkey)
shared_key = private_key.exchange(ec.ECDH(), server_public_key)
server_random = base64.b64decode(resp.random)
# 派生加密密钥
salt = client_random + server_random
encryption_key = self._derive_key(shared_key, salt, HKDF_INFO)
# 保存会话,计算绝对过期时间(提前 60 秒刷新)
self.session = WaveSession(
ticket=resp.ticket,
ticket_long=resp.ticket_long,
encryption_key=encryption_key,
client_random=client_random,
server_random=server_random,
shared_key=shared_key,
ticket_exp=resp.ticket_exp,
ticket_long_exp=resp.ticket_long_exp,
expires_at=time.time() + resp.ticket_exp - 60,
)
if self._on_session_update:
self._on_session_update(self.session)
return True
def _ensure_session(self) -> None:
"""确保会话有效,如果过期则自动刷新"""
if self.session is None or self.session.is_expired():
if not self.handshake():
raise RuntimeError("Failed to establish/refresh session")
def prepare_request(self, plaintext: bytes, extra_headers: Optional[dict] = None) -> tuple[bytes, dict]:
"""
准备加密请求
:param plaintext: 明文数据
:param extra_headers: 额外的请求头
:return: (密文, headers) 元组
"""
self._ensure_session()
nonce = secrets.token_bytes(12)
ciphertext = self._chacha20_crypt(self.session.encryption_key, nonce, plaintext)
stub = hashlib.md5(ciphertext).hexdigest().upper()
headers = {
"Content-Type": "application/json",
"x-tt-e-b": "1",
"x-tt-e-t": self.session.ticket,
"x-tt-e-p": base64.b64encode(nonce).decode(),
"x-ss-stub": stub,
}
if extra_headers:
headers.update(extra_headers)
return ciphertext, headers
def decrypt(self, ciphertext: bytes, nonce: bytes) -> bytes:
"""
解密数据
:param ciphertext: 密文数据
:param nonce: 12 字节 nonce
:return: 明文数据
"""
if not self.session:
raise RuntimeError("No active session. Call handshake() first.")
return self._chacha20_crypt(self.session.encryption_key, nonce, ciphertext)
@@ -2,8 +2,8 @@
"domain": "hass_stt_doubao", "domain": "hass_stt_doubao",
"name": "Doubao Speech to Text", "name": "Doubao Speech to Text",
"version": "1.0.0", "version": "1.0.0",
"documentation": "https://github.com/yourusername/hass_stt_doubao", "documentation": "https://github.com/xyzmos/hass_stt_doubao",
"issue_tracker": "https://github.com/yourusername/hass_stt_doubao/issues", "issue_tracker": "https://github.com/xyzmos/hass_stt_doubao/issues",
"dependencies": [], "dependencies": [],
"codeowners": ["@xyzmos"], "codeowners": ["@xyzmos"],
"requirements": [ "requirements": [
@@ -14,5 +14,5 @@
"websockets>=12.0" "websockets>=12.0"
], ],
"config_flow": true, "config_flow": true,
"iot_class": "cloud_polling" "iot_class": "cloud_push"
} }