plugins/content/mcp.py

"""
内容提取插件 MCP工具定义
"""

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

from .core import ContentExtractorCore


class ExtractInput(BaseModel):
    """内容提取工具输入参数"""
    url: str = Field(description="内容链接（目前支持小红书、微博）")
    include_ocr: bool = Field(
        default=True,
        description="是否对图片进行OCR识别"
    )
    cookies: Optional[dict] = Field(
        default=None,
        description="Cookie（可选，用于应对反爬）"
    )


class ExtractOutput(BaseModel):
    """内容提取工具输出结果"""
    success: bool = Field(description="操作是否成功")
    title: str = Field(description="内容标题")
    content: str = Field(description="正文内容")
    images_text: str = Field(description="图片OCR识别的文本")
    source_type: str = Field(description="内容来源类型")
    error: Optional[str] = Field(default=None, description="错误信息（如果失败）")


# 全局核心逻辑实例
_core: ContentExtractorCore = None


def _get_core() -> ContentExtractorCore:
    """获取核心逻辑实例"""
    global _core
    if _core is None:
        _core = ContentExtractorCore()
    return _core


@mcp_tool(
    name="content-extract",
    title="内容提取",
    description="从支持的链接中提取内容，包括标题、正文和图片文字（OCR）",
    annotations={
        "readOnlyHint": False,
        "destructiveHint": False,
    }
)
async def extract_content(params: ExtractInput) -> ExtractOutput:
    """
    从链接中提取内容

    支持的平台：
    - 小红书：提取笔记标题、正文和图片中的文字
    - 微博：提取微博正文、标题和图片中的文字

    Args:
        params: 包含URL和配置参数

    Returns:
        ExtractOutput: 提取的内容
    """
    core = _get_core()
    result = await core.extract(
        url=params.url,
        include_ocr=params.include_ocr,
        cookies=params.cookies
    )

    return ExtractOutput(**result)