264 lines
8.7 KiB
Python
264 lines
8.7 KiB
Python
"""MiMo ASR adapter — ``POST /v1/chat/completions`` with ``input_audio``.
|
||
|
||
MiMo 的 ASR 不走独立音频端点，而是通过聊天补全接口扩展实现：
|
||
|
||
POST {base_url}/chat/completions
|
||
Content-Type: application/json
|
||
Authorization: Bearer {api_key}
|
||
|
||
{
|
||
"model": "mimo-v2.5-asr",
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "input_audio",
|
||
"input_audio": {
|
||
"data": "data:audio/wav;base64,<base64_encoded_audio>"
|
||
}
|
||
}
|
||
]
|
||
}
|
||
],
|
||
"asr_options": {
|
||
"language": "auto"
|
||
}
|
||
}
|
||
|
||
响应：:
|
||
|
||
{
|
||
"choices": [{
|
||
"message": {
|
||
"content": "转写后的文本"
|
||
}
|
||
}]
|
||
}
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import time
|
||
|
||
import httpx
|
||
|
||
from app.ai_provider.audio.base import (
|
||
ASRProtocol,
|
||
ASRRequest,
|
||
ASRResponse,
|
||
AudioAdapterBase,
|
||
AudioServiceError,
|
||
)
|
||
from app.core.logging import logger
|
||
|
||
_MIME_MAP = {
|
||
"wav": "audio/wav",
|
||
"mp3": "audio/mpeg",
|
||
"m4a": "audio/mp4",
|
||
"flac": "audio/flac",
|
||
}
|
||
|
||
# ASR 缓存：默认最多 128 条，TTL 1 小时
|
||
_ASR_CACHE_MAX_SIZE = 128
|
||
_ASR_CACHE_TTL_SECONDS = 3600
|
||
|
||
|
||
class _ASRCacheEntry:
|
||
__slots__ = ("text", "stored_at")
|
||
|
||
def __init__(self, text: str) -> None:
|
||
self.text = text
|
||
self.stored_at = time.monotonic()
|
||
|
||
|
||
_asr_cache: dict[str, _ASRCacheEntry] = {}
|
||
|
||
|
||
def _asr_cache_key(model: str, audio_data: bytes, language: str, fmt: str) -> str:
|
||
"""生成 ASR 缓存键：SHA256(model + audio_bytes + language + format)。"""
|
||
raw = f"{model}\0{language}\0{fmt}\0"
|
||
raw_bytes = raw.encode() + audio_data
|
||
return hashlib.sha256(raw_bytes).hexdigest()
|
||
|
||
|
||
def _asr_cache_get(key: str) -> str | None:
|
||
entry = _asr_cache.get(key)
|
||
if entry is None:
|
||
return None
|
||
if time.monotonic() - entry.stored_at > _ASR_CACHE_TTL_SECONDS:
|
||
del _asr_cache[key]
|
||
return None
|
||
return entry.text
|
||
|
||
|
||
def _asr_cache_set(key: str, text: str) -> None:
|
||
if len(_asr_cache) >= _ASR_CACHE_MAX_SIZE:
|
||
sorted_keys = sorted(_asr_cache.keys(), key=lambda k: _asr_cache[k].stored_at)
|
||
for old_key in sorted_keys[: _ASR_CACHE_MAX_SIZE // 2]:
|
||
del _asr_cache[old_key]
|
||
_asr_cache[key] = _ASRCacheEntry(text)
|
||
|
||
|
||
class HunjingASROpenAI(AudioAdapterBase, ASRProtocol):
|
||
"""MiMo ASR 适配器，通过 ``POST /chat/completions`` 携带 ``input_audio`` 实现。
|
||
|
||
将音频字节包装为 chat 消息中的 ``input_audio`` 内容类型的 data URL，
|
||
然后从响应中提取转写文本。
|
||
"""
|
||
|
||
_adapter_type = "asr"
|
||
_default_timeout = 300
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
base_url: str,
|
||
api_key: str,
|
||
model: str = "mimo-v2.5-asr",
|
||
request_timeout_seconds: int = 300,
|
||
) -> None:
|
||
self._base_url = base_url.rstrip("/")
|
||
self._api_key = api_key
|
||
self._model = model
|
||
self._timeout = request_timeout_seconds
|
||
self._client: httpx.AsyncClient | None = None
|
||
|
||
# ── 客户端 ───────────────────────────────────────────────────────────
|
||
|
||
@property
|
||
def _http(self) -> httpx.AsyncClient:
|
||
if self._client is None or self._client.is_closed:
|
||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
|
||
return self._client
|
||
|
||
async def close(self) -> None:
|
||
if self._client is not None and not self._client.is_closed:
|
||
await self._client.aclose()
|
||
|
||
# ── ASRProtocol ──────────────────────────────────────────────────────
|
||
|
||
async def transcribe(self, request: ASRRequest) -> ASRResponse:
|
||
"""通过 ``POST /chat/completions`` 转写音频。
|
||
|
||
将音频编码为 data URL 并以 ``input_audio`` 内容块的形式发送到 chat 消息中。
|
||
"""
|
||
# ── 缓存查找 ──────────────────────────────────────────────────
|
||
cache_key = _asr_cache_key(
|
||
self._model, request.audio_data, request.language or "auto", request.format
|
||
)
|
||
cached = _asr_cache_get(cache_key)
|
||
if cached is not None:
|
||
logger.info("audio_asr.cache_hit", model=self._model, text_length=len(cached))
|
||
return ASRResponse(text=cached)
|
||
|
||
# ── 构造请求 ──────────────────────────────────────────────────
|
||
mime = _MIME_MAP.get(request.format, "audio/wav")
|
||
b64_audio = base64.b64encode(request.audio_data).decode()
|
||
data_url = f"data:{mime};base64,{b64_audio}"
|
||
|
||
payload: dict = {
|
||
"model": self._model,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "input_audio",
|
||
"input_audio": {
|
||
"data": data_url,
|
||
},
|
||
},
|
||
],
|
||
},
|
||
],
|
||
"asr_options": {
|
||
"language": request.language or "auto",
|
||
},
|
||
}
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {self._api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
url = f"{self._base_url}/chat/completions"
|
||
|
||
logger.info(
|
||
"audio_asr.request",
|
||
model=self._model,
|
||
format=request.format,
|
||
language=request.language,
|
||
url=url,
|
||
size_bytes=len(request.audio_data),
|
||
)
|
||
|
||
# ── API 调用 ────────────────────────────────────────────────
|
||
try:
|
||
response = await self._http.post(url, json=payload, headers=headers)
|
||
except httpx.TimeoutException as exc:
|
||
raise AudioServiceError(
|
||
status_code=504,
|
||
code="audio_asr_timeout",
|
||
message="ASR request timed out calling chat completions.",
|
||
) from exc
|
||
except httpx.RequestError as exc:
|
||
raise AudioServiceError(
|
||
status_code=503,
|
||
code="audio_asr_connection_error",
|
||
message=f"ASR connection error: {exc}",
|
||
) from exc
|
||
|
||
if response.status_code != 200:
|
||
body = response.text[:500]
|
||
raise AudioServiceError(
|
||
status_code=response.status_code,
|
||
code="audio_asr_api_error",
|
||
message=f"ASR API returned {response.status_code}: {body}",
|
||
)
|
||
|
||
# 解析响应，提取转写文本
|
||
text = self._extract_text(response.json())
|
||
if not text:
|
||
raise AudioServiceError(
|
||
status_code=500,
|
||
code="audio_asr_empty_response",
|
||
message="ASR API returned empty transcription.",
|
||
)
|
||
|
||
# ── 写入缓存 ────────────────────────────────────────────────
|
||
_asr_cache_set(cache_key, text)
|
||
|
||
logger.info(
|
||
"audio_asr.success",
|
||
model=self._model,
|
||
text_length=len(text),
|
||
cache_size=len(_asr_cache),
|
||
)
|
||
return ASRResponse(text=text)
|
||
|
||
# ── 辅助方法 ──────────────────────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _extract_text(body: dict) -> str:
|
||
"""从 chat completion 响应中提取转写文本。"""
|
||
try:
|
||
choices = body.get("choices", [])
|
||
if not choices:
|
||
raise AudioServiceError(
|
||
status_code=500,
|
||
code="audio_asr_no_choices",
|
||
message="ASR response has no choices.",
|
||
)
|
||
return (choices[0].get("message") or {}).get("content") or ""
|
||
except AudioServiceError:
|
||
raise
|
||
except Exception as exc:
|
||
raise AudioServiceError(
|
||
status_code=500,
|
||
code="audio_asr_response_parse_error",
|
||
message=f"Failed to parse ASR response: {exc}",
|
||
) from exc
|