data-collector-mcp/data_collector_mcp/common_server.py

from __future__ import annotations

from typing import Annotated, Any

from pydantic import BaseModel, ConfigDict, Field

from .auth import load_projects_config
from .collector_api import (
    connect_device as api_connect_device,
    disconnect_device as api_disconnect_device,
    list_device_points as api_list_device_points,
    list_devices as api_list_devices,
)
from .mcp_app import mcp


JsonScalar = str | int | float | bool | None

ProjectKeyParam = Annotated[str, Field(description="项目标识，来自 project.list 返回的 project_key。")]
DeviceIdParam = Annotated[int, Field(description="汇采设备 id，来自 collector.device_list 返回的设备对象。")]
DeviceTypeParam = Annotated[
    str,
    Field(description="设备协议类型，可传 modbus、s7、bacnet、ethernet-ip、opc-ua、opc-da、snmp、iec104。"),
]
DeviceGroupIdParam = Annotated[int, Field(description="设备点位分组 id，默认 0 表示查询设备下默认分组。")]


class FlexibleModel(BaseModel):
    model_config = ConfigDict(extra="allow")


class ProjectListItem(BaseModel):
    project_key: str = Field(description="项目标识，用于其他采集工具的 project_key 参数。")
    project_name: str = Field(description="项目显示名称。")


class ProjectListOutput(BaseModel):
    projects: list[ProjectListItem] = Field(description="当前 MCP 服务可用的已启用项目列表。")
    total: int = Field(description="可用项目数量。")


class CollectorBaseOutput(FlexibleModel):
    state: int | str | None = Field(default=None, description="汇采状态码；0 表示成功。")
    state_info: str | None = Field(default=None, description="汇采状态说明。")


class CollectorDeviceItem(FlexibleModel):
    id: int | None = Field(default=None, description="设备 id，可用于连接、断开和查询点位。")
    name: str | None = Field(default=None, description="设备名称。")
    type: str | None = Field(default=None, description="设备类型；devicegroup 表示设备分组。")
    device_type: int | str | None = Field(default=None, description="设备子类型或协议内设备类型。")
    ip: str | None = Field(default=None, description="设备地址。")
    port: int | None = Field(default=None, description="设备端口。")
    num_points: int | None = Field(default=None, description="设备点位数量；num_points=true 时返回。")
    status: int | None = Field(default=None, description="连接状态：1=未连接，2=已连接，3=连接异常。")
    running_status: int | None = Field(default=None, description="运行状态：0=未采集，1=采集中，2=采集异常。")
    group_id: int | None = Field(default=None, description="所属设备分组 id。")
    groups: list["CollectorDeviceItem"] | None = Field(default=None, description="设备分组下的设备列表。")


class CollectorDeviceListOutput(CollectorBaseOutput):
    devices: list[CollectorDeviceItem] = Field(default_factory=list, description="设备列表；type=devicegroup 时继续检查 groups。")


class CollectorDeviceConnectData(FlexibleModel):
    status: int | None = Field(default=None, description="操作后的设备连接状态：1=未连接，2=已连接，3=连接异常。")
    running_status: int | None = Field(default=None, description="操作后的设备运行状态：0=未采集，1=采集中，2=采集异常。")
    msg: str | None = Field(default=None, description="连接或断开失败时的错误信息；成功时通常为空。")


class CollectorDeviceConnectOutput(CollectorBaseOutput):
    data: CollectorDeviceConnectData | None = Field(default=None, description="设备连接或断开后的状态。")


class CollectorDevicePointItem(FlexibleModel):
    id: int | None = Field(default=None, description="采集点位 id，可用于对应协议的点位编辑工具。")
    node_id: str | None = Field(default=None, description="采集节点 id。")
    name: str | None = Field(default=None, description="点位名称。")
    point_id: str | None = Field(default=None, description="外部点位编码。")
    data_type: str | None = Field(default=None, description="数据类型。")
    present_value: JsonScalar = Field(default=None, description="当前内存中的点位最新值。")
    status: int | None = Field(default=None, description="点位状态：0=未采集，1=采集正常，2=采集异常。")
    group_id: int | None = Field(default=None, description="所属点位分组 id。")
    created_time: str | None = Field(default=None, description="创建时间。")
    updated_time: str | None = Field(default=None, description="更新时间。")


class CollectorDevicePointsData(FlexibleModel):
    point: list[CollectorDevicePointItem] = Field(default_factory=list, description="设备下的采集点位列表。")
    total: int | None = Field(default=None, description="点位总数。")


class CollectorDevicePointsOutput(CollectorBaseOutput):
    data: CollectorDevicePointsData | None = Field(default=None, description="设备点位查询结果。")


@mcp.tool(
    name="project.list",
    title="Project List",
    description=(
        "列出当前 MCP 服务可用的已启用采集项目。使用其他采集工具前，应先调用本工具选择 project_key。"
    ),
    tags={"project", "list"},
)
def project_list() -> ProjectListOutput:
    projects = load_projects_config()
    result = [
        {
            "project_key": item["project_key"],
            "project_name": item["project_name"],
        }
        for item in projects
        if item["enabled"]
    ]
    result.sort(key=lambda item: item["project_key"])
    return {"projects": result, "total": len(result)}


@mcp.tool(
    name="collector.device_list",
    description=(
        "汇采-返回目前所有设备及其详情信息。type=devicegroup 时表示设备分组，需要继续检查 groups 下的设备。"
    ),
)
def collector_device_list(
    project_key: ProjectKeyParam,
    num_points: Annotated[bool, Field(description="是否返回设备点位数量；默认 false。")]= False,
) -> CollectorDeviceListOutput:
    return api_list_devices(project_key, num_points=num_points)


@mcp.tool(
    name="collector.device_connect",
    description=(
        "汇采-连接设备。连接成功不代表正在采集，设备采集状态请看响应 data.running_status。"
    ),
)
def collector_device_connect(
    project_key: ProjectKeyParam,
    device_id: DeviceIdParam,
    device_type: DeviceTypeParam,
) -> CollectorDeviceConnectOutput:
    return api_connect_device(project_key, device_id=device_id, device_type=device_type)


@mcp.tool(
    name="collector.device_disconnect",
    description=(
        "汇采-断开设备。用于停止指定设备连接或采集相关状态，使设备回到未连接或空闲状态。"
    ),
)
def collector_device_disconnect(
    project_key: ProjectKeyParam,
    device_id: DeviceIdParam,
    device_type: DeviceTypeParam = "modbus",
) -> CollectorDeviceConnectOutput:
    return api_disconnect_device(project_key, device_id=device_id, device_type=device_type)


@mcp.tool(
    name="collector.device_points",
    description=(
        "汇采-查询设备下的所有采集点位。返回的点位 id 可用于对应协议的点位编辑工具。"
    ),
)
def collector_device_points(
    project_key: ProjectKeyParam,
    device_id: DeviceIdParam,
    device_type: DeviceTypeParam,
    group_id: DeviceGroupIdParam = 0,
) -> CollectorDevicePointsOutput:
    return api_list_device_points(
        project_key,
        device_id=device_id,
        device_type=device_type,
        group_id=group_id,
    )