data-collector-mcp/data_collector_mcp/modbus_server.py

from __future__ import annotations

from typing import Annotated, Any

from pydantic import BaseModel, ConfigDict, Field

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


JsonScalar = str | int | float | bool | None

ProjectKeyParam = Annotated[str, Field(description="项目标识，来自 project.list 返回的 project_key。")]
ModbusIpParam = Annotated[str, Field(description="Modbus TCP 设备地址。")]
ModbusPortParam = Annotated[int, Field(description="Modbus TCP 端口。")]
ModbusSlaveIdParam = Annotated[int, Field(description="Modbus 从站 ID。")]
ModbusDeviceTypeParam = Annotated[str, Field(description="网关设备类型，默认 ModbusTCP。")]
ModbusWordByteOrderParam = Annotated[str, Field(description="字节/字顺序，可选 ABCD、BADC、CDAB、DCBA。")]
ModbusAddressBaseParam = Annotated[int, Field(description="地址基准或地址偏移，默认 0。")]
ModbusFunctionCodeParam = Annotated[int, Field(description="功能码：1=线圈，2=离散输入，3=保持寄存器，4=输入寄存器。")]
ModbusDataTypeParam = Annotated[
    str,
    Field(description="数据类型；支持 bool、int16、uint16、int32、uint32、int64、uint64、float32、float64 及常见别名。"),
]


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


class ModbusDeviceCreateItem(BaseModel):
    name: str = Field(description="设备名称。")
    ip: str = Field(description="Modbus TCP/UDP 设备地址；RTU 设备可为空。")
    port: int = Field(description="Modbus TCP/UDP 端口；RTU 设备可为 0。")
    slave_id: int = Field(description="Modbus 从站 ID。")
    byte_order: int = Field(description="字节顺序：1=Big Endian，2=Small Endian。")
    word_order: int = Field(description="字顺序：1=Big Endian，2=Small Endian。")
    address_base: int = Field(description="地址基准；会转换为汇采接口 address_offset。")
    serial_port: str = Field(default="", description="RTU 串口名。")
    device_type: int = Field(default=1, description="协议类型：1=TCP，2=RTU，3=UDP，4=RTU OVER TCP，5=RTU OVER UDP。")
    timeout: int = Field(default=3, description="连接超时时间。")
    is_persistent: bool = Field(default=False, description="是否持久化连接。")
    baud_rate: int = Field(default=0, description="RTU 波特率。")
    data_bit: int = Field(default=0, description="RTU 数据位。")
    parity: int = Field(default=0, description="RTU 校验位。")
    stop_bit: int = Field(default=0, description="RTU 停止位。")
    mode: int = Field(default=0, description="连接模式。")
    retry_times: int = Field(default=0, description="重试次数。")
    group_id: int = Field(default=0, description="设备分组 id。")
    alarm_interval: int = Field(default=90, description="告警间隔，单位秒。")
    collect_interval: int = Field(default=5, description="采集周期，单位秒。")


class ModbusPointCreateItem(BaseModel):
    device_id: int = Field(description="所属 Modbus 设备 id，来自 collector.device_list 或设备创建结果。")
    name: str = Field(description="点位名称。")
    address: int = Field(description="寄存器地址。")
    type: ModbusDataTypeParam
    func_code: int | None = Field(default=3, description="功能码或寄存器类型：1=Read Coils/线圈/0x，2=Read Discrete Inputs/离散输入/1x，3=Read Holding Registers/保持寄存器/4x，4=Read Input Registers/输入寄存器/3x。")
    register_type: str = Field(default="", description="寄存器类型；可用 coil、discrete_input、holding_register、input_register。")
    point_id: str = Field(default="", description="外部点位编码。")
    scale_ratio: float = Field(default=1, description="缩放系数。")
    value_offset: float = Field(default=0, description="值偏移。")
    group_id: int = Field(default=0, description="点位分组 id。")
    invalid_values: str = Field(default="", description="无效值列表，多个值用逗号分隔。")
    valid_range_start: float | None = Field(default=None, description="合法范围最小值。")
    valid_range_end: float | None = Field(default=None, description="合法范围最大值。")
    bit: int = Field(default=0, description="位索引；用于从寄存器中取指定 bit。")
    describe: str = Field(default="", description="点位描述。")


class ModbusRawReadSpec(BaseModel):
    function_code: ModbusFunctionCodeParam
    address: int = Field(description="起始地址。")
    quantity: int = Field(description="读取数量，范围 1..125。")


class ModbusPointCollectSpec(FlexibleModel):
    function_code: int | None = Field(default=None, description="功能码或寄存器类型：1=Read Coils/线圈/0x，2=Read Discrete Inputs/离散输入/1x，3=Read Holding Registers/保持寄存器/4x，4=Read Input Registers/输入寄存器/3x。")
    func_code: int | None = Field(default=None, description="功能码或寄存器类型：1=Read Coils/线圈/0x，2=Read Discrete Inputs/离散输入/1x，3=Read Holding Registers/保持寄存器/4x，4=Read Input Registers/输入寄存器/3x。")
    address: int = Field(description="寄存器地址。")
    type: ModbusDataTypeParam
    bit: int | None = Field(default=None, description="位索引；用于读取 bool 位。")
    name: str | None = Field(default=None, description="点位名称；测试读取时可不传。")


class GatewayBaseOutput(FlexibleModel):
    code: int | str | None = Field(default=None, description="网关状态码；0 表示成功。")
    msg: str | None = Field(default=None, description="网关状态说明或错误信息。")


class GatewayCommunicationData(FlexibleModel):
    communication: list[str] = Field(default_factory=list, description="本次请求期间捕获的通信报文或通信过程记录。")


class ModbusRawReadOutput(GatewayBaseOutput):
    data: GatewayCommunicationData | None = Field(default=None, description="Modbus 原始读取结果。")


class ModbusGatewayPoint(FlexibleModel):
    name: str | None = Field(default=None, description="点位名称。")
    address: int | None = Field(default=None, description="寄存器地址。")
    type: str | None = Field(default=None, description="数据类型。")
    value: JsonScalar = Field(default=None, description="转换后的点位值。")


class ModbusPointCollectData(GatewayCommunicationData):
    points: list[ModbusGatewayPoint] = Field(default_factory=list, description="转换后的 Modbus 点位值列表。")


class ModbusPointCollectOutput(GatewayBaseOutput):
    data: ModbusPointCollectData | None = Field(default=None, description="Modbus 点位读取结果。")


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


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


class BatchError(FlexibleModel):
    index: int | None = Field(default=None, description="输入数组中的下标。")
    name: str | None = Field(default=None, description="输入项名称。")
    stage: str | None = Field(default=None, description="失败阶段。")
    error: str | None = Field(default=None, description="失败原因。")


class DeviceCreateSummary(FlexibleModel):
    total: int | None = Field(default=None, description="输入设备总数。")
    created: int | None = Field(default=None, description="汇采创建设备成功数量。")
    matched: int | None = Field(default=None, description="创建后从设备列表匹配到设备 id 的数量。")
    failed: int | None = Field(default=None, description="失败数量。")


class DeviceCreateResult(FlexibleModel):
    index: int | None = Field(default=None, description="输入 devices 数组中的下标。")
    name: str | None = Field(default=None, description="设备名称。")
    device_id: int | None = Field(default=None, description="匹配到的汇采设备 id。")
    create_response: dict[str, Any] | None = Field(default=None, description="汇采创建设备接口原始响应。")
    matched_device: dict[str, Any] | None = Field(default=None, description="设备列表中匹配到的设备详情。")


class ModbusDeviceCreateOutput(FlexibleModel):
    state: int | str | None = Field(default=None, description="批量创建状态；0 表示全部成功，1 表示存在失败。")
    summary: DeviceCreateSummary | None = Field(default=None, description="批量创建设备汇总。")
    results: list[DeviceCreateResult] = Field(default_factory=list, description="每个设备的创建和匹配结果。")
    errors: list[BatchError] = Field(default_factory=list, description="失败项列表。")


class PointCreateSummary(FlexibleModel):
    total: int | None = Field(default=None, description="输入点位总数。")
    success: int | None = Field(default=None, description="汇采创建点位成功数量。")
    failed: int | None = Field(default=None, description="失败数量。")


class PointCreateResult(FlexibleModel):
    index: int | None = Field(default=None, description="输入 points 数组中的下标。")
    name: str | None = Field(default=None, description="点位名称。")
    device_id: int | None = Field(default=None, description="所属 Modbus 设备 id。")
    response: dict[str, Any] | None = Field(default=None, description="汇采创建点位接口原始响应。")


class ModbusPointCreateOutput(FlexibleModel):
    state: int | str | None = Field(default=None, description="批量创建状态；0 表示全部成功，1 表示存在失败。")
    summary: PointCreateSummary | None = Field(default=None, description="批量创建点位汇总。")
    results: list[PointCreateResult] = Field(default_factory=list, description="每个点位的创建结果。")
    errors: list[BatchError] = Field(default_factory=list, description="失败项列表。")


@mcp.tool(
    name="modbus.raw_read",
    description=(
        "采集网关-通过 Modbus TCP 读取原始数据，只返回通信报文，不解析业务值。"
        "如需转换后的点位值，请使用 modbus.point_collect_test。"
    ),
)
def modbus_raw_read(
    project_key: ProjectKeyParam,
    ip: ModbusIpParam,
    port: ModbusPortParam,
    slave_id: ModbusSlaveIdParam,
    read: Annotated[ModbusRawReadSpec, Field(description="原始读取参数。")],
    device_type: ModbusDeviceTypeParam = "ModbusTCP",
    word_byte_order: ModbusWordByteOrderParam = "ABCD",
    address_base: ModbusAddressBaseParam = 0,
) -> ModbusRawReadOutput:
    return api_modbus_raw_read(
        project_key,
        ip=ip,
        port=port,
        slave_id=slave_id,
        read=_dump_model_or_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 点位并转换为业务值。"
    ),
)
def modbus_point_collect_test(
    project_key: ProjectKeyParam,
    ip: ModbusIpParam,
    port: ModbusPortParam,
    slave_id: ModbusSlaveIdParam,
    points: Annotated[list[ModbusPointCollectSpec], Field(description="要读取的 Modbus 点位列表。")],
    device_type: ModbusDeviceTypeParam = "ModbusTCP",
    word_byte_order: ModbusWordByteOrderParam = "ABCD",
    address_base: ModbusAddressBaseParam = 0,
) -> ModbusPointCollectOutput:
    return api_modbus_point_collect_test(
        project_key,
        ip=ip,
        port=port,
        slave_id=slave_id,
        points=[_dump_model_or_dict(item) for item in points],
        device_type=device_type,
        word_byte_order=word_byte_order,
        address_base=address_base,
    )


@mcp.tool(
    name="collector.modbus_device_create",
    description=(
        "汇采-批量创建 Modbus 设备。创建后会读取设备列表并匹配返回设备 id。"
        "word_byte_order 映射到 byte_order/word_order：ABCD=>1/1，BADC=>2/1，CDAB=>1/2，DCBA=>2/2。"
    ),
)
def collector_modbus_device_create(
    project_key: ProjectKeyParam,
    devices: Annotated[list[ModbusDeviceCreateItem], Field(description="要创建的 Modbus 设备列表。")],
) -> ModbusDeviceCreateOutput:
    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 设备。编辑前设备不能处于已连接状态；"
        "若已连接，请先调用 collector.device_disconnect。"
    ),
)
def collector_modbus_device_edit(
    project_key: ProjectKeyParam,
    id: Annotated[int, Field(description="要编辑的原汇采设备 id。")],
    name: Annotated[str, Field(description="设备名称。")],
    device_type: Annotated[int, Field(description="协议类型：1=TCP，2=RTU，3=UDP，4=RTU OVER TCP，5=RTU OVER UDP。")],
    slave_id: ModbusSlaveIdParam,
    byte_order: Annotated[int, Field(description="字节顺序：1=Big Endian，2=Small Endian。")],
    word_order: Annotated[int, Field(description="字顺序：1=Big Endian，2=Small Endian。")],
    ip: Annotated[str, Field(description="Modbus TCP/UDP 设备地址；RTU 设备可为空。")]= "",
    port: Annotated[int, Field(description="Modbus TCP/UDP 端口；RTU 设备可为 0。")]= 0,
    serial_port: Annotated[str, Field(description="RTU 串口名。")]= "",
    timeout: Annotated[int, Field(description="连接超时时间。")]= 3,
    is_persistent: Annotated[bool, Field(description="是否持久化连接。")]= False,
    baud_rate: Annotated[int, Field(description="RTU 波特率。")]= 0,
    data_bit: Annotated[int, Field(description="RTU 数据位。")]= 0,
    parity: Annotated[int, Field(description="RTU 校验位。")]= 0,
    stop_bit: Annotated[int, Field(description="RTU 停止位。")]= 0,
    mode: Annotated[int, Field(description="连接模式。")]= 0,
    address_offset: Annotated[int, Field(description="地址偏移。")]= 0,
    retry_times: Annotated[int, Field(description="重试次数。")]= 0,
    device_group_id: Annotated[int, Field(description="设备分组 id。")]= 0,
    alarm_interval: Annotated[int, Field(description="告警间隔，单位秒。")]= 90,
    collect_interval: Annotated[int, Field(description="采集周期，单位秒。")]= 5,
) -> CollectorBaseOutput:
    return api_edit_modbus_device(
        project_key,
        {
            "ori_id": 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 采集点位。数据类型和寄存器类型会在调用汇采接口前规范化。"
    ),
)
def collector_modbus_point_create(
    project_key: ProjectKeyParam,
    points: Annotated[list[ModbusPointCreateItem], Field(description="要创建的 Modbus 采集点位列表。")],
) -> ModbusPointCreateOutput:
    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 采集点位。数据类型和寄存器类型会在调用汇采接口前规范化。"
    ),
)
def collector_modbus_point_edit(
    project_key: ProjectKeyParam,
    id: Annotated[int, Field(description="要编辑的原采集点位 id，对应 collector.device_points 返回的点位 id。")],
    name: Annotated[str, Field(description="点位名称。")],
    address: Annotated[int, Field(description="寄存器地址。")],
    data_type: ModbusDataTypeParam,
    func_code: Annotated[int, Field(description="功能码或寄存器类型：1=Read Coils/线圈/0x，2=Read Discrete Inputs/离散输入/1x，3=Read Holding Registers/保持寄存器/4x，4=Read Input Registers/输入寄存器/3x。")],
    point_id: Annotated[str, Field(description="外部点位编码。")]= "",
    scale_ratio: Annotated[float, Field(description="缩放系数。")]= 1,
    value_offset: Annotated[float, Field(description="值偏移。")]= 0,
    group_id: Annotated[int, Field(description="点位分组 id。")]= 0,
    invalid_values: Annotated[str, Field(description="无效值列表，多个值用逗号分隔。")]= "",
    valid_range_start: Annotated[float | None, Field(description="合法范围最小值。")]= None,
    valid_range_end: Annotated[float | None, Field(description="合法范围最大值。")]= None,
    bit: Annotated[int, Field(description="位索引；用于从寄存器中取指定 bit。")]= 0,
    describe: Annotated[str, Field(description="点位描述。")]= "",
) -> CollectorPointEditOutput:
    payload: dict[str, Any] = {
        "ori_id": 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,
    }
    payload["func_code"] = func_code
    return api_edit_modbus_point(project_key, payload)