plugins/doc/mcp.py

"""
文档合并拆分插件 MCP工具定义
"""

from app.mcp.decorators import mcp_tool
from pydantic import BaseModel, Field
from typing import List

from .core import do_merge, do_split


class MergeInput(BaseModel):
    """文档合并输入参数"""
    files_base64: List[str] | None = Field(
        default=None, description="待合并文件的Base64编码列表"
    )
    files_url: List[str] | None = Field(
        default=None, description="待合并文件的URL列表"
    )
    output_format: str = Field(default="markdown", description="输出格式：markdown、text")


class SplitInput(BaseModel):
    """文档拆分输入参数"""
    content_base64: str | None = Field(default=None, description="待拆分文档的Base64编码")
    content_url: str | None = Field(default=None, description="待拆分文档的URL")
    split_by: str = Field(default="line", description="拆分方式：line、paragraph、chunk")
    chunk_size: int = Field(default=100, ge=10, le=10000, description="chunk模式下的块大小")


class MergeOutput(BaseModel):
    """合并输出结果"""
    success: bool = Field(description="操作是否成功")
    output_base64: str | None = Field(description="合并结果的Base64编码")
    message: str = Field(description="结果说明")


class SplitOutput(BaseModel):
    """拆分输出结果"""
    success: bool = Field(description="操作是否成功")
    chunks_base64: List[str] | None = Field(description="拆分片段的Base64编码列表")
    message: str = Field(description="结果说明")


@mcp_tool(name="doc-merge", title="文档合并", description="合并多个文档内容")
async def doc_merge(params: MergeInput) -> MergeOutput:
    result = await do_merge(params.files_base64, params.output_format)
    return MergeOutput(**result)


@mcp_tool(name="doc-split", title="文档拆分", description="拆分文档为多个片段")
async def doc_split(params: SplitInput) -> SplitOutput:
    result = await do_split(params.content_base64, params.split_by, params.chunk_size)
    return SplitOutput(**result)