data-collector-mcp/data_collector_mcp/modbus_server.py

from __future__ import annotations

from typing import Any, NotRequired, TypedDict

from pydantic import BaseModel

from .collector_api import (
    create_modbus_devices as api_create_modbus_devices,
    create_modbus_points as api_create_modbus_points,
    edit_modbus_device as api_edit_modbus_device,
    edit_modbus_point as api_edit_modbus_point,
)
from .gateway_api import (
    modbus_point_collect_test as api_modbus_point_collect_test,
    modbus_raw_read as api_modbus_raw_read,
)
from .mcp_app import mcp


class ModbusDeviceCreateItem(BaseModel):
    name: str
    ip: str
    port: int
    slave_id: int
    byte_order: int
    word_order: int
    address_base: int
    serial_port: str = ""
    device_type: int = 1
    timeout: int = 3
    is_persistent: bool = False
    baud_rate: int = 0
    data_bit: int = 0
    parity: int = 0
    stop_bit: int = 0
    mode: int = 0
    retry_times: int = 0
    group_id: int = 0
    alarm_interval: int = 90
    collect_interval: int = 5


class ModbusPointCreateItem(BaseModel):
    device_id: int
    name: str
    address: int
    type: str
    func_code: int
    register_type: str = ""
    point_id: str = ""
    scale_ratio: float = 1
    value_offset: float = 0
    group_id: int = 0
    invalid_values: str = ""
    valid_range_start: float | None = None
    valid_range_end: float | None = None
    bit: int = 0
    describe: str = ""


class ModbusRawReadSpec(TypedDict):
    function_code: int
    address: int
    quantity: int


@mcp.tool(
    name="modbus.raw_read",
    description=(
        "采集网关-通过Modbus TCP 读取原始数据。"
        "device_type 设备类型，默认 ModbusTCP；word_byte_order 默认 ABCD，可选 ABCD、BADC、CDAB、DCBA；"
        "address_base 地址偏移，默认 0。"
        "read 必须包含 function_code、address、quantity；function_code: "
        "1=Read Coils/线圈，2=Read Discrete Inputs/离散输入，"
        "3=Read Holding Registers/保持寄存器，4=Read Input Registers/输入寄存器；"
        "quantity 范围 1..125。"
        "不要传name字段"
    ),
)
def modbus_raw_read(
    project_key: str,
    ip: str,
    port: int,
    slave_id: int,
    read: ModbusRawReadSpec,
    device_type: str = "ModbusTCP",
    word_byte_order: str = "ABCD",
    address_base: int = 0,
) -> dict[str, Any]:
    return api_modbus_raw_read(
        project_key,
        ip=ip,
        port=port,
        slave_id=slave_id,
        read=dict(read),
        device_type=device_type,
        word_byte_order=word_byte_order,
        address_base=address_base,
    )


@mcp.tool(
    name="modbus.point_collect_test",
    description=(
        "通过采集网关读取 Modbus TCP 点位的数据"
        "寄存器类型可选：1=Read Coils/线圈/0x，2=Read Discrete Inputs/离散输入/1x，"
        "3=Read Holding Registers/保持寄存器/4x，4=Read Input Registers/输入寄存器/3x。"
        "word_byte_order 可选 ABCD、BADC、CDAB、DCBA"
    ),
)
def modbus_point_collect_test(
    project_key: str,
    ip: str,
    port: int,
    slave_id: int,
    points: list[dict[str, Any]],
    device_type: str = "ModbusTCP",
    word_byte_order: str = "ABCD",
    address_base: int = 0,
) -> dict[str, Any]:
    return api_modbus_point_collect_test(
        project_key,
        ip=ip,
        port=port,
        slave_id=slave_id,
        points=points,
        device_type=device_type,
        word_byte_order=word_byte_order,
        address_base=address_base,
    )


@mcp.tool(
    name="collector.modbus_device_create",
    description=(
        "汇采-批量创建 Modbus 设备。"
        "默认参数: type=modbus, timeout=3, is_persistent=false, group_id=0, "
        "alarm_interval=90, collect_interval=5, retry_times=0。"
        "address_base 地址偏移/address_offset。"
        "byte_order: 1=Big Endian, 2=Small Endian。word_order: 1=Big Endian, 2=Small Endian。"
        "word_byte_order 映射: ABCD=>1/1, BADC=>2/1, CDAB=>1/2, DCBA=>2/2。"
    ),
)
def collector_modbus_device_create(
    project_key: str,
    devices: list[ModbusDeviceCreateItem],
) -> dict[str, Any]:
    return api_create_modbus_devices(project_key, [_dump_model_or_dict(item) for item in devices])


@mcp.tool(
    name="collector.modbus_device_edit",
    description=(
        "汇采-编辑 Modbus 设备。必须传 ori_id 原设备 id、name、slave_id、"
        "word_order、byte_order、ip 和 port；"
        "编辑前设备不能处于已连接状态；若已连接，请先调用 collector.device_disconnect。"
        "默认参数: ip='', port=0, serial_port='', timeout=3, is_persistent=false, "
        "baud_rate=0, data_bit=0, parity=0, stop_bit=0, mode=0, address_offset=0, "
        "retry_times=0, device_group_id=0, alarm_interval=90, collect_interval=5。"
        "byte_order: 1=Big Endian, 2=Small Endian；word_order: 1=Big Endian, 2=Small Endian。"
    ),
)
def collector_modbus_device_edit(
    project_key: str,
    ori_id: int,
    name: str,
    device_type: int,
    slave_id: int,
    byte_order: int,
    word_order: int,
    ip: str = "",
    port: int = 0,
    serial_port: str = "",
    timeout: int = 3,
    is_persistent: bool = False,
    baud_rate: int = 0,
    data_bit: int = 0,
    parity: int = 0,
    stop_bit: int = 0,
    mode: int = 0,
    address_offset: int = 0,
    retry_times: int = 0,
    device_group_id: int = 0,
    alarm_interval: int = 90,
    collect_interval: int = 5,
) -> dict[str, Any]:
    return api_edit_modbus_device(
        project_key,
        {
            "ori_id": ori_id,
            "name": name,
            "device_type": device_type,
            "ip": ip,
            "port": port,
            "slave_id": slave_id,
            "byte_order": byte_order,
            "word_order": word_order,
            "serial_port": serial_port,
            "timeout": timeout,
            "is_persistent": is_persistent,
            "baud_rate": baud_rate,
            "data_bit": data_bit,
            "parity": parity,
            "stop_bit": stop_bit,
            "mode": mode,
            "address_offset": address_offset,
            "retry_times": retry_times,
            "device_group_id": device_group_id,
            "alarm_interval": alarm_interval,
            "collect_interval": collect_interval,
        },
    )


@mcp.tool(
    name="collector.modbus_point_create",
    description=(
        "汇采-批量创建 Modbus 采集点位。"
        "points 每项必须传 device_id、name 名称、address 寄存器地址、"
        "type 数据类型，以及 func_code 寄存器类型。"
        "func_code: 1=Read Coils/线圈/0x，2=Read Discrete Inputs/离散输入/1x，"
        "3=Read Holding Registers/保持寄存器/4x，4=Read Input Registers/输入寄存器/3x。"
        "数据类型应使用汇采类型: bool, int16, uint16, int32, uint32, int64, uint64, float32, float64。"
        "常见点表类型映射: BOOL=>bool, SHORT=>int16, WORD=>uint16, LONG=>int32, "
        "DWORD=>uint32, FLOAT/REAL=>float32, DOUBLE=>float64, LONGLONG=>int64, QWORD=>uint64。"
    ),
)
def collector_modbus_point_create(
    project_key: str,
    points: list[ModbusPointCreateItem],
) -> dict[str, Any]:
    return api_create_modbus_points(project_key, [_dump_model_or_dict(item) for item in points])


def _dump_model_or_dict(item: Any) -> dict[str, Any]:
    if isinstance(item, BaseModel):
        return item.model_dump(exclude_none=True)
    return dict(item)


@mcp.tool(
    name="collector.modbus_point_edit",
    description=(
        "汇采-编辑 Modbus 采集点位。调用 "
        "必须传 ori_id 原点位 id、name 名称、address 寄存器地址、type 数据类型，"
        "以及 func_code 寄存器类型。"
        "ori_id 对应 collector.device_points 返回的 data.point[].id。"
        "默认参数: point_id='', scale_ratio=1, value_offset=0, group_id=0, "
        "invalid_values='', valid_range_start=null, valid_range_end=null, bit=0。"
        "func_code: 1=Read Coils/线圈/0x，2=Read Discrete Inputs/离散输入/1x，"
        "3=Read Holding Registers/保持寄存器/4x，4=Read Input Registers/输入寄存器/3x。"
        "register_type 可用 coil、discrete_input、holding_register、input_register。"
        "数据类型应使用汇采类型: bool, int16, uint16, int32, uint32, int64, uint64, float32, float64。"
        "常见点表类型映射: BOOL=>bool, SHORT=>int16, WORD=>uint16, LONG=>int32, "
        "DWORD=>uint32, FLOAT/REAL=>float32, DOUBLE=>float64, LONGLONG=>int64, QWORD=>uint64。"
    ),
)
def collector_modbus_point_edit(
    project_key: str,
    ori_id: int,
    name: str,
    address: int,
    data_type: str,
    func_code: int = 0,
    register_type: str = "",
    point_id: str = "",
    scale_ratio: float = 1,
    value_offset: float = 0,
    group_id: int = 0,
    invalid_values: str = "",
    valid_range_start: float | None = None,
    valid_range_end: float | None = None,
    bit: int = 0,
    describe: str = "",
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "ori_id": ori_id,
        "name": name,
        "address": address,
        "type": data_type,
        "point_id": point_id,
        "scale_ratio": scale_ratio,
        "value_offset": value_offset,
        "group_id": group_id,
        "invalid_values": invalid_values,
        "valid_range_start": valid_range_start,
        "valid_range_end": valid_range_end,
        "bit": bit,
        "describe": describe,
    }
    if func_code:
        payload["func_code"] = func_code
    else:
        payload["register_type"] = register_type
    return api_edit_modbus_point(project_key, payload)