m2-mcp/m2_mcp/server.py

from __future__ import annotations

import os
from typing import Any

from fastmcp import FastMCP

from .auth import load_projects_config
from .config_api import (
    get_command_log as api_get_command_log,
    get_ai_online_v2 as api_get_ai_online_v2,
    search_ai_rcmd_operations as api_search_ai_rcmd_operations,
    search_ai_systems as api_search_ai_systems,
)


mcp = FastMCP("m2")


@mcp.tool(
    name="project.list",
    title="Project List",
    description="List enabled projects configured in sys_config for m2 tools. 在执行其他工具前，询问用户要查询哪一个项目，根据用户的选择查询对应项目。",
    tags={"project", "list"},
)
def project_list() -> dict[str, Any]:
    projects = load_projects_config()

    result: list[dict[str, Any]] = []
    for item in projects:
        if not item["enabled"]:
            continue
        result.append(
            {
                "project_key": item["project_key"],
                "project_name": item["project_name"],
            }
        )

    result.sort(key=lambda item: item["project_key"])
    return {
        "projects": result,
        "total": len(result),
    }


def _append_next_page_hint(payload: Any, page_num: int) -> Any:
    if not isinstance(payload, dict):
        return payload

    total_page = None
    data = payload.get("data")
    if isinstance(data, dict):
        total_page = data.get("total_page")
    if total_page is None:
        total_page = payload.get("total_pages")
    if total_page is None:
        total_page = payload.get("total_page")

    if isinstance(total_page, int) and total_page > page_num:
        payload.setdefault(
            "mcp_note",
            f"Current result is page {page_num}. If the target was not found, continue with page_num={page_num + 1}.",
        )
    return payload


@mcp.tool()
def search_ai_systems(
    project_key: str,
    keyword: str = "",
    page_size: int = 20,
    page_num: int = 1,
    order_by: list[str] | None = None,
) -> Any:
    """Search AI systems and return system codes. Default page_size is 20 unless explicitly overridden."""
    payload = api_search_ai_systems(
        project_key,
        keyword=keyword,
        page_size=page_size,
        page_num=page_num,
        order_by=order_by,
    )
    return _append_next_page_hint(payload, page_num)


@mcp.tool()
def search_ai_rcmd_operations(
    project_key: str,
    codes: list[str],
    end: str,
    page_size: int = 10,
    page_num: int = 1,
    order: str = "-create_time",
) -> Any:
    """Search AI strategy records by AI system code. Default page_size is 10 unless explicitly overridden. auto_exec=true means auto control, otherwise recommendation only."""
    payload = api_search_ai_rcmd_operations(
        project_key,
        codes=codes,
        end=end,
        page_size=page_size,
        page_num=page_num,
        order=order,
    )
    return _append_next_page_hint(payload, page_num)


@mcp.tool()
def get_ai_online_v2(project_key: str, codes: list[str]) -> Any:
    """Get AI control mode and status by AI system code. control_mode=1 means auto control; status=0 means normal and 1 means abnormal."""
    return api_get_ai_online_v2(project_key, codes=codes)


@mcp.tool()
def get_command_log(
    project_key: str,
    point_id: str,
    begin: int,
    end: int,
    page_size: int = 20,
    page_num: int = 1,
    export: bool = False,
) -> Any:
    """Query command delivery logs by point_id and time range. In returned data, status=1 means the point was delivered successfully."""
    payload = api_get_command_log(
        project_key,
        point_id=point_id,
        begin=begin,
        end=end,
        page_size=page_size,
        page_num=page_num,
        export=export,
    )
    return _append_next_page_hint(payload, page_num)


def main() -> None:
    host = os.getenv("MCP_HOST", "0.0.0.0").strip() or "0.0.0.0"
    port = int(os.getenv("MCP_PORT", "8500"))
    path = os.getenv("MCP_PATH", "/mcp").strip() or "/mcp"
    mcp.run(transport="http", host=host, port=port, path=path)


if __name__ == "__main__":
    main()