ai-video-fullstack/backend/routes/voice_webrtc.py

"""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

from db.session import SessionLocal
from fastapi import APIRouter, WebSocket
from loguru import logger
from models import AssistantConfig, SignalingOffer
from services.config_resolver import resolve_runtime_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"])


@router.get("/api/webrtc/ice-servers")
async def ice_servers():
    """Browser fetches STUN/TURN config (with ephemeral TURN creds when configured)."""
    return {"iceServers": client_ice_servers()}


@router.websocket("/ws/voice")
async def voice_signaling(websocket: WebSocket):
    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": 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:
            return await resolve_runtime_config(session, offer.assistant_id)
    if offer.inline_config:
        return offer.inline_config
    raise ValueError("offer 缺少 assistant_id 或 inline_config")


def _apply_vision_model(cfg: AssistantConfig) -> None:
    cfg.model = cfg.vision_model
    cfg.llm_interface_type = cfg.vision_llm_interface_type
    cfg.llm_values = cfg.vision_llm_values
    cfg.llm_secrets = cfg.vision_llm_secrets
    cfg.llm_support_image_input = cfg.vision_llm_support_image_input
    cfg.llm_api_key = cfg.vision_llm_api_key
    cfg.llm_base_url = cfg.vision_llm_base_url


async def _handle_offer(websocket, 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

    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)  # 解析放在建连前,配置错就别建连
        vision_enabled = offer.vision_enabled or cfg.vision_enabled
        if vision_enabled:
            if not cfg.vision_llm_support_image_input:
                raise ValueError(
                    "当前视觉模型不支持图片输入,请在模型资源中选择支持图片输入的 LLM"
                )
            _apply_vision_model(cfg)
        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

        @pc.event_handler("closed")
        async def _on_closed(conn: SmallWebRTCConnection):
            peers.pop(conn.pc_id, None)

        # 后台跑管线:WebRTC transport + 解析出的运行时配置
        transport = build_webrtc_transport(
            pc,
            video_in_enabled=vision_enabled,
        )
        asyncio.create_task(
            run_pipeline(transport, cfg, vision_enabled=vision_enabled)
        )

    answer = pc.get_answer()
    if websocket.application_state == WebSocketState.CONNECTED:
        await websocket.send_json(
            {
                "type": "answer",
                "payload": {
                    "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}")