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