audio_codebuddy/mcp-for-xiaozhi/audio_server.py

"""
audio_server.py — AI有声书 MCP 工具集（纯模块，无 I/O）

为什么从可执行脚本变成纯模块：
  Windows + Python 3.14 上，mcp_pipe.py 通过 subprocess.Popen + stdio PIPE
  启动子进程时，child 的 stdin 永远读不到数据（pipe 怪行为）。所以本工具
  改成纯模块，让 mcp_pipe.py 直接 import，避免任何跨进程通信。

提供的接口：
  - TOOLS：MCP 工具 schema 列表
  - handle_request(req) -> response | None：处理 xiaozhi 发来的 JSON-RPC 请求
    （notifications 返回 None，request 返回 response dict）
"""
from __future__ import annotations

import json
import logging
import os
import sys
import urllib.parse
import urllib.request
from typing import Any

API_BASE = os.environ.get("AUDIOBOOK_API_BASE", "http://127.0.0.1:3000")
API_TIMEOUT = float(os.environ.get("AUDIOBOOK_API_TIMEOUT", "15"))
MAX_RESP_BYTES = 900

logger = logging.getLogger("audiobook_mcp")


# ---------------- 同步 HTTP 客户端 ----------------
def _http_get(path: str, params: dict | None = None) -> dict:
    url = f"{API_BASE}{path}"
    if params:
        url += "?" + urllib.parse.urlencode(params)
    try:
        with urllib.request.urlopen(url, timeout=API_TIMEOUT) as r:
            return json.loads(r.read().decode("utf-8"))
    except Exception as e:  # noqa: BLE001
        return {"code": -1, "message": f"http error: {e}"}


def _http_post(path: str, body: dict | None = None) -> dict:
    url = f"{API_BASE}{path}"
    data = json.dumps(body or {}).encode("utf-8")
    try:
        req = urllib.request.Request(url, data=data, method="POST",
                                     headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=API_TIMEOUT) as r:
            return json.loads(r.read().decode("utf-8"))
    except Exception as e:  # noqa: BLE001
        return {"code": -1, "message": f"http error: {e}"}


# ---------------- 工具实现 ----------------
def _truncate(text: str, max_bytes: int = MAX_RESP_BYTES) -> str:
    """仅在可信的短文本上使用——不要对 JSON 整体截断。"""
    enc = text.encode("utf-8")
    if len(enc) <= max_bytes:
        return text
    return enc[:max_bytes].decode("utf-8", errors="ignore") + "..."


def _ok(payload: Any) -> dict:
    # 不截断外层 JSON——工具层已控制内容长度（text 截断到 250~300 字符）
    return {"success": True, "result": json.dumps(payload, ensure_ascii=False, separators=(",", ":"))}


def _err(msg: str) -> dict:
    return {"success": False, "error": _truncate(msg, 400)}


def tool_search_audiobooks(args: dict) -> dict:
    keywords = (args.get("keywords") or "").strip()
    if not keywords:
        return _err("搜索关键词不能为空")
    data = _http_get("/api/search", {"q": keywords, "limit": "8"})
    items = (data.get("data") or [])[:8]
    return _ok({
        "count": len(items),
        "items": [{"id": it.get("id") or it.get("audioId") or it.get("bookId"),
                   "title": it.get("title") or it.get("name") or "",
                   "desc": (it.get("description") or it.get("summary") or "")[:120]}
                  for it in items],
        "tip": "用 get_book_details(book_id) 查看章节",
    })


def tool_list_categories(args: dict) -> dict:
    data = _http_get("/api/categories")
    items = data.get("data") or []
    return _ok({"count": len(items),
                "categories": [{"id": it.get("id"), "name": it.get("name") or it.get("title")}
                               for it in items]})


def tool_list_books_by_category(args: dict) -> dict:
    cid = args.get("category_id")
    if not cid:
        return _err("category_id 必填")
    data = _http_get(f"/api/categories/{int(cid)}",
                     {"page": str(args.get("page") or 1), "pageSize": "10"})
    info = data.get("data") or {}
    items = info.get("list") if isinstance(info, dict) else info
    if not isinstance(items, list):
        items = []
    return _ok({"category_id": cid, "count": len(items[:10]),
                "books": [{"id": it.get("id") or it.get("bookId"),
                           "title": it.get("title") or "",
                           "desc": (it.get("description") or "")[:120]}
                          for it in items[:10]]})


def tool_get_book_details(args: dict) -> dict:
    bid = args.get("book_id")
    if not bid:
        return _err("book_id 必填")
    data = _http_get(f"/api/book-generator/books/{int(bid)}")
    if data.get("code") != 0:
        return _err("书籍不存在或无权访问，可用 search_audiobooks 重新搜索或用 generate_audiobook 生成新书")
    info = data.get("data") or {}
    chapters = info.get("chapters") or []
    stage = info.get("genStage") or ""
    progress = info.get("progress") or 0
    # 阶段→中文提示
    stage_hint = {"outlining": "正在生成大纲", "content_generating": "正在写内容",
                  "content_completed": "内容已生成", "audio_generating": "正在生成音频",
                  "audio_completed": "全部完成"}.get(stage, stage)
    result = {
        "id": info.get("id"),
        "title": info.get("title"),
        "progress": progress,
        "stage": stage_hint,
        "total_chapters": info.get("totalChapters"),
        "chapters": [{"id": ch.get("id"), "title": ch.get("title"),
                      "has_audio": bool(ch.get("audioUrl") and not str(ch.get("audioUrl", "")).startswith("/uploads/"))}
                     for ch in chapters[:30]],
    }
    if chapters:
        result["action"] = "有%d章内容，调get_chapter_text(chapter_id=%d)开始朗读" % (len(chapters), chapters[0]["id"])
    else:
        result["action"] = "还没生成完(%s %d%%)，告诉用户'正在%s，进度%d%%，说继续查看'。用户说继续后再调本工具" % (stage_hint, progress, stage_hint, progress)
    return _ok(result)


def _chapter_nav(book_id: int, chapter_id: int) -> dict:
    """获取章节导航：第几章/共几章/下一章ID/上一章ID"""
    nav = {"chapter_index": 1, "total_chapters": 1, "next_chapter_id": None, "prev_chapter_id": None}
    try:
        bdata = _http_get(f"/api/book-generator/books/{book_id}")
        chapters = (bdata.get("data") or {}).get("chapters") or []
        ids = [c["id"] for c in chapters if c.get("id")]
        nav["total_chapters"] = len(ids)
        if chapter_id in ids:
            idx = ids.index(chapter_id)
            nav["chapter_index"] = idx + 1
            if idx > 0:
                nav["prev_chapter_id"] = ids[idx - 1]
            if idx < len(ids) - 1:
                nav["next_chapter_id"] = ids[idx + 1]
    except Exception:  # noqa: BLE001
        pass
    return nav


def tool_get_chapter_audio_url(args: dict) -> dict:
    """【核心】获取章节音频URL。音箱用此URL直接播放。
    返回章节音频地址+导航信息（下一章ID、共几章）。播完当前章节后，
    请自动调用 get_chapter_audio_url(chapter_id=下一章ID) 继续播放下一章。"""
    cid = args.get("chapter_id")
    if not cid:
        return _err("chapter_id 必填")
    data = _http_get(f"/api/book-generator/books/chapters/{int(cid)}")
    info = data.get("data") or {}
    audio_url = info.get("audioUrl") or ""
    book_id = info.get("bookId")
    if not audio_url or audio_url.startswith("/uploads/"):
        # 无音频 → 自动降级为文本朗读，不让小智卡在错误上
        return tool_get_chapter_text({"chapter_id": cid, "max_chars": 200})
    nav = _chapter_nav(int(book_id) if book_id else 0, int(cid))
    result = {
        "chapter_id": cid,
        "title": info.get("title"),
        "audio_url": audio_url,
        "duration": info.get("audioDuration"),
        "book_id": book_id,
        **nav,
    }
    if nav["next_chapter_id"]:
        result["action"] = "播完本章后请自动调 get_chapter_audio_url(chapter_id=%s)" % nav["next_chapter_id"]
    else:
        result["action"] = "最后一章播完后告诉用户全部播放完毕并总结"
    return _ok(result)


def tool_get_chapter_text(args: dict) -> dict:
    """读正文自己TTS（推荐用 read_chapter）。"""
    cid = args.get("chapter_id")
    if not cid:
        return _err("chapter_id 必填")
    cid = int(cid)
    data = _http_get(f"/api/book-generator/books/chapters/{cid}/read",
                     {"format": "text", "max_chars": "300"})
    if data.get("code") != 0:
        return _err(data.get("message") or "读取失败")
    info = (data.get("data") or {})
    title = info.get("title") or ""
    text = info.get("text") or ""
    if not text:
        return _err("该章节暂无文本内容")
    nav = info.get("nav") or {}
    return {
        "action": "RESPONSE",
        "response": "%s。%s" % (title, text),
        "chapter_index": nav.get("chapter_index"),
        "total_chapters": nav.get("total_chapters"),
        "has_more": info.get("has_more"),
        **({"next_chapter_id": nav["next_chapter_id"]} if nav.get("next_chapter_id") else {}),
    }


def tool_read_chapter(args: dict) -> dict:
    """统一读章节：format=audio 返回URL播放；text 返回正文自己TTS；auto 智能选（推荐）。"""
    cid = args.get("chapter_id")
    if not cid:
        return _err("chapter_id 必填")
    fmt = args.get("format") or "auto"
    if fmt not in ("audio", "text", "auto"):
        return _err("format 必须是 audio/text/auto")

    # 调新端点
    data = _http_get(f"/api/book-generator/books/chapters/{int(cid)}/read",
                     {"format": fmt, "max_chars": "300"})
    code = data.get("code")

    # 2002 = on-demand TTS 失败 → 自动降级到 text 让音箱自己读
    if code == 2002:
        logger.warning("read_chapter TTS failed, fallback to text: %s", data.get("message"))
        data = _http_get(f"/api/book-generator/books/chapters/{int(cid)}/read",
                         {"format": "text", "max_chars": "300"})
    elif code != 0:
        return _err(data.get("message") or "章节读取失败")

    info = (data.get("data") or {})
    nav = (info.get("nav") or {})

    if info.get("mode") == "audio":
        result = {
            "chapter_id": info["chapter_id"],
            "title": info.get("title"),
            "audio_url": info.get("audio_url"),
            "duration": info.get("audio_duration"),
            "audio_source": info.get("audio_source"),  # 'full' 或 'on_demand'
            "book_id": info.get("book_id"),
            **nav,
        }
        next_id = nav.get("next_chapter_id")
        if next_id:
            result["action"] = "播完调 read_chapter(chapter_id=%s)" % next_id
        else:
            result["action"] = "最后一章，播完总结"
        return _ok(result)

    # text 模式
    text = info.get("text", "")
    title = info.get("title", "")
    return {
        "action": "RESPONSE",
        "response": "%s。%s" % (title, text),
        "has_more": info.get("has_more", False),
        "next_offset": info.get("next_offset"),
        "chapter_id": info["chapter_id"],
        **({"next_chapter_id": nav["next_chapter_id"]} if nav.get("next_chapter_id") else {}),
    }


def tool_list_recent_history(args: dict) -> dict:
    limit = max(1, min(int(args.get("limit") or 10), 30))
    data = _http_get("/api/history/list", {"page": "1", "pageSize": str(limit)})
    info = data.get("data") or {}
    items = info.get("list") or []
    return _ok({"count": len(items[:limit]),
                "history": [{"id": it.get("id"), "title": it.get("title"),
                             "type": it.get("type", "audio"),
                             "created_at": it.get("createdAt")}
                            for it in items[:limit]]})


def _bg_poll_and_push(book_id: int, title: str, device_id: str):
    """后台轮询生成进度，完成后主动 push 到小智音箱"""
    import time as _time
    import urllib.request as _rq
    push_url = os.environ.get("XIAOZHI_PUSH_URL", "http://192.168.31.155:58003/xiaozhi/push")
    deadline = _time.time() + 1800
    interval = 10
    while _time.time() < deadline:
        _time.sleep(interval)
        detail = _http_get(f"/api/book-generator/books/{book_id}")
        info = detail.get("data") or {}
        chapters = info.get("chapters") or []
        progress = info.get("progress", 0)
        interval = 5 if progress > 30 else 10
        if chapters and progress >= 30:
            first_c = chapters[0]
            fd = _http_get(f"/api/book-generator/books/chapters/{first_c['id']}")
            first_content = ((fd.get("data") or {}).get("content") or (fd.get("data") or {}).get("text") or "")
            if first_content and len(first_content) > 50:
                clean = first_content.replace("#", "").replace("*", "").strip()
                first_title = first_c.get("title") or ""
                book_title = info.get("title") or title
                push_text = "《%s》已生成！共%d章。%s。%s" % (book_title, len(chapters), first_title, clean[:200])
                try:
                    data = json.dumps({"device_id": device_id, "text": push_text, "type": "tts"}).encode()
                    req = _rq.Request(push_url, data=data, method="POST",
                                      headers={"Content-Type": "application/json"})
                    _rq.urlopen(req, timeout=10)
                    logger.info("push done: book_id=%d", book_id)
                except Exception as e:
                    logger.error("push failed: %s", e)
                return
    logger.warning("generate timeout: book_id=%d", book_id)


def tool_generate_audiobook(args: dict) -> dict:
    """一键生成有声书。创建书+后台生成→立即返回。完成后自动推送通知。"""
    import threading
    title = (args.get("title") or "").strip()
    if not title:
        return _err("title 必填")

    create = _http_post("/api/book-generator/books", {"title": title, "description": args.get("description") or ""})
    book_id = ((create.get("data") or {}) or {}).get("id")
    if not book_id:
        return _err(f"创建失败: {create.get('message') or create}")

    _http_post(f"/api/book-generator/books/{book_id}/batch-generate",
               {"steps": ["generate_content", "generate_audio"]})

    # 通知 xiaozhi-server 注册后台轮询任务
    xz_url = os.environ.get("XIAOZHI_SERVER_URL", "http://host.docker.internal:58003/xiaozhi/book-callback")
    device_id = (args.get("device_id") or "").strip()
    if device_id:
        try:
            cb = json.dumps({"book_id": book_id, "device_id": device_id, "title": title}).encode()
            urllib.request.urlopen(urllib.request.Request(xz_url, data=cb, method="POST",
                headers={"Content-Type": "application/json"}), timeout=5)
        except Exception:
            pass

    return {
        "action": "RESPONSE",
        "response": "已开始生成《%s》（书号%d），大约需要5-10分钟。" % (title, book_id),
        "book_id": book_id,
    }


# 精简工具 schema：7 个核心工具，目标 < 1024 字节。
TOOLS: list[dict] = [
    {"name": "search_audiobooks",
     "description": "搜书。",
     "inputSchema": {"type": "object",
                     "properties": {"keywords": {"type": "string"}},
                     "required": ["keywords"]}},
    {"name": "list_categories",
     "description": "分类。",
     "inputSchema": {"type": "object", "properties": {}}},
    {"name": "get_book_details",
     "description": "书详情+章节。",
     "inputSchema": {"type": "object",
                     "properties": {"book_id": {"type": "integer"}},
                     "required": ["book_id"]}},
    {"name": "read_chapter",
     "description": "audio返URL播，text返正文，auto智能选。",
     "inputSchema": {"type": "object",
                     "properties": {
                         "chapter_id": {"type": "integer"},
                         "format": {"type": "string", "enum": ["audio", "text", "auto"]},
                     },
                     "required": ["chapter_id"]}},
    {"name": "generate_audiobook",
     "description": "AI生成新书。异步，立即返回，完成后自动推送通知。",
     "inputSchema": {"type": "object",
                     "properties": {"title": {"type": "string"}, "device_id": {"type": "string"}},
                     "required": ["title"]}},
]

TOOL_DISPATCH = {
    "search_audiobooks": tool_search_audiobooks,
    "list_categories": tool_list_categories,
    "get_book_details": tool_get_book_details,
    "read_chapter": tool_read_chapter,
    "generate_audiobook": tool_generate_audiobook,
}

SERVER_INFO = {"name": "AudioBook", "version": "1.0.0"}
SERVER_CAPABILITIES = {"tools": {"listChanged": False}}
PROTOCOL_VERSION = "2024-11-05"


def handle_request(req: dict) -> dict | None:
    """处理一个 JSON-RPC 请求。
    - notification（无 id）返回 None
    - request 返回 response dict
    """
    method = req.get("method")
    req_id = req.get("id")
    params = req.get("params") or {}

    if req_id is None and method and not method.startswith("notifications/"):
        return {"jsonrpc": "2.0", "id": None,
                "error": {"code": -32600, "message": "invalid request"}}

    if method == "initialize":
        return {"jsonrpc": "2.0", "id": req_id, "result": {
            "protocolVersion": PROTOCOL_VERSION,
            "capabilities": SERVER_CAPABILITIES,
            "serverInfo": SERVER_INFO,
        }}
    if method == "ping":
        return {"jsonrpc": "2.0", "id": req_id, "result": {}}
    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}}
    if method == "tools/call":
        name = params.get("name")
        args = params.get("arguments") or {}
        handler = TOOL_DISPATCH.get(name)
        if not handler:
            return {"jsonrpc": "2.0", "id": req_id,
                    "error": {"code": -32602, "message": f"unknown tool: {name}"}}
        try:
            result = handler(args)
            text = json.dumps(result, ensure_ascii=False, separators=(",", ":"))
            return {"jsonrpc": "2.0", "id": req_id, "result": {
                "content": [{"type": "text", "text": text}],
                "isError": not result.get("success", True),
            }}
        except Exception as e:  # noqa: BLE001
            logger.exception("tool %s failed", name)
            return {"jsonrpc": "2.0", "id": req_id, "result": {
                "content": [{"type": "text", "text": json.dumps(_err(str(e)), ensure_ascii=False)}],
                "isError": True,
            }}

    if method and method.startswith("notifications/"):
        return None

    return {"jsonrpc": "2.0", "id": req_id,
            "error": {"code": -32601, "message": f"method not found: {method}"}}


# ---- stdio 入口（保留用于本地 stdio 调用，不通过 mcp_pipe） ----
def main_stdio():
    """直接 stdio 模式启动一个 MCP server（用于 verify_local.py / test_stdio.py）"""
    logging.basicConfig(
        level=os.environ.get("MCP_LOG_LEVEL", "INFO"),
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        stream=sys.stderr,
    )
    logger.info("AudioBook MCP server (stdio) starting, api_base=%s", API_BASE)
    while True:
        line = sys.stdin.buffer.readline()
        if not line:
            return
        try:
            req = json.loads(line.decode("utf-8").strip())
        except Exception as e:  # noqa: BLE001
            err = {"jsonrpc": "2.0", "id": None, "error": {"code": -32700, "message": f"parse error: {e}"}}
            sys.stdout.buffer.write((json.dumps(err, ensure_ascii=False) + "\n").encode("utf-8"))
            sys.stdout.buffer.flush()
            continue
        resp = handle_request(req)
        if resp is not None:
            sys.stdout.buffer.write((json.dumps(resp, ensure_ascii=False, separators=(",", ":")) + "\n").encode("utf-8"))
            sys.stdout.buffer.flush()


if __name__ == "__main__":
    main_stdio()