355 lines
12 KiB
Python
355 lines
12 KiB
Python
"""WebRTC 输出:SmallWebRTC 信令握手。
|
|
|
|
参考 dograh 的 webrtc_signaling.py,砍掉鉴权/配额/DB/org/ICE 过滤策略/TURN。
|
|
握手消息:
|
|
client → {type:"offer", payload:{pc_id, sdp, type, assistant_id}}
|
|
server → {type:"answer", payload:{pc_id, sdp, type}}
|
|
both → {type:"ice-candidate", payload:{pc_id, candidate:{...}}}
|
|
server → {type:"error", payload:{message}}
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
from collections.abc import Coroutine
|
|
from typing import Any
|
|
|
|
from db.session import SessionLocal
|
|
from fastapi import APIRouter, Body, Depends, Request, WebSocket
|
|
from loguru import logger
|
|
from models import AssistantConfig, SignalingOffer
|
|
from services.auth import require_admin, require_admin_websocket
|
|
from services.config_resolver import resolve_runtime_config
|
|
from services.runtime_variables import DynamicVariableError, prepare_dynamic_config
|
|
from starlette.websockets import WebSocketDisconnect, WebSocketState
|
|
|
|
from services.webrtc_ice import aiortc_ice_servers, client_ice_servers
|
|
|
|
# 注意:pipecat 是重依赖(语音才用),在 _handle_offer 等处惰性导入。
|
|
|
|
router = APIRouter(tags=["voice"])
|
|
_http_peers: dict[str, object] = {}
|
|
_active_connections: set[object] = set()
|
|
_pipeline_tasks: set[asyncio.Task[None]] = set()
|
|
_connection_tasks: dict[object, asyncio.Task[None]] = {}
|
|
PIPELINE_CLOSE_GRACE_SECONDS = 10.0
|
|
|
|
|
|
def _consume_pipeline_result(task: asyncio.Task[None], connection: object) -> None:
|
|
"""Retrieve background exceptions and release the task's strong reference."""
|
|
_pipeline_tasks.discard(task)
|
|
if _connection_tasks.get(connection) is task:
|
|
_connection_tasks.pop(connection, None)
|
|
try:
|
|
error = task.exception()
|
|
except asyncio.CancelledError:
|
|
logger.info(f"WebRTC pipeline 已取消: task={task.get_name()}")
|
|
return
|
|
if error is not None:
|
|
logger.opt(exception=error).error(
|
|
f"WebRTC pipeline 异常结束: task={task.get_name()}"
|
|
)
|
|
|
|
|
|
def _start_pipeline_task(
|
|
connection: object,
|
|
coroutine: Coroutine[Any, Any, None],
|
|
) -> asyncio.Task[None]:
|
|
"""Start one strongly referenced pipeline task for a WebRTC connection."""
|
|
connection_id = str(getattr(connection, "pc_id", "unknown"))
|
|
task = asyncio.create_task(
|
|
coroutine,
|
|
name=f"webrtc-pipeline:{connection_id}",
|
|
)
|
|
_active_connections.add(connection)
|
|
_pipeline_tasks.add(task)
|
|
_connection_tasks[connection] = task
|
|
task.add_done_callback(
|
|
lambda completed, connection=connection: _consume_pipeline_result(
|
|
completed,
|
|
connection,
|
|
)
|
|
)
|
|
return task
|
|
|
|
|
|
async def _wait_for_pipeline_close(
|
|
task: asyncio.Task[None] | None,
|
|
*,
|
|
connection_id: str,
|
|
) -> None:
|
|
"""Let transport disconnect finish normally, then cancel a stuck pipeline."""
|
|
if task is None:
|
|
return
|
|
try:
|
|
await asyncio.wait_for(
|
|
asyncio.shield(task),
|
|
timeout=PIPELINE_CLOSE_GRACE_SECONDS,
|
|
)
|
|
return
|
|
except TimeoutError:
|
|
logger.warning(
|
|
f"WebRTC pipeline 关闭超过 {PIPELINE_CLOSE_GRACE_SECONDS:g} 秒，"
|
|
f"执行取消: pc_id={connection_id}"
|
|
)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
# The done callback owns exception reporting and retrieval.
|
|
return
|
|
|
|
task.cancel()
|
|
done, _pending = await asyncio.wait(
|
|
{task},
|
|
timeout=PIPELINE_CLOSE_GRACE_SECONDS,
|
|
)
|
|
if not done:
|
|
logger.error(f"WebRTC pipeline 取消后仍未退出: pc_id={connection_id}")
|
|
|
|
|
|
async def shutdown_active_sessions() -> None:
|
|
"""Close active peers and drain every managed pipeline during app shutdown."""
|
|
connections = list(_active_connections)
|
|
if connections:
|
|
await asyncio.gather(
|
|
*(connection.disconnect() for connection in connections),
|
|
return_exceptions=True,
|
|
)
|
|
_active_connections.difference_update(connections)
|
|
_http_peers.clear()
|
|
|
|
tasks = list(_pipeline_tasks)
|
|
if not tasks:
|
|
return
|
|
done, pending = await asyncio.wait(
|
|
tasks,
|
|
timeout=PIPELINE_CLOSE_GRACE_SECONDS,
|
|
)
|
|
for task in pending:
|
|
task.cancel()
|
|
if pending:
|
|
cancelled, stuck = await asyncio.wait(
|
|
pending,
|
|
timeout=PIPELINE_CLOSE_GRACE_SECONDS,
|
|
)
|
|
if stuck:
|
|
logger.error(f"应用关闭时仍有 {len(stuck)} 个 WebRTC pipeline 未退出")
|
|
else:
|
|
cancelled = set()
|
|
logger.info(
|
|
f"WebRTC 会话清理完成: normal={len(done)} cancelled={len(cancelled)}"
|
|
)
|
|
|
|
|
|
@router.get("/api/webrtc/ice-servers", dependencies=[Depends(require_admin)])
|
|
async def ice_servers():
|
|
"""Browser fetches STUN/TURN config (with ephemeral TURN creds when configured)."""
|
|
return {"iceServers": client_ice_servers()}
|
|
|
|
|
|
@router.post("/api/webrtc/offer", dependencies=[Depends(require_admin)])
|
|
async def http_offer(request: Request, payload: dict = Body(...)):
|
|
"""Official SmallWebRTC JS transport offer endpoint."""
|
|
request_data = payload.get("requestData") or {}
|
|
encoded_variables = request.headers.get("x-pipecat-dynamic-variables", "")
|
|
header_variables = {}
|
|
if encoded_variables:
|
|
padding = "=" * (-len(encoded_variables) % 4)
|
|
header_variables = json.loads(
|
|
base64.urlsafe_b64decode(encoded_variables + padding).decode("utf-8")
|
|
)
|
|
offer_payload = {
|
|
**payload,
|
|
"assistant_id": request.headers.get("x-pipecat-assistant-id"),
|
|
"dynamic_variables": header_variables,
|
|
**request_data,
|
|
}
|
|
answer = await _handle_offer_payload(offer_payload, _http_peers)
|
|
return answer
|
|
|
|
|
|
@router.patch("/api/webrtc/offer", dependencies=[Depends(require_admin)])
|
|
async def http_ice_candidates(payload: dict = Body(...)):
|
|
"""Accept the batched trickle ICE format used by the official JS client."""
|
|
pc_id = payload.get("pc_id")
|
|
pc = _http_peers.get(pc_id) if pc_id else None
|
|
if not pc:
|
|
return {"ok": False}
|
|
|
|
from aiortc.sdp import candidate_from_sdp
|
|
|
|
for item in payload.get("candidates") or []:
|
|
candidate = candidate_from_sdp(item.get("candidate", ""))
|
|
candidate.sdpMid = item.get("sdp_mid")
|
|
candidate.sdpMLineIndex = item.get("sdp_mline_index")
|
|
await pc.add_ice_candidate(candidate)
|
|
return {"ok": True}
|
|
|
|
|
|
@router.websocket("/ws/voice")
|
|
async def voice_signaling(websocket: WebSocket):
|
|
if not await require_admin_websocket(websocket):
|
|
return
|
|
await websocket.accept()
|
|
peers: dict = {}
|
|
try:
|
|
while True:
|
|
message = await websocket.receive_json()
|
|
try:
|
|
if message.get("type") == "offer":
|
|
await _handle_offer(websocket, message.get("payload", {}), peers)
|
|
elif message.get("type") == "ice-candidate":
|
|
await _handle_ice(message.get("payload", {}), peers)
|
|
except Exception as e:
|
|
logger.exception(f"处理 WebRTC 信令消息失败: {e}")
|
|
if websocket.application_state == WebSocketState.CONNECTED:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "error",
|
|
"payload": {
|
|
"message": (
|
|
str(e)
|
|
if isinstance(e, DynamicVariableError)
|
|
else f"语音会话启动失败: {type(e).__name__}"
|
|
)
|
|
},
|
|
}
|
|
)
|
|
except WebSocketDisconnect:
|
|
logger.info("WebRTC 信令断开")
|
|
except Exception as e:
|
|
logger.error(f"WebRTC 信令出错: {e}")
|
|
finally:
|
|
# disconnect() triggers the registered closed callback, which removes
|
|
# the peer from this dict. Iterate over a snapshot to avoid mutation.
|
|
for pc in list(peers.values()):
|
|
await pc.disconnect()
|
|
|
|
|
|
async def _resolve_config(offer: SignalingOffer) -> AssistantConfig:
|
|
"""优先用 assistant_id 从 DB 解析(含真 key);否则用调试内联配置。"""
|
|
if offer.assistant_id:
|
|
async with SessionLocal() as session:
|
|
cfg = await resolve_runtime_config(session, offer.assistant_id)
|
|
return prepare_dynamic_config(
|
|
cfg,
|
|
offer.dynamic_variables,
|
|
assistant_id=offer.assistant_id,
|
|
)
|
|
if offer.inline_config:
|
|
return prepare_dynamic_config(
|
|
offer.inline_config,
|
|
offer.dynamic_variables,
|
|
assistant_id=None,
|
|
)
|
|
raise ValueError("offer 缺少 assistant_id 或 inline_config")
|
|
|
|
|
|
async def _handle_offer(websocket, payload, peers):
|
|
answer = await _handle_offer_payload(payload, peers)
|
|
if websocket.application_state == WebSocketState.CONNECTED:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "answer",
|
|
"payload": answer,
|
|
}
|
|
)
|
|
|
|
|
|
async def _handle_offer_payload(payload, peers):
|
|
from pipecat.transports.smallwebrtc.connection import SmallWebRTCConnection
|
|
from services.pipecat.pipeline import run_pipeline
|
|
from services.pipecat.transports import build_webrtc_transport
|
|
|
|
offer = SignalingOffer(**payload)
|
|
pc_id = offer.pc_id
|
|
|
|
restart_pc = bool(payload.get("restart_pc"))
|
|
if restart_pc and pc_id and pc_id in peers:
|
|
old_pc = peers.pop(pc_id)
|
|
old_task = _connection_tasks.get(old_pc)
|
|
await old_pc.disconnect()
|
|
await _wait_for_pipeline_close(
|
|
old_task,
|
|
connection_id=str(getattr(old_pc, "pc_id", pc_id)),
|
|
)
|
|
|
|
if pc_id and pc_id in peers:
|
|
pc = peers[pc_id]
|
|
await pc.renegotiate(sdp=offer.sdp, type=offer.type, restart_pc=False)
|
|
else:
|
|
cfg = await _resolve_config(offer) # 解析放在建连前,配置错就别建连
|
|
# 服务端助手配置是视觉理解的唯一授权来源；客户端 offer 只负责携带媒体轨。
|
|
if cfg.type == "workflow":
|
|
from services.workflow_engine import WorkflowEngine
|
|
|
|
vision_enabled = WorkflowEngine(cfg.graph).uses_vision()
|
|
else:
|
|
vision_enabled = cfg.vision_enabled
|
|
if vision_enabled and cfg.type != "workflow":
|
|
has_native_vision = (
|
|
not cfg.vision_model_resource_id and cfg.llm_support_image_input
|
|
)
|
|
has_aux_vision_model = (
|
|
bool(cfg.vision_model_resource_id)
|
|
and cfg.vision_llm_support_image_input
|
|
)
|
|
if not (has_native_vision or has_aux_vision_model):
|
|
raise ValueError(
|
|
"当前模型不支持图片输入,请在模型资源中选择支持图片输入的视觉模型"
|
|
)
|
|
pc = SmallWebRTCConnection(ice_servers=aiortc_ice_servers())
|
|
if pc_id:
|
|
pc._pc_id = pc_id
|
|
await pc.initialize(sdp=offer.sdp, type=offer.type)
|
|
peers[pc.pc_id] = pc
|
|
|
|
# 后台跑管线:WebRTC transport + 解析出的运行时配置
|
|
transport = build_webrtc_transport(
|
|
pc,
|
|
video_in_enabled=vision_enabled,
|
|
)
|
|
pipeline_task = _start_pipeline_task(
|
|
pc,
|
|
run_pipeline(
|
|
transport,
|
|
cfg,
|
|
vision_enabled=vision_enabled,
|
|
assistant_id=offer.assistant_id,
|
|
channel="webrtc",
|
|
),
|
|
)
|
|
|
|
@pc.event_handler("closed")
|
|
async def _on_closed(conn: SmallWebRTCConnection):
|
|
peers.pop(conn.pc_id, None)
|
|
_active_connections.discard(conn)
|
|
await _wait_for_pipeline_close(
|
|
pipeline_task,
|
|
connection_id=conn.pc_id,
|
|
)
|
|
|
|
answer = pc.get_answer()
|
|
return {
|
|
"pc_id": answer["pc_id"],
|
|
"sdp": answer["sdp"],
|
|
"type": answer["type"],
|
|
}
|
|
|
|
|
|
async def _handle_ice(payload, peers):
|
|
from aiortc.sdp import candidate_from_sdp
|
|
|
|
pc_id = payload.get("pc_id")
|
|
candidate_data = payload.get("candidate")
|
|
pc = peers.get(pc_id) if pc_id else None
|
|
if not pc or not candidate_data:
|
|
return
|
|
try:
|
|
candidate = candidate_from_sdp(candidate_data.get("candidate", ""))
|
|
candidate.sdpMid = candidate_data.get("sdpMid")
|
|
candidate.sdpMLineIndex = candidate_data.get("sdpMLineIndex")
|
|
await pc.add_ice_candidate(candidate)
|
|
except Exception as e:
|
|
logger.error(f"添加 ICE candidate 失败: {e}")
|