server.py

"""mcp-video-transcript —— 视频/音频转文字的 MCP 服务(Streamable-HTTP)。

通用工具,与任何业务无关:给一个媒体 URL,返回纯文本。可被 weft 的批量流水线定向调用,
也可被 Claude Desktop 等任意 MCP 客户端使用。

长任务约定(重要):MCP 是请求/响应模型、客户端超时通常 30~120s,而转写要几分钟到几十分钟,
所以走 **异步 submit / poll**:
    submit_transcribe(media_url) -> {job_id}      # 秒回
    get_transcribe_result(job_id) -> {status,...} # 轮询到 done/failed

运行:  python server.py            (默认 127.0.0.1:8010)
配置:  见 config.py(env VT_* 优先,其次 config.toml)
"""
from __future__ import annotations

import argparse
import contextlib
import json
import logging

from mcp.server.fastmcp import FastMCP

import engines
from config import CFG
from jobqueue import QUEUE

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("vt.server")

mcp = FastMCP("video-transcript", host=CFG.host, port=CFG.port)


@mcp.tool()
async def submit_transcribe(media_url: str, title: str = "", engine: str = "",
                            model: str = "", hotwords: str = "") -> str:
    """提交一个视频/音频转文字任务,立即返回 job_id(不等待完成)。

    转写耗时可达几十分钟,故本工具**不阻塞**:拿到 job_id 后用 get_transcribe_result 轮询。

    Args:
        media_url: 媒体地址(http/https,或 file:// 本地路径)。需要鉴权的地址由服务端配置的 key 携带。
        title: 可选,仅用于日志辨识。
        engine: 可选,ASR 引擎;留空用服务端默认(sensevoice)。
        model: 可选,ASR 模型 ID(见 list_asr_models);留空用服务端当前默认模型。
            调用方按会议/视频的语种或专业领域指定,不必改服务端配置。
        hotwords: 可选,空格分隔的热词,用于把专业名词在**识别阶段**就偏置进去。
            **只对支持上下文偏置的模型生效**(如 SeACo-Paraformer);SenseVoice 不支持,会被忽略。
    """
    QUEUE.start()
    if not (media_url or "").strip():
        return json.dumps({"ok": False, "error": "media_url 不能为空"}, ensure_ascii=False)
    job = QUEUE.submit(media_url.strip(), title=title, engine=engine,
                       model=model, hotwords=hotwords)
    logger.info("收到任务 job=%s title=%s", job.id, (title or "")[:40])
    return json.dumps({"ok": True, "job_id": job.id, "status": job.status}, ensure_ascii=False)


@mcp.tool()
async def get_transcribe_result(job_id: str) -> str:
    """查询转写任务状态/结果。

    status: queued(排队中) | running(转写中) | done(完成,text 为正文) |
            failed(失败,error 说明原因) | canceled(已取消)
    调用方应在 done/failed 之前按秒级~分钟级间隔轮询。
    """
    job = QUEUE.get((job_id or "").strip())
    if job is None:
        return json.dumps({"ok": False, "error": "job 不存在或已过期"}, ensure_ascii=False)
    return json.dumps({"ok": True, **job.public()}, ensure_ascii=False)


@mcp.tool()
async def cancel_transcribe(job_id: str) -> str:
    """取消尚未开始的转写任务(已在转写中的不会中断)。"""
    ok = QUEUE.cancel((job_id or "").strip())
    return json.dumps({"ok": ok, "message": "已取消" if ok else "任务不存在或已开始/已结束"},
                      ensure_ascii=False)


@mcp.tool()
async def transcribe_status() -> str:
    """服务自身状态:引擎/模型/设备、可用引擎、队列积压、已加载模型。用于连通性与健康检查。

    **不会触发模型加载** —— 只报告已经在内存里的实例。健康检查若会顺手加载 1GB 模型,
    那每次后台点"测试连接"都要等几十秒,还可能把内存打满。
    """
    import models as catalog
    import platform
    info = {
        "engines_available": engines.available(),
        "default_engine": CFG.engine,
        "default_model": CFG.model,
        "default_model_info": catalog.describe(CFG.model),
        "loaded_models": engines.loaded(),          # 空列表 = 还没跑过任务,属正常冷启动
        "queue": QUEUE.stats(),
        "ffmpeg": CFG.ffmpeg_bin,
        "host": f"{CFG.host}:{CFG.port}",
        "concurrency": CFG.concurrency,
        "platform": f"{platform.system()} {platform.machine()}",
    }
    try:
        import torch
        info["torch"] = {"version": torch.__version__,
                         "mps": bool(torch.backends.mps.is_available()),
                         "cuda": bool(torch.cuda.is_available())}
    except Exception:                      # noqa: BLE001
        info["torch"] = {"version": "(未安装)"}
    return json.dumps({"ok": True, **info}, ensure_ascii=False)


@mcp.tool()
async def list_asr_models() -> str:
    """可用 ASR 模型目录:ID、覆盖语言/方言、体积、是否支持热词、适用场景,**以及本地下载状态**。

    切模型前先看这个 —— 各模型的差别不在新旧,而在**覆盖面与代价**
    (中文最优 vs 多语种、快 vs 准、是否支持把专业词表喂进识别阶段)。

    每条附带 `disk`:{downloaded, bytes, path},以及 `downloading` 进度(若正在下)。
    判"已下载"按目录内**实际字节数**而非目录是否存在 —— 失败的下载会留下空壳目录。
    """
    import models as catalog
    out = []
    for m in catalog.CATALOG:
        d = dict(m)
        d["disk"] = catalog.disk_state(m["id"])
        d["downloading"] = DOWNLOADS.get(m["id"], {}).get("status") == "running"
        d["download_error"] = DOWNLOADS.get(m["id"], {}).get("error", "")
        out.append(d)
    return json.dumps({"ok": True, "current": CFG.model, "models": out}, ensure_ascii=False)


# 模型下载进度表。下载一个模型动辄几百 MB / 几分钟,**远超 MCP 的超时上限**,
# 故只能"发起 → 立刻返回 → 由 list_asr_models 报进度",与转写任务同一套异步思路。
DOWNLOADS: dict = {}


def _do_download(model_id: str) -> None:
    import models as catalog
    try:
        from modelscope.hub.snapshot_download import snapshot_download
        snapshot_download(model_id)
        st = catalog.disk_state(model_id)
        if not st["downloaded"]:
            raise RuntimeError(f"下载后本地仍不完整({st['bytes']} 字节)")
        DOWNLOADS[model_id] = {"status": "done", "error": "", "bytes": st["bytes"]}
        logger.info("模型下载完成 %s(%.0f MB)", model_id, st["bytes"] / 1e6)
    except Exception as e:                 # noqa: BLE001
        DOWNLOADS[model_id] = {"status": "failed", "error": str(e)[:300]}
        logger.warning("模型下载失败 %s:%s", model_id, e)
        # 失败会留下空壳目录,顺手清掉,免得下次被误判成"已下载"
        with contextlib.suppress(Exception):
            catalog.remove_local(model_id)


@mcp.tool()
async def download_asr_model(model: str) -> str:
    """下载模型到本地(**后台进行,立即返回**)。用 list_asr_models 看进度与结果。

    已下载则直接返回 already=True,不重复下。
    """
    import asyncio
    import models as catalog
    model = (model or "").strip()
    if not model or model == "custom":
        return json.dumps({"ok": False, "error": "model 不能为空"}, ensure_ascii=False)
    if catalog.disk_state(model)["downloaded"]:
        return json.dumps({"ok": True, "already": True, "message": "本地已有该模型"},
                          ensure_ascii=False)
    if DOWNLOADS.get(model, {}).get("status") == "running":
        return json.dumps({"ok": True, "already": False, "message": "正在下载中"},
                          ensure_ascii=False)
    DOWNLOADS[model] = {"status": "running", "error": ""}
    asyncio.get_running_loop().run_in_executor(None, _do_download, model)
    return json.dumps({"ok": True, "already": False, "message": "已开始后台下载"},
                      ensure_ascii=False)


@mcp.tool()
async def delete_asr_model(model: str) -> str:
    """删除本地模型文件(**目录条目仍会在 list_asr_models 里列出**,只是标记为未下载)。

    正在被加载使用的模型会先从内存卸载再删 —— 否则文件删了、进程里那份还在跑,
    状态就和磁盘对不上了。
    """
    import models as catalog
    model = (model or "").strip()
    if not model or model == "custom":
        return json.dumps({"ok": False, "error": "model 不能为空"}, ensure_ascii=False)
    if model == CFG.model:
        return json.dumps({"ok": False, "error": "该模型正在使用中,请先切换到别的模型再删除"},
                          ensure_ascii=False)
    engines.evict(model)                   # 先卸载内存里的实例
    r = catalog.remove_local(model)
    DOWNLOADS.pop(model, None)
    return json.dumps(r, ensure_ascii=False)


@mcp.tool()
async def set_asr_model(model: str, preload: bool = False) -> str:
    """切换服务端**默认**模型。preload=True 时立刻加载(几十秒),否则下个任务用时再加载。

    注意作用域:这只改服务端默认值,**进程重启后回到配置文件的值**。
    调用方(weft)应把自己的选择持久化,并在 submit_transcribe 里显式传 model —— 那才是权威来源。
    """
    model = (model or "").strip()
    if not model:
        return json.dumps({"ok": False, "error": "model 不能为空"}, ensure_ascii=False)
    old = CFG.model
    CFG.model = model
    out = {"ok": True, "old": old, "current": model, "preloaded": False}
    if preload:
        try:
            eng = await engines.get(CFG.engine, model)
            out["preloaded"] = True
            out["engine_info"] = eng.info()
        except Exception as e:             # noqa: BLE001
            CFG.model = old                # 加载不了就回滚,别把默认值改成一个用不了的模型
            return json.dumps({"ok": False, "error": f"模型加载失败,已回滚到 {old}:{str(e)[:300]}"},
                              ensure_ascii=False)
    return json.dumps(out, ensure_ascii=False)


# ── 低延迟同步端点(给实时语音客服用,不走 MCP)────────────────────────────────
# 为什么单开一条路:MCP 那套是 submit → 轮询 → 拉媒体 URL,批处理语义,一来一回秒级起步;
# 实时对话要的是"几百毫秒内把这 3 秒音频变成文字"。而模型**已经常驻在本进程里**(RSS 2.3GB),
# 让 weft 再装一份 torch + 再加载一份模型是纯浪费(内存翻倍、还要多付 20 秒加载)。
# 于是复用同一个进程、同一份模型,只是换一个**同步、直传 PCM、不落盘**的入口。

@mcp.custom_route("/asr", methods=["POST"])
async def asr_sync(request):
    """裸 PCM(16k 单声道 s16le)进,文本出。**同步返回**,给实时对话用。

    只接受**短音频**(默认 ≤60 秒):这是对话里的一轮,不是一整个视频。
    长音频请走 submit_transcribe —— 那边有队列与超时保护,不会把这个端点堵死。
    """
    from starlette.responses import JSONResponse
    import time as _t
    raw = await request.body()
    if not raw:
        return JSONResponse({"ok": False, "error": "空音频"}, status_code=400)
    secs = len(raw) / 32000                      # 16000Hz × 1ch × 2byte
    if secs > 60:
        return JSONResponse({"ok": False, "error": f"音频 {secs:.0f}s 超过 60s 上限,请走 submit_transcribe"},
                            status_code=413)
    model = request.query_params.get("model", "") or CFG.model
    hotwords = request.query_params.get("hotwords", "")
    t0 = _t.time()
    try:
        eng = await engines.get(CFG.engine, model)
        text = await eng.transcribe_pcm(raw, hotwords)
    except Exception as e:                        # noqa: BLE001
        logger.warning("同步转写失败:%s", e)
        return JSONResponse({"ok": False, "error": f"{type(e).__name__}: {str(e)[:200]}"},
                            status_code=500)
    el = _t.time() - t0
    logger.info("同步转写 %.1fs 音频 → %.0fms → %d 字", secs, el * 1000, len(text))
    return JSONResponse({"ok": True, "text": text, "audio_s": round(secs, 2),
                         "asr_ms": round(el * 1000)})


def main() -> None:
    ap = argparse.ArgumentParser(description="视频转写 MCP 服务")
    ap.add_argument("--host", default=None)
    ap.add_argument("--port", type=int, default=None)
    ap.add_argument("--preload", action="store_true", help="启动时就加载 ASR 模型(否则首个任务时才加载)")
    a = ap.parse_args()
    if a.host:
        CFG.host = mcp.settings.host = a.host
    if a.port:
        CFG.port = mcp.settings.port = a.port

    if a.preload:
        import asyncio
        logger.info("预加载引擎 %s …", CFG.engine)
        asyncio.run(engines.get(CFG.engine))

    logger.info("MCP 服务监听 http://%s:%s/mcp (引擎=%s, ffmpeg=%s)",
                CFG.host, CFG.port, CFG.engine, CFG.ffmpeg_bin)
    mcp.run(transport="streamable-http")


if __name__ == "__main__":
    main()