plugins/json/mcp.py

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

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


class ExtractFromJsonInput(BaseModel):
    """从JSON数据提取content的输入参数"""
    json_data: Any = Field(description="JSON数据对象")
    output_file: Optional[str] = Field(
        default=None,
        description="输出文件路径（可选，不提供则返回格式化文本）"
    )


class ExtractFromJsonOutput(BaseModel):
    """从JSON数据提取content的输出结果"""
    success: bool = Field(description="操作是否成功")
    contents: List[str] = Field(default=[], description="提取的content列表")
    formatted_output: str = Field(default="", description="格式化后的输出")
    count: int = Field(default=0, description="提取的content数量")
    error: Optional[str] = Field(default=None, description="错误信息")


class ExtractFromFileInput(BaseModel):
    """从JSON文件提取content的输入参数"""
    file_path: str = Field(description="JSON文件路径")
    output_file: Optional[str] = Field(
        default=None,
        description="输出文件路径（可选）"
    )


class ExtractFromFileOutput(BaseModel):
    """从JSON文件提取content的输出结果"""
    success: bool = Field(description="操作是否成功")
    file_path: str = Field(description="源文件路径")
    contents: List[str] = Field(default=[], description="提取的content列表")
    formatted_output: str = Field(default="", description="格式化后的输出")
    count: int = Field(default=0, description="提取的content数量")
    error: Optional[str] = Field(default=None, description="错误信息")


class ExtractFromDirectoryInput(BaseModel):
    """从目录批量提取content的输入参数"""
    dir_path: str = Field(description="包含JSON文件的目录路径")
    output_file: Optional[str] = Field(
        default=None,
        description="输出文件路径（可选）"
    )


class ExtractFromDirectoryOutput(BaseModel):
    """从目录批量提取content的输出结果"""
    success: bool = Field(description="操作是否成功")
    dir_path: str = Field(description="源目录路径")
    file_count: int = Field(default=0, description="处理的文件数量")
    result: str = Field(default="", description="提取结果")
    error: Optional[str] = Field(default=None, description="错误信息")


class ConversationMessage(BaseModel):
    """对话消息结构"""
    role: str = Field(description="消息角色（system/user/assistant）")
    content: str = Field(description="消息内容")


class ExtractConversationFromJsonInput(BaseModel):
    """从JSON数据提取对话的输入参数"""
    json_data: Any = Field(description="JSON数据对象")


class ExtractConversationFromJsonOutput(BaseModel):
    """从JSON数据提取对话的输出结果"""
    success: bool = Field(description="操作是否成功")
    is_conversation: bool = Field(description="是否为对话格式")
    messages: List[ConversationMessage] = Field(default=[], description="对话消息列表")
    formatted_output: str = Field(default="", description="格式化后的对话记录")
    count: int = Field(default=0, description="消息数量")
    error: Optional[str] = Field(default=None, description="错误信息")


class ExtractConversationFromFileInput(BaseModel):
    """从JSON文件提取对话的输入参数"""
    file_path: str = Field(description="JSON文件路径")


class ExtractConversationFromFileOutput(BaseModel):
    """从JSON文件提取对话的输出结果"""
    success: bool = Field(description="操作是否成功")
    file_path: str = Field(description="源文件路径")
    is_conversation: bool = Field(description="是否为对话格式")
    messages: List[ConversationMessage] = Field(default=[], description="对话消息列表")
    formatted_output: str = Field(default="", description="格式化后的对话记录")
    count: int = Field(default=0, description="消息数量")
    error: Optional[str] = Field(default=None, description="错误信息")


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


def _get_core():
    """获取核心逻辑实例"""
    global _core
    if _core is None:
        from .core import JsonContentExtractorCore
        _core = JsonContentExtractorCore()
    return _core


@mcp_tool(
    name="json-content",
    title="从JSON数据提取content",
    description="从JSON数据对象中递归提取所有content字段的内容。如果是OpenAI对话格式，会自动识别并按对话格式输出。",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
    }
)
async def extract_from_json(params: ExtractFromJsonInput) -> ExtractFromJsonOutput:
    """
    从JSON数据中提取content字段

    支持递归提取嵌套的content字段，包括处理OpenAI API格式的content数组。
    如果是OpenAI对话格式（包含messages数组），会自动识别并按对话格式输出。

    Args:
        params: 包含JSON数据

    Returns:
        ExtractFromJsonOutput: 提取的内容
    """
    core = _get_core()

    try:
        # 检测是否为对话格式
        if core._is_conversation_format(params.json_data):
            formatted_output = core._format_conversation(params.json_data)
            messages = core._extract_conversation(params.json_data)
            contents = [msg["content"] for msg in messages]
        else:
            contents = core.extract_content_from_json(params.json_data)

            formatted_contents = []
            for i, content in enumerate(contents, 1):
                formatted = core.format_content(content)
                formatted_contents.append(f"=== Content {i} ===\n{formatted}")

            formatted_output = "\n\n".join(formatted_contents) if formatted_contents else "未找到content内容"

        # 如果指定了输出文件，保存结果
        if params.output_file and contents:
            try:
                with open(params.output_file, 'w', encoding='utf-8') as f:
                    f.write(formatted_output)
            except Exception as e:
                return ExtractFromJsonOutput(
                    success=False,
                    contents=contents,
                    formatted_output=formatted_output,
                    count=len(contents),
                    error=f"保存文件失败: {str(e)}"
                )

        return ExtractFromJsonOutput(
            success=True,
            contents=contents,
            formatted_output=formatted_output,
            count=len(contents)
        )

    except Exception as e:
        return ExtractFromJsonOutput(
            success=False,
            error=str(e)
        )


@mcp_tool(
    name="json-file",
    title="从JSON文件提取content",
    description="从JSON文件中递归提取所有content字段的内容。如果是OpenAI对话格式，会自动识别并按对话格式输出。",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
    }
)
async def extract_from_file(params: ExtractFromFileInput) -> ExtractFromFileOutput:
    """
    从JSON文件中提取content字段

    如果是OpenAI对话格式，会自动识别并按对话格式输出。

    Args:
        params: 包含文件路径

    Returns:
        ExtractFromFileOutput: 提取的内容
    """
    core = _get_core()

    try:
        import os
        if not os.path.exists(params.file_path):
            return ExtractFromFileOutput(
                success=False,
                file_path=params.file_path,
                error=f"文件不存在: {params.file_path}"
            )

        contents = core.extract_content_from_file(params.file_path)
        result = core.process_single_file(params.file_path, output_file=params.output_file)

        return ExtractFromFileOutput(
            success=True,
            file_path=params.file_path,
            contents=contents,
            formatted_output=result,
            count=len(contents)
        )

    except Exception as e:
        return ExtractFromFileOutput(
            success=False,
            file_path=params.file_path,
            error=str(e)
        )


@mcp_tool(
    name="json-dir",
    title="从目录批量提取content",
    description="从目录中所有JSON文件递归提取content字段的内容",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
    }
)
async def extract_from_directory(params: ExtractFromDirectoryInput) -> ExtractFromDirectoryOutput:
    """
    从目录中批量提取content字段

    Args:
        params: 包含目录路径

    Returns:
        ExtractFromDirectoryOutput: 提取的结果
    """
    core = _get_core()

    try:
        import os
        import glob

        if not os.path.exists(params.dir_path):
            return ExtractFromDirectoryOutput(
                success=False,
                dir_path=params.dir_path,
                error=f"目录不存在: {params.dir_path}"
            )

        json_files = glob.glob(os.path.join(params.dir_path, "*.json"))
        result = core.process_directory(params.dir_path, output_file=params.output_file)

        return ExtractFromDirectoryOutput(
            success=True,
            dir_path=params.dir_path,
            file_count=len(json_files),
            result=result
        )

    except Exception as e:
        return ExtractFromDirectoryOutput(
            success=False,
            dir_path=params.dir_path,
            error=str(e)
        )


@mcp_tool(
    name="json-chat",
    title="从JSON数据提取对话",
    description="从JSON数据中提取OpenAI格式的对话消息，返回结构化的对话记录",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
    }
)
async def extract_conversation_from_json(params: ExtractConversationFromJsonInput) -> ExtractConversationFromJsonOutput:
    """
    从JSON数据中提取对话消息

    专门用于处理OpenAI格式的对话JSON，返回结构化的消息列表。
    如果不是对话格式，is_conversation会返回false。

    Args:
        params: 包含JSON数据

    Returns:
        ExtractConversationFromJsonOutput: 结构化的对话消息
    """
    core = _get_core()

    try:
        messages = core.extract_conversation_from_json(params.json_data)
        is_conversation = len(messages) > 0

        if is_conversation:
            formatted_output = core._format_conversation(params.json_data)
        else:
            formatted_output = "该数据不是对话格式"

        return ExtractConversationFromJsonOutput(
            success=True,
            is_conversation=is_conversation,
            messages=[ConversationMessage(**msg) for msg in messages],
            formatted_output=formatted_output,
            count=len(messages)
        )

    except Exception as e:
        return ExtractConversationFromJsonOutput(
            success=False,
            error=str(e)
        )


@mcp_tool(
    name="json-chatfile",
    title="从JSON文件提取对话",
    description="从JSON文件中提取OpenAI格式的对话消息，返回结构化的对话记录",
    annotations={
        "readOnlyHint": True,
        "destructiveHint": False,
    }
)
async def extract_conversation_from_file(params: ExtractConversationFromFileInput) -> ExtractConversationFromFileOutput:
    """
    从JSON文件中提取对话消息

    专门用于处理OpenAI格式的对话JSON，返回结构化的消息列表。
    如果不是对话格式，is_conversation会返回false。

    Args:
        params: 包含文件路径

    Returns:
        ExtractConversationFromFileOutput: 结构化的对话消息
    """
    core = _get_core()

    try:
        import os
        if not os.path.exists(params.file_path):
            return ExtractConversationFromFileOutput(
                success=False,
                file_path=params.file_path,
                error=f"文件不存在: {params.file_path}"
            )

        messages = core.extract_conversation_from_file(params.file_path)
        is_conversation = len(messages) > 0

        if is_conversation:
            import json
            with open(params.file_path, 'r', encoding='utf-8') as f:
                json_data = json.load(f)
            formatted_output = core._format_conversation(json_data)
        else:
            formatted_output = "该文件不是对话格式"

        return ExtractConversationFromFileOutput(
            success=True,
            file_path=params.file_path,
            is_conversation=is_conversation,
            messages=[ConversationMessage(**msg) for msg in messages],
            formatted_output=formatted_output,
            count=len(messages)
        )

    except Exception as e:
        return ExtractConversationFromFileOutput(
            success=False,
            file_path=params.file_path,
            error=str(e)
        )