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 配置
This commit is contained in:
2026-05-01 22:32:19 +08:00
parent 6f28d70b0a
commit 2ffc20b14c
11 changed files with 962 additions and 1111 deletions
+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."""
hass.data.setdefault(DOMAIN, {})
# 获取配置
credential_path = entry.data.get(CONF_CREDENTIAL_PATH, DEFAULT_CREDENTIAL_PATH)
enable_punctuation = entry.data.get(CONF_ENABLE_PUNCTUATION, DEFAULT_ENABLE_PUNCTUATION)
_update_entry_data(hass, entry)
# 将相对路径转换为绝对路径
if not Path(credential_path).is_absolute():
credential_path = hass.config.path(credential_path)
entry.async_on_unload(
entry.add_update_listener(_async_update_entry_listener)
)
# 存储配置供 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)
_LOGGER.info("Doubao STT 集成已设置完成")
@@ -47,6 +38,29 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
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:
"""Unload a config entry."""
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]:
"""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)
enable_punctuation = data.get(CONF_ENABLE_PUNCTUATION, DEFAULT_ENABLE_PUNCTUATION)
# 将相对路径转换为 Home Assistant 配置目录下的绝对路径
if not Path(credential_path).is_absolute():
credential_path = hass.config.path(credential_path)
# 尝试初始化配置并验证凭据
try:
config = ASRConfig(
credential_path=credential_path,
enable_punctuation=enable_punctuation,
)
# 确保凭据已初始化(会自动注册设备如果需要)
await config.async_ensure_credentials()
# 简单验证:检查是否成功获取了 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
cred_path = Path(credential_path)
if cred_path.exists():
try:
cred_data = json_module.loads(cred_path.read_text(encoding="utf-8"))
if not isinstance(cred_data, dict):
raise CannotConnect("凭据文件格式不正确")
except (json_module.JSONDecodeError, OSError) as err:
raise CannotConnect(f"凭据文件读取失败: {err}") from err
# 返回用户可读的标题信息
return {
"title": "Doubao STT",
"credential_path": credential_path,
@@ -128,7 +117,6 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
errors: dict[str, str] = {}
if user_input is not None:
# 验证新的配置
try:
await validate_input(self.hass, user_input)
except CannotConnect:
@@ -137,6 +125,9 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
_LOGGER.exception("Unexpected exception")
errors["base"] = "unknown"
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)
# 获取当前配置值
@@ -1,5 +1,10 @@
from .asr import (
from .client import (
DoubaoASR,
transcribe,
transcribe_stream,
transcribe_realtime,
)
from .models import (
ASRResponse,
ASRResult,
ASRAlternative,
@@ -9,9 +14,6 @@ from .asr import (
ASRError,
ResponseType,
AudioChunk,
transcribe,
transcribe_stream,
transcribe_realtime,
)
from .config import ASRConfig
from .ner import NerResponse, NerResult, NerWord, ner, async_ner
@@ -1,789 +1,33 @@
from __future__ import annotations
import asyncio
import contextlib
from dataclasses import dataclass, field
from enum import Enum, auto
import json
from pathlib import Path
import ssl
import time
from typing import AsyncIterator, Callable, List, Optional, Union
import uuid
from pydantic import BaseModel, Field
import websockets
from websockets import ClientConnection
from .config import ASRConfig, SessionConfig
from .audio import AudioEncoder
from .asr_pb2 import AsrRequest, AsrResponse as AsrResponsePb, FrameState
# PCM 音频数据的类型别名
AudioChunk = bytes
class ResponseType(Enum):
"""
ASR 响应类型
"""
TASK_STARTED = auto()
SESSION_STARTED = auto()
SESSION_FINISHED = auto()
VAD_START = auto()
INTERIM_RESULT = auto()
FINAL_RESULT = auto()
HEARTBEAT = auto()
ERROR = auto()
UNKNOWN = auto()
@dataclass
class ASRWord:
"""单词级别的识别结果"""
word: str
start_time: float
end_time: float
@dataclass
class OIDecodingInfo:
"""OI 解码信息"""
oi_former_word_num: int = 0
oi_latter_word_num: int = 0
oi_words: Optional[List] = None
@dataclass
class ASRAlternative:
"""识别候选结果"""
text: str
start_time: float
end_time: float
words: List[ASRWord] = field(default_factory=list)
semantic_related_to_prev: Optional[bool] = None
oi_decoding_info: Optional[OIDecodingInfo] = None
@dataclass
class ASRResult:
"""单条识别结果"""
text: str
start_time: float
end_time: float
confidence: float = 0.0
alternatives: List[ASRAlternative] = field(default_factory=list)
is_interim: bool = True
is_vad_finished: bool = False
index: int = 0
@dataclass
class ASRExtra:
"""响应附加信息"""
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
@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
results: List[ASRResult] = field(default_factory=list)
extra: Optional[ASRExtra] = 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
def _create_ssl_context() -> ssl.SSLContext:
"""
在线程中创建并初始化 SSL 上下文,避免阻塞事件循环
"""
ctx = ssl.create_default_context()
return ctx
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 _get_ssl_context(self) -> ssl.SSLContext:
"""
获取预初始化的 SSL 上下文(懒加载,在线程中创建)
"""
if self._ssl_context is None:
self._ssl_context = await asyncio.to_thread(_create_ssl_context)
return self._ssl_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 = self._encoder.pcm_to_opus_frames(pcm_data)
state = _SessionState()
ws_url = await self.config.async_ws_url()
ssl_context = await self._get_ssl_context()
try:
async with websockets.connect(
ws_url,
additional_headers=self.config.headers,
open_timeout=self.config.connect_timeout,
ssl=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()
ws_url = await self.config.async_ws_url()
ssl_context = await self._get_ssl_context()
try:
async with websockets.connect(
ws_url,
additional_headers=self.config.headers,
open_timeout=self.config.connect_timeout,
ssl=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 数据并实时发送
"""
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 = self._encoder.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 = 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:
# 没有剩余数据,但需要发送一个 LAST 帧标记
# 发送一个空的 LAST 帧(静音)
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)
# FinishSession
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]:
"""
初始化 ASR 会话
"""
token = await self.config.async_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_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)
# FinishSession
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:
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_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]:
"""解析 OI 解码信息"""
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:
"""解析 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_raw = json_data.get("results")
extra_raw = json_data.get("extra", {})
# 解析为强类型
parsed_extra = _parse_extra(extra_raw)
# 无 results,可能是心跳包
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,
)
# 解析 results
parsed_results = [_parse_result(r) for r in results_raw]
# VAD 开始
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
nonstream_result = False
for r in results_raw:
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,
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,
)
# =============
# 便捷函数
# =============
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
from .client import (
DoubaoASR,
transcribe,
transcribe_stream,
transcribe_realtime,
)
from .models import (
ASRResponse,
ASRResult,
ASRAlternative,
ASRWord,
ASRExtra,
OIDecodingInfo,
ASRError,
ResponseType,
AudioChunk,
)
__all__ = [
"DoubaoASR",
"ASRResponse",
"ASRResult",
"ASRAlternative",
"ASRWord",
"ASRExtra",
"OIDecodingInfo",
"ASRError",
"ResponseType",
"AudioChunk",
"transcribe",
"transcribe_stream",
"transcribe_realtime",
]
@@ -0,0 +1,339 @@
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 = ssl.create_default_context()
async def __aenter__(self) -> DoubaoASR:
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,34 +1,13 @@
import asyncio
import base64
from dataclasses import dataclass, field
import json
from pathlib import Path
import time
from typing import Optional, Union
from pydantic import BaseModel
from .constants import WEBSOCKET_URL, USER_AGENT, AID
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:
"""
检查 JWT token 是否已过期(提前 margin 秒视为过期)
"""
try:
payload_b64 = token.split(".")[1]
# JWT base64url 需要补齐 padding
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
from .credential import CredentialManager
from .device import DeviceCredentials
class _AudioInfo(BaseModel):
@@ -47,9 +26,6 @@ class _SessionExtraConfig(BaseModel):
class SessionConfig(BaseModel):
"""
ASR 任务开始前需要初始化 Session 的配置
"""
audio_info: _AudioInfo
enable_punctuation: bool
enable_speech_rejection: bool
@@ -58,223 +34,61 @@ class SessionConfig(BaseModel):
@dataclass
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
aid: str = AID
user_agent: str = USER_AGENT
device_id: Optional[str] = None # 空则自动获取
token: Optional[str] = None # 空则自动获取
device_id: Optional[str] = None
token: Optional[str] = None
credential_path: Union[str, Path, None] = None
"""
凭据文件路径
"""
# 这些都是客户端给的默认值,挺通用的。其实我也没尝试过改了服务器会不会认
# 音频配置
sample_rate: int = 16000
channels: int = 1
frame_duration_ms: int = 20
# 会话配置
enable_punctuation: bool = True
enable_speech_rejection: bool = False
enable_asr_twopass: bool = True
enable_asr_threepass: bool = True
# 这里是输入法当前作用在哪个应用上
# 可能服务器会根据当前所使用的程序调整不同的语音识别策略??
# 这里用 Chrome 浏览器,算是比较通用的了吧?
app_name: str = "com.android.chrome"
# 连接配置
connect_timeout: float = 10.0
recv_timeout: float = 10.0
# 内部状态
_credentials: Optional[DeviceCredentials] = field(default=None, repr=None)
_initialized: bool = field(default=False, repr=False)
_credential_manager: Optional[CredentialManager] = field(default=None, repr=False)
_wave_client: Optional[object] = field(default=None, repr=False)
def _load_credentials_from_file_sync(self) -> Optional[DeviceCredentials]:
"""
同步的文件读取实现(在线程中调用)
"""
if self.credential_path is None:
return None
@property
def _cred_mgr(self) -> CredentialManager:
if self._credential_manager is None:
self._credential_manager = CredentialManager(
credential_path=self.credential_path,
device_id=self.device_id,
token=self.token,
)
return self._credential_manager
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)
def _sync_from_mgr(self) -> None:
self.device_id = self._cred_mgr.device_id
self.token = self._cred_mgr.token
except (json.JSONDecodeError, OSError):
return None
def _save_credentials_to_file_sync(self, creds: DeviceCredentials):
"""
同步的文件写入实现(在线程中调用)
"""
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_credentials_from_file(self) -> Optional[DeviceCredentials]:
"""
从缓存文件中加载凭据信息(异步,不阻塞事件循环)
"""
return await asyncio.to_thread(self._load_credentials_from_file_sync)
async def _save_credentials_to_file(self, creds: DeviceCredentials):
"""
保存凭据至缓存文件(异步,不阻塞事件循环)
"""
await asyncio.to_thread(self._save_credentials_to_file_sync, creds)
def ensure_credentials(self):
"""
确保凭据已初始化(同步版本,仅用于非 async 上下文)
"""
if self._initialized:
return
# 保存直接通过参数传入的凭据,用于进行覆盖
user_device_id = self.device_id
user_token = self.token
self._cred_mgr.ensure_sync()
self._sync_from_mgr()
# 尝试从文件中加载凭据
file_creds = self._load_credentials_from_file_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
# 如果 device_id 仍为 None, 则注册设备
need_save = False
if self.device_id is None:
self._credentials = 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
self.token = get_asr_token(self.device_id, cdid)
# 如果指定了 credential_path 且有新注册的凭据,则保存至文件
if self.credential_path and need_save and self._credentials:
self._credentials.token = self.token
self._save_credentials_to_file_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 async_ensure_credentials(self):
"""
确保凭据已初始化(异步版本,不阻塞事件循环)
await self._cred_mgr.ensure_async()
self._sync_from_mgr()
优先级:
1. 直接传入的 device_id/token 参数(最高优先级)
2. credential_path 文件中的值
3. 自动注册获取(最低优先级)
如果指定了 credential_path 且文件不存在,会注册设备并保存到文件。
"""
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, 则注册设备(网络 I/O 在线程中执行)
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
# 如果 token 仍为 None, 则获取 token(网络 I/O 在线程中执行)
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)
# 如果指定了 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
@property
def ws_url(self) -> str:
"""
获取 WebSocket URL(同步版本,需确保凭据已初始化)
"""
self.ensure_credentials()
return f'{self.url}?aid={self.aid}&device_id={self.device_id}'
async def async_ws_url(self) -> str:
"""
获取 WebSocket URL(异步版本,不阻塞事件循环)
"""
await self.async_ensure_credentials()
return f'{self.url}?aid={self.aid}&device_id={self.device_id}'
@property
def headers(self) -> dict[str, str]:
return {
@@ -284,23 +98,14 @@ class ASRConfig:
}
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()
return self._build_session_config()
def _build_session_config(self) -> SessionConfig:
"""
构建会话配置(内部方法,需确保凭据已初始化后调用)
"""
audio_info = _AudioInfo(
channel=self.channels,
format="speech_opus",
@@ -321,41 +126,31 @@ class ASRConfig:
enable_speech_rejection=self.enable_speech_rejection,
extra=extra,
)
def get_token(self) -> str:
"""
获取 token(同步版本,需确保凭据已初始化)
"""
self.ensure_credentials()
return self.token
async def async_get_token(self) -> str:
"""
获取 token(异步版本,不阻塞事件循环)
"""
await self.async_ensure_credentials()
return self.token
def _on_wave_session_update(self, session) -> None:
"""
Wave 会话更新时的回调,将新会话同步到凭据缓存
"""
if self._credentials:
self._credentials.wave_session = session.to_dict()
self._save_credentials_to_file_sync(self._credentials)
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):
"""
获取 WaveClient 实例(同步版本,懒加载,自动管理握手)
"""
from .wave_client import WaveClient, WaveSession
self.ensure_credentials()
if self._wave_client is None:
cached_session = None
if self._credentials and self._credentials.wave_session:
creds = self._cred_mgr.credentials
if creds and creds.wave_session:
try:
session = WaveSession.from_dict(self._credentials.wave_session)
session = WaveSession.from_dict(creds.wave_session)
if not session.is_expired():
cached_session = session
except (KeyError, ValueError):
@@ -370,17 +165,15 @@ class ASRConfig:
return self._wave_client
async def async_get_wave_client(self):
"""
获取 WaveClient 实例(异步版本,不阻塞事件循环)
"""
from .wave_client import WaveClient, WaveSession
await self.async_ensure_credentials()
if self._wave_client is None:
cached_session = None
if self._credentials and self._credentials.wave_session:
creds = self._cred_mgr.credentials
if creds and creds.wave_session:
try:
session = WaveSession.from_dict(self._credentials.wave_session)
session = WaveSession.from_dict(creds.wave_session)
if not session.is_expired():
cached_session = session
except (KeyError, ValueError):
@@ -395,41 +188,9 @@ class ASRConfig:
return self._wave_client
def get_sami_token(self) -> str:
"""
获取 SAMI token(同步版本,需确保凭据已初始化)
"""
self.ensure_credentials()
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_credentials_to_file_sync(self._credentials)
return sami_token
return self._cred_mgr.get_sami_token_sync()
async def async_get_sami_token(self) -> str:
"""
获取 SAMI token(异步版本,不阻塞事件循环)
"""
await self.async_ensure_credentials()
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_credentials_to_file(self._credentials)
return sami_token
return await self._cred_mgr.get_sami_token_async()
@@ -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
@@ -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,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()
@@ -2,8 +2,8 @@
"domain": "hass_stt_doubao",
"name": "Doubao Speech to Text",
"version": "1.0.0",
"documentation": "https://github.com/yourusername/hass_stt_doubao",
"issue_tracker": "https://github.com/yourusername/hass_stt_doubao/issues",
"documentation": "https://github.com/xyzmos/hass_stt_doubao",
"issue_tracker": "https://github.com/xyzmos/hass_stt_doubao/issues",
"dependencies": [],
"codeowners": ["@xyzmos"],
"requirements": [
@@ -14,5 +14,5 @@
"websockets>=12.0"
],
"config_flow": true,
"iot_class": "cloud_polling"
"iot_class": "cloud_push"
}