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="项目列表",
    description="列出目前用户可查询的项目列表。执行其他工具前，如前置未提到要查询项目的project_key，应先让用户选择要查询的项目，再使用对应的 project_key 调用后续工具。",
    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:
    """搜索AI控制系统并返回系统code。默认 page_size 为 20，除非显式传入其他值。"""
    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:
    """按AI控制系统的code查询策略记录。默认 page_size 为 10，除非显式传入其他值。auto_exec=true 表示自动控制，否则表示推荐参考。"""
    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:
    """按AI控制系统的code查询控制模式和运行状态。control_mode=1 表示自动控制；status=0 表示正常，1 表示异常。"""
    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:
    """按 point_id 和时间范围查询指令下发日志。返回结果中 status=1 表示该点位下发成功。"""
    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()