mirror of
https://github.com/dptech-corp/Uni-Lab-OS.git
synced 2025-12-17 04:51:10 +00:00
357 lines
12 KiB
Python
357 lines
12 KiB
Python
"""
|
||
JSON 工作流转换模块
|
||
|
||
提供从多种 JSON 格式转换为统一工作流格式的功能。
|
||
支持的格式：
|
||
1. workflow/reagent 格式
|
||
2. steps_info/labware_info 格式
|
||
"""
|
||
|
||
import json
|
||
from os import PathLike
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Set, Tuple, Union
|
||
|
||
from unilabos.workflow.common import WorkflowGraph, build_protocol_graph
|
||
from unilabos.registry.registry import lab_registry
|
||
|
||
|
||
def get_action_handles(resource_name: str, template_name: str) -> Dict[str, List[str]]:
|
||
"""
|
||
从 registry 获取指定设备和动作的 handles 配置
|
||
|
||
Args:
|
||
resource_name: 设备资源名称，如 "liquid_handler.prcxi"
|
||
template_name: 动作模板名称，如 "transfer_liquid"
|
||
|
||
Returns:
|
||
包含 source 和 target handler_keys 的字典:
|
||
{"source": ["sources_out", "targets_out", ...], "target": ["sources", "targets", ...]}
|
||
"""
|
||
result = {"source": [], "target": []}
|
||
|
||
device_info = lab_registry.device_type_registry.get(resource_name, {})
|
||
if not device_info:
|
||
return result
|
||
|
||
action_mappings = device_info.get("class", {}).get("action_value_mappings", {})
|
||
action_config = action_mappings.get(template_name, {})
|
||
handles = action_config.get("handles", {})
|
||
|
||
if isinstance(handles, dict):
|
||
# 处理 input handles (作为 target)
|
||
for handle in handles.get("input", []):
|
||
handler_key = handle.get("handler_key", "")
|
||
if handler_key:
|
||
result["source"].append(handler_key)
|
||
# 处理 output handles (作为 source)
|
||
for handle in handles.get("output", []):
|
||
handler_key = handle.get("handler_key", "")
|
||
if handler_key:
|
||
result["target"].append(handler_key)
|
||
|
||
return result
|
||
|
||
|
||
def validate_workflow_handles(graph: WorkflowGraph) -> Tuple[bool, List[str]]:
|
||
"""
|
||
校验工作流图中所有边的句柄配置是否正确
|
||
|
||
Args:
|
||
graph: 工作流图对象
|
||
|
||
Returns:
|
||
(is_valid, errors): 是否有效，错误信息列表
|
||
"""
|
||
errors = []
|
||
nodes = graph.nodes
|
||
|
||
for edge in graph.edges:
|
||
left_uuid = edge.get("source")
|
||
right_uuid = edge.get("target")
|
||
# target_handle_key是target, right的输入节点（入节点）
|
||
# source_handle_key是source, left的输出节点（出节点）
|
||
right_source_conn_key = edge.get("target_handle_key", "")
|
||
left_target_conn_key = edge.get("source_handle_key", "")
|
||
|
||
# 获取源节点和目标节点信息
|
||
left_node = nodes.get(left_uuid, {})
|
||
right_node = nodes.get(right_uuid, {})
|
||
|
||
left_res_name = left_node.get("resource_name", "")
|
||
left_template_name = left_node.get("template_name", "")
|
||
right_res_name = right_node.get("resource_name", "")
|
||
right_template_name = right_node.get("template_name", "")
|
||
|
||
# 获取源节点的 output handles
|
||
left_node_handles = get_action_handles(left_res_name, left_template_name)
|
||
target_valid_keys = left_node_handles.get("target", [])
|
||
target_valid_keys.append("ready")
|
||
|
||
# 获取目标节点的 input handles
|
||
right_node_handles = get_action_handles(right_res_name, right_template_name)
|
||
source_valid_keys = right_node_handles.get("source", [])
|
||
source_valid_keys.append("ready")
|
||
|
||
# 如果节点配置了 output handles，则 source_port 必须有效
|
||
if not right_source_conn_key:
|
||
node_name = left_node.get("name", left_uuid[:8])
|
||
errors.append(f"源节点 '{node_name}' 的 source_handle_key 为空，" f"应设置为: {source_valid_keys}")
|
||
elif right_source_conn_key not in source_valid_keys:
|
||
node_name = left_node.get("name", left_uuid[:8])
|
||
errors.append(
|
||
f"源节点 '{node_name}' 的 source 端点 '{right_source_conn_key}' 不存在，" f"支持的端点: {source_valid_keys}"
|
||
)
|
||
|
||
# 如果节点配置了 input handles，则 target_port 必须有效
|
||
if not left_target_conn_key:
|
||
node_name = right_node.get("name", right_uuid[:8])
|
||
errors.append(f"目标节点 '{node_name}' 的 target_handle_key 为空，" f"应设置为: {target_valid_keys}")
|
||
elif left_target_conn_key not in target_valid_keys:
|
||
node_name = right_node.get("name", right_uuid[:8])
|
||
errors.append(
|
||
f"目标节点 '{node_name}' 的 target 端点 '{left_target_conn_key}' 不存在，"
|
||
f"支持的端点: {target_valid_keys}"
|
||
)
|
||
|
||
return len(errors) == 0, errors
|
||
|
||
|
||
# action 到 resource_name 的映射
|
||
ACTION_RESOURCE_MAPPING: Dict[str, str] = {
|
||
# 生物实验操作
|
||
"transfer_liquid": "liquid_handler.prcxi",
|
||
"transfer": "liquid_handler.prcxi",
|
||
"incubation": "incubator.prcxi",
|
||
"move_labware": "labware_mover.prcxi",
|
||
"oscillation": "shaker.prcxi",
|
||
# 有机化学操作
|
||
"HeatChillToTemp": "heatchill.chemputer",
|
||
"StopHeatChill": "heatchill.chemputer",
|
||
"StartHeatChill": "heatchill.chemputer",
|
||
"HeatChill": "heatchill.chemputer",
|
||
"Dissolve": "stirrer.chemputer",
|
||
"Transfer": "liquid_handler.chemputer",
|
||
"Evaporate": "rotavap.chemputer",
|
||
"Recrystallize": "reactor.chemputer",
|
||
"Filter": "filter.chemputer",
|
||
"Dry": "dryer.chemputer",
|
||
"Add": "liquid_handler.chemputer",
|
||
}
|
||
|
||
|
||
def normalize_steps(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||
"""
|
||
将不同格式的步骤数据规范化为统一格式
|
||
|
||
支持的输入格式：
|
||
- action + parameters
|
||
- action + action_args
|
||
- operation + parameters
|
||
|
||
Args:
|
||
data: 原始步骤数据列表
|
||
|
||
Returns:
|
||
规范化后的步骤列表，格式为 [{"action": str, "parameters": dict, "description": str?, "step_number": int?}, ...]
|
||
"""
|
||
normalized = []
|
||
for idx, step in enumerate(data):
|
||
# 获取动作名称（支持 action 或 operation 字段）
|
||
action = step.get("action") or step.get("operation")
|
||
if not action:
|
||
continue
|
||
|
||
# 获取参数（支持 parameters 或 action_args 字段）
|
||
raw_params = step.get("parameters") or step.get("action_args") or {}
|
||
params = dict(raw_params)
|
||
|
||
# 规范化 source/target -> sources/targets
|
||
if "source" in raw_params and "sources" not in raw_params:
|
||
params["sources"] = raw_params["source"]
|
||
if "target" in raw_params and "targets" not in raw_params:
|
||
params["targets"] = raw_params["target"]
|
||
|
||
# 获取描述（支持 description 或 purpose 字段）
|
||
description = step.get("description") or step.get("purpose")
|
||
|
||
# 获取步骤编号（优先使用原始数据中的 step_number，否则使用索引+1）
|
||
step_number = step.get("step_number", idx + 1)
|
||
|
||
step_dict = {"action": action, "parameters": params, "step_number": step_number}
|
||
if description:
|
||
step_dict["description"] = description
|
||
|
||
normalized.append(step_dict)
|
||
|
||
return normalized
|
||
|
||
|
||
def normalize_labware(data: List[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
|
||
"""
|
||
将不同格式的 labware 数据规范化为统一的字典格式
|
||
|
||
支持的输入格式：
|
||
- reagent_name + material_name + positions
|
||
- name + labware + slot
|
||
|
||
Args:
|
||
data: 原始 labware 数据列表
|
||
|
||
Returns:
|
||
规范化后的 labware 字典，格式为 {name: {"slot": int, "labware": str, "well": list, "type": str, "role": str, "name": str}, ...}
|
||
"""
|
||
labware = {}
|
||
for item in data:
|
||
# 获取 key 名称（优先使用 reagent_name，其次是 material_name 或 name）
|
||
reagent_name = item.get("reagent_name")
|
||
key = reagent_name or item.get("material_name") or item.get("name")
|
||
if not key:
|
||
continue
|
||
|
||
key = str(key)
|
||
|
||
# 处理重复 key，自动添加后缀
|
||
idx = 1
|
||
original_key = key
|
||
while key in labware:
|
||
idx += 1
|
||
key = f"{original_key}_{idx}"
|
||
|
||
labware[key] = {
|
||
"slot": item.get("positions") or item.get("slot"),
|
||
"labware": item.get("material_name") or item.get("labware"),
|
||
"well": item.get("well", []),
|
||
"type": item.get("type", "reagent"),
|
||
"role": item.get("role", ""),
|
||
"name": key,
|
||
}
|
||
|
||
return labware
|
||
|
||
|
||
def convert_from_json(
|
||
data: Union[str, PathLike, Dict[str, Any]],
|
||
workstation_name: str = "PRCXi",
|
||
validate: bool = True,
|
||
) -> WorkflowGraph:
|
||
"""
|
||
从 JSON 数据或文件转换为 WorkflowGraph
|
||
|
||
支持的 JSON 格式：
|
||
1. {"workflow": [...], "reagent": {...}} - 直接格式
|
||
2. {"steps_info": [...], "labware_info": [...]} - 需要规范化的格式
|
||
|
||
Args:
|
||
data: JSON 文件路径、字典数据、或 JSON 字符串
|
||
workstation_name: 工作站名称，默认 "PRCXi"
|
||
validate: 是否校验句柄配置，默认 True
|
||
|
||
Returns:
|
||
WorkflowGraph: 构建好的工作流图
|
||
|
||
Raises:
|
||
ValueError: 不支持的 JSON 格式 或 句柄校验失败
|
||
FileNotFoundError: 文件不存在
|
||
json.JSONDecodeError: JSON 解析失败
|
||
"""
|
||
# 处理输入数据
|
||
if isinstance(data, (str, PathLike)):
|
||
path = Path(data)
|
||
if path.exists():
|
||
with path.open("r", encoding="utf-8") as fp:
|
||
json_data = json.load(fp)
|
||
elif isinstance(data, str):
|
||
# 尝试作为 JSON 字符串解析
|
||
json_data = json.loads(data)
|
||
else:
|
||
raise FileNotFoundError(f"文件不存在: {data}")
|
||
elif isinstance(data, dict):
|
||
json_data = data
|
||
else:
|
||
raise TypeError(f"不支持的数据类型: {type(data)}")
|
||
|
||
# 根据格式解析数据
|
||
if "workflow" in json_data and "reagent" in json_data:
|
||
# 格式1: workflow/reagent（已经是规范格式）
|
||
protocol_steps = json_data["workflow"]
|
||
labware_info = json_data["reagent"]
|
||
elif "steps_info" in json_data and "labware_info" in json_data:
|
||
# 格式2: steps_info/labware_info（需要规范化）
|
||
protocol_steps = normalize_steps(json_data["steps_info"])
|
||
labware_info = normalize_labware(json_data["labware_info"])
|
||
elif "steps" in json_data and "labware" in json_data:
|
||
# 格式3: steps/labware（另一种常见格式）
|
||
protocol_steps = normalize_steps(json_data["steps"])
|
||
if isinstance(json_data["labware"], list):
|
||
labware_info = normalize_labware(json_data["labware"])
|
||
else:
|
||
labware_info = json_data["labware"]
|
||
else:
|
||
raise ValueError(
|
||
"不支持的 JSON 格式。支持的格式：\n"
|
||
"1. {'workflow': [...], 'reagent': {...}}\n"
|
||
"2. {'steps_info': [...], 'labware_info': [...]}\n"
|
||
"3. {'steps': [...], 'labware': [...]}"
|
||
)
|
||
|
||
# 构建工作流图
|
||
graph = build_protocol_graph(
|
||
labware_info=labware_info,
|
||
protocol_steps=protocol_steps,
|
||
workstation_name=workstation_name,
|
||
action_resource_mapping=ACTION_RESOURCE_MAPPING,
|
||
)
|
||
|
||
# 校验句柄配置
|
||
if validate:
|
||
is_valid, errors = validate_workflow_handles(graph)
|
||
if not is_valid:
|
||
import warnings
|
||
|
||
for error in errors:
|
||
warnings.warn(f"句柄校验警告: {error}")
|
||
|
||
return graph
|
||
|
||
|
||
def convert_json_to_node_link(
|
||
data: Union[str, PathLike, Dict[str, Any]],
|
||
workstation_name: str = "PRCXi",
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
将 JSON 数据转换为 node-link 格式的字典
|
||
|
||
Args:
|
||
data: JSON 文件路径、字典数据、或 JSON 字符串
|
||
workstation_name: 工作站名称，默认 "PRCXi"
|
||
|
||
Returns:
|
||
Dict: node-link 格式的工作流数据
|
||
"""
|
||
graph = convert_from_json(data, workstation_name)
|
||
return graph.to_node_link_dict()
|
||
|
||
|
||
def convert_json_to_workflow_list(
|
||
data: Union[str, PathLike, Dict[str, Any]],
|
||
workstation_name: str = "PRCXi",
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
将 JSON 数据转换为工作流列表格式
|
||
|
||
Args:
|
||
data: JSON 文件路径、字典数据、或 JSON 字符串
|
||
workstation_name: 工作站名称，默认 "PRCXi"
|
||
|
||
Returns:
|
||
List: 工作流节点列表
|
||
"""
|
||
graph = convert_from_json(data, workstation_name)
|
||
return graph.to_dict()
|
||
|
||
|
||
# 为了向后兼容，保留下划线前缀的别名
|
||
_normalize_steps = normalize_steps
|
||
_normalize_labware = normalize_labware
|