book-mcp/book_mcp_server.py

"""book-mcp — 教辅处理 MCP 服务（v1 最小可用，MongoDB 后端）

启动：/tmp/wordparse_venv/bin/python book_mcp_server.py
工具数：8 + 1（health）
"""
import os
import sys
import logging

# 让 services / db 包能 import
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import yaml
from fastmcp import FastMCP

from db.database import load_config
from db import get_collection, COL_DOC, health as db_health
from services import (
    book_service, document_service, file_service,
    rustfs_service, wordparse,
)

# ---- 日志 ----
_cfg = load_config()
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s [%(levelname)s] %(name)s: %(message)s')
log = logging.getLogger('book-mcp')

# ---- MCP server ----
mcp = FastMCP('book-mcp')


# ============================================================
# 1. create_book — 创建教辅
# ============================================================
@mcp.tool()
def create_book(name: str, zip_path: str, zip_sha256: str, zip_size: int = 0,
                subject: str | None = None, grade: str | None = None,
                version: str | None = None, type: str | None = None,
                created_by: str | None = None) -> dict:
    """创建教辅（按 zip_sha256 幂等）。

    Args:
        name: 教辅名（"2027版.曲一线3•2物理.广西版.一轮复习-突破训练"）
        zip_path: 对象存储里的 zip 路径
        zip_sha256: zip 文件 sha256（去重用）
        zip_size: 字节数
        subject/grade/version/type: 学科/年级/版本/类型（可选）
        created_by: 创建者（Agent 名 / 用户名）
    Returns: {ok, book_id, status, total_count, idempotent}
    """
    return book_service.create_book(
        name=name, zip_path=zip_path, zip_sha256=zip_sha256, zip_size=zip_size,
        subject=subject, grade=grade, version=version, type_=type,
        created_by=created_by,
    )


# ============================================================
# 2. add_document — 添加文档
# ============================================================
@mcp.tool()
def add_document(book_id: int, original_filename: str, original_path: str,
                 title: str | None = None, doc_type: str | None = None) -> dict:
    """向教辅添加一份文档（按 (book_id, original_path) 幂等）。

    Args:
        book_id: 教辅 ID
        original_filename: 原始文件名（含扩展名）
        original_path: zip 内的相对路径
        title: 文档标题（可选）
        doc_type: 文档类型（自动从文件名推断）
    Returns: {ok, doc_id, order_index, status, doc_type, idempotent}
    """
    return document_service.add_document(
        book_id=book_id, original_filename=original_filename,
        original_path=original_path, title=title, doc_type=doc_type,
    )


# ============================================================
# 3. mark_html_ready — 标记 HTML 已上传
# ============================================================
@mcp.tool()
def mark_html_ready(doc_id: str, html_url: str, html_size: int,
                    formulas: int = 0, images: int = 0, chars: int = 0) -> dict:
    """标记文档 HTML 已就绪。

    Args:
        doc_id: 文档 ID
        html_url: HTML 的对象存储 URL
        html_size: 字节数
        formulas/images/chars: 统计信息
    Returns: {ok, file_id, object_key, url}
    """
    stats = {'formulas': formulas, 'images': images, 'chars': chars}
    r = document_service.mark_html_ready(
        doc_id=doc_id, html_url=html_url, html_size=html_size, stats=stats,
    )
    if r.get('ok'):
        d = get_collection(COL_DOC).find_one({'_id': doc_id})
        if d:
            book_service.recompute_book_stats(d['book_id'])
    return r


# ============================================================
# 4. mark_json_ready — 标记 JSON 已上传
# ============================================================
@mcp.tool()
def mark_json_ready(doc_id: str, json_url: str, json_size: int,
                    schema_version: str | None = None) -> dict:
    """标记文档结构化 JSON 已就绪。

    Args:
        doc_id: 文档 ID
        json_url: JSON 的对象存储 URL
        json_size: 字节数
        schema_version: JSON schema 版本（可选）
    Returns: {ok, file_id, object_key, url}
    """
    return document_service.mark_json_ready(
        doc_id=doc_id, json_url=json_url, json_size=json_size,
        schema_version=schema_version,
    )


# ============================================================
# 5. mark_failed — 标记文档处理失败
# ============================================================
@mcp.tool()
def mark_failed(doc_id: str, error_msg: str, stage: str = 'html') -> dict:
    """标记文档处理失败。

    Args:
        doc_id: 文档 ID
        error_msg: 失败原因
        stage: 失败阶段（html / json / upload）
    Returns: {ok}
    """
    r = document_service.mark_failed(doc_id=doc_id, error_msg=error_msg, stage=stage)
    if r.get('ok'):
        d = get_collection(COL_DOC).find_one({'_id': doc_id})
        if d:
            book_service.recompute_book_stats(d['book_id'])
    return r


# ============================================================
# 6. update_book_status — 更新教辅状态
# ============================================================
@mcp.tool()
def update_book_status(book_id: int, status: str, error_msg: str | None = None) -> dict:
    """更新教辅状态。

    Args:
        book_id: 教辅 ID
        status: pending/processing/done/failed/partial
        error_msg: 错误信息（可选）
    Returns: {ok}
    """
    ok = book_service.update_book_status(book_id, status, error_msg)
    return {'ok': ok}


# ============================================================
# 7. get_book — 查教辅（含所有 doc + file）
# ============================================================
@mcp.tool()
def get_book(book_id: int) -> dict:
    """查教辅完整信息（教辅 + 所有文档 + 每个文档的所有文件）。

    Args:
        book_id: 教辅 ID
    Returns: 完整 book dict（documents 嵌套 files），或 None
    """
    return book_service.get_book(book_id, include_documents=True)


# ============================================================
# 8. list_books — 列表查询（分页）
# ============================================================
@mcp.tool()
def list_books(status: str | None = None, page: int = 1, page_size: int = 20,
               keyword: str | None = None) -> dict:
    """列出教辅（分页）。

    Args:
        status: 状态过滤（pending/processing/done/failed/partial）
        page: 页码（从 1 开始）
        page_size: 每页条数
        keyword: 关键词（name/subject/grade/version 模糊匹配）
    Returns: {books[], total, page, page_size}
    """
    return book_service.list_books(status=status, page=page,
                                   page_size=page_size, keyword=keyword)


# ============================================================
# 9. get_document — 查单文档（含所有 file）
# ============================================================
@mcp.tool()
def get_document(doc_id: str) -> dict:
    """查单文档（所有 file 一起返）。

    Args:
        doc_id: 文档 ID（16 字符 sha1）
    Returns: 完整 document dict（files[]），或 None
    """
    return document_service.get_document(doc_id, include_files=True)


# ============================================================
# 10. list_documents — 文档列表（按 book_id 分页）
# ============================================================
@mcp.tool()
def list_documents(book_id: int, status: str | None = None,
                   page: int = 1, page_size: int = 50) -> dict:
    """列出教辅的所有文档（按 order_index 排序，分页）。

    Args:
        book_id: 教辅 ID
        status: 状态过滤（pending/html_ready/json_ready/failed）
        page: 页码（从 1 开始）
        page_size: 每页条数
    Returns: {documents[], total, page, page_size}
    """
    return document_service.list_documents(book_id, status=status,
                                            page=page, page_size=page_size)


# ============================================================
# 11. get_document_url — 查文档某个 file 的 URL
# ============================================================
@mcp.tool()
def get_document_url(doc_id: str, file_type: str) -> dict:
    """查文档某个 file（html/json/docx/error_log）的 URL。

    Args:
        doc_id: 文档 ID
        file_type: html / json / docx / error_log
    Returns: {file_type, url, size, checksum} 或 None
    """
    f = file_service.get_file_by_type(doc_id, file_type)
    if not f:
        return None
    return {
        'file_type': f['file_type'],
        'url': f['url'],
        'size': f['size'],
        'checksum': f['checksum'],
    }


# ============================================================
# 12. upload_to_rustfs — 直接上传本地文件到对象存储
# ============================================================
@mcp.tool()
def upload_to_rustfs(local_path: str, key: str,
                     bucket: str | None = None) -> dict:
    """上传本地文件到 RustFS（Agent 用的上传助手）。

    Args:
        local_path: 本地文件路径
        key: 对象 key（如 'book/123/docs/abc/output.html'）
        bucket: 桶名（默认 config.yaml 的 default_bucket）
    Returns: {url, key, bucket, size, checksum}
    """
    if not os.path.exists(local_path):
        return {'ok': False, 'error': f'file not found: {local_path}'}
    return rustfs_service.upload(local_path, key, bucket=bucket)


# ============================================================
# 13. soft_delete_book — 软删教辅
# ============================================================
@mcp.tool()
def soft_delete_book(book_id: int) -> dict:
    """软删教辅（标记 is_deleted=1，不删对象存储）。

    Args:
        book_id: 教辅 ID
    Returns: {ok}
    """
    ok = book_service.soft_delete_book(book_id)
    return {'ok': ok}


# ============================================================
# 14. retry_failed — 重试失败的文档
# ============================================================
@mcp.tool()
def retry_failed(book_id: int) -> dict:
    """把教辅下所有 failed 文档改回 pending，方便重跑。

    Args:
        book_id: 教辅 ID
    Returns: {ok, retried_count}
    """
    from db import get_collection, COL_DOC
    col = get_collection(COL_DOC)
    r = col.update_many(
        {'book_id': book_id, 'status': 'failed'},
        {'$set': {'status': 'pending', 'error_msg': None}}
    )
    book_service.recompute_book_stats(book_id)
    return {'ok': True, 'retried_count': r.modified_count}


# ============================================================
# health — 健康检查
# ============================================================
@mcp.tool()
def health() -> dict:
    """健康检查：wordparse 服务 + mongo"""
    return {
        'wordparse': wordparse.health(),
        'mongo': db_health(),
    }


# ---- main ----
if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--port', type=int, default=8100)
    parser.add_argument('--host', default='127.0.0.1')
    parser.add_argument('--init-db', action='store_true', help='初始化 mongo 索引')
    args = parser.parse_args()

    if args.init_db:
        from db import init_db
        init_db()
        sys.exit(0)

    log.info(f"book-mcp 启动：http://{args.host}:{args.port}")
    log.info(f"wordparse 服务: {_cfg['wordparse']['url']}")
    # stateless_http + json_response：让前端跨域 JSON-RPC 调用不需要 SSE
    mcp.run(transport='http', host=args.host, port=args.port,
            stateless_http=True, json_response=True)