zcbot/tools/software_jobs.py

273 lines
10 KiB
Python

"""Agent 可调用的专业软件任务工具。"""
from __future__ import annotations
import json
from uuid import UUID, uuid4
from core.software_jobs import (
SoftwareJobError,
create_job,
get_job,
list_jobs,
request_job_cancel,
)
from core.software_nodes import SUPPORTED_CAPABILITIES, list_nodes
from .base import Tool
class _SoftwareJobTool(Tool):
def __init__(self, user_id: UUID, task_id: UUID, **kwargs) -> None:
super().__init__(**kwargs)
self.user_id = user_id
self.task_id = task_id
class SoftwareCapabilityListTool(_SoftwareJobTool):
name = "software_capability_list"
description = "List professional software capabilities available through managed Windows nodes."
parameters = {"type": "object", "properties": {}, "additionalProperties": False}
def execute(self) -> str:
nodes = list_nodes()
items = [{
"capability": item,
"display_name": "Origin 科研绘图" if item == "origin.plot@v2" else item,
"available_nodes": sum(
1
for node in nodes
if node["status"] == "online"
and item in (node.get("capabilities") or [])
and (node.get("runtime") or {}).get("available_slots", 0) > 0
),
} for item in sorted(SUPPORTED_CAPABILITIES)]
return json.dumps({"capabilities": items}, ensure_ascii=False)
class SoftwareJobSubmitTool(_SoftwareJobTool):
name = "software_job_submit"
description = (
"Submit an Origin 2D, error-bar, contour, 3D, ternary, or heatmap plot job using "
"one or more registered CSV, XLSX, or JSON artifacts. "
"Call register_artifact first for each workspace file. "
"Return immediately with job_id; do not poll continuously or wait for completion."
)
_axis_schema = {
"type": "object",
"properties": {
"title": {"type": "string"},
"unit": {"type": "string"},
"scale": {"type": "string", "enum": ["linear"]},
},
"additionalProperties": False,
}
_plot_schema = {
"type": "object",
"properties": {
"type": {"type": "string", "enum": [
"line", "scatter", "line_scatter", "column", "bar", "grouped_column",
"y_error", "contour", "surface_3d", "ternary", "heatmap",
]},
"series": {
"type": "array",
"minItems": 1,
"maxItems": 16,
"items": {
"type": "object",
"properties": {
"input": {"type": "string", "pattern": "^[a-z][a-z0-9_]{0,31}$"},
"x": {"type": "string", "minLength": 1, "maxLength": 128},
"y": {"type": "string", "minLength": 1, "maxLength": 128},
"z": {"type": "string", "minLength": 1, "maxLength": 128},
"y_error": {"type": "string", "minLength": 1, "maxLength": 128},
"label": {"type": "string", "minLength": 1, "maxLength": 200},
},
"required": ["input"],
"additionalProperties": False,
},
"description": (
"Typed data roles. XY plots require x/y; y_error requires x/y/y_error; "
"contour, surface_3d, ternary, and heatmap require x/y/z."
),
},
"template": {"type": "string", "enum": ["publication_double_column"]},
"title": {"type": "string", "maxLength": 500},
"x_axis": _axis_schema,
"y_axis": _axis_schema,
"z_axis": _axis_schema,
"legend": {
"type": "object",
"properties": {
"enabled": {"type": "boolean", "enum": [True]},
"position": {"type": "string", "enum": ["top_right"]},
},
"additionalProperties": False,
},
"error_bars": {"type": "null"},
},
"required": ["type", "series"],
"additionalProperties": False,
}
_input_schema = {
"type": "object",
"properties": {
"key": {"type": "string", "pattern": "^[a-z][a-z0-9_]{0,31}$"},
"artifact_id": {
"type": "string",
"description": "Artifact UUID returned by register_artifact.",
},
"selector": {
"type": "object",
"properties": {
"sheet": {"type": "string", "minLength": 1, "maxLength": 128},
},
"required": ["sheet"],
"additionalProperties": False,
},
},
"required": ["key", "artifact_id"],
"additionalProperties": False,
}
_output_schema = {
"type": "object",
"properties": {
"key": {
"type": "string",
"enum": ["project", "figure_png", "figure_svg", "figure_pdf"],
},
"type": {"type": "string", "enum": ["project", "figure"]},
"format": {"type": "string", "enum": ["opju", "png", "svg", "pdf"]},
"options": {
"type": "object",
"properties": {"dpi": {"type": "integer", "minimum": 72, "maximum": 1200}},
"required": ["dpi"],
"additionalProperties": False,
"description": "Only valid for the PNG figure output.",
},
},
"required": ["key", "type", "format"],
"additionalProperties": False,
}
parameters = {
"type": "object",
"properties": {
"capability": {"type": "string", "enum": sorted(SUPPORTED_CAPABILITIES)},
"inputs": {
"type": "array",
"minItems": 1,
"maxItems": 16,
"items": _input_schema,
},
"operation": {
"type": "object",
"properties": {"plot": _plot_schema},
"required": ["plot"],
"additionalProperties": False,
},
"outputs": {
"type": "array",
"minItems": 1,
"maxItems": 16,
"uniqueItems": True,
"items": _output_schema,
"description": (
"Requested deliverables. Use project/project/opju, "
"figure_png/figure/png, figure_svg/figure/svg, or figure_pdf/figure/pdf."
),
},
"idempotency_key": {
"type": "string",
"description": "Stable unique key for this exact submission; omit to generate one.",
},
},
"required": ["capability", "inputs", "operation", "outputs"],
"additionalProperties": False,
}
def execute(
self,
capability: str,
inputs: list[dict] | None = None,
operation: dict | None = None,
outputs: list[dict] | None = None,
idempotency_key: str = "",
) -> str:
try:
if not inputs or operation is None or not outputs:
return "[Error] inputs, operation, and outputs are required for a software job"
canonical_inputs = []
for item in inputs:
if not isinstance(item, dict):
return "[Error] each input must be an object"
try:
artifact_id = str(UUID(str(item.get("artifact_id") or "")))
except ValueError:
return "[Error] every artifact_id must be a UUID; call register_artifact first"
canonical = {**item, "artifact_id": artifact_id}
canonical_inputs.append(canonical)
normalized_request = {
"schema_version": 2,
"inputs": canonical_inputs,
"operation": operation,
"outputs": outputs,
}
job, created = create_job(
self.user_id,
self.task_id,
idempotency_key=idempotency_key.strip() or str(uuid4()),
capability=capability,
request=normalized_request,
)
return json.dumps({**job, "created": created}, ensure_ascii=False)
except (SoftwareJobError, ValueError) as exc:
return f"[Error] {exc}"
class SoftwareJobStatusTool(_SoftwareJobTool):
name = "software_job_status"
description = (
"Check one software job, or list recent jobs in the current task when job_id is omitted. "
"A succeeded job returns output_dir as the authoritative starting directory; "
"inspect or search within that directory to analyze its files."
)
parameters = {
"type": "object",
"properties": {"job_id": {"type": "string"}},
"additionalProperties": False,
}
def execute(self, job_id: str = "") -> str:
try:
if job_id.strip():
item = get_job(self.user_id, UUID(job_id.strip()))
if item is None or item["task_id"] != str(self.task_id):
return "[Error] software job not found"
return json.dumps(item, ensure_ascii=False)
return json.dumps(
{"results": list_jobs(self.user_id, task_id=self.task_id, limit=20)},
ensure_ascii=False,
)
except ValueError:
return "[Error] invalid job_id"
class SoftwareJobCancelTool(_SoftwareJobTool):
name = "software_job_cancel"
description = "Request cancellation of a software job in the current task after the user asks to stop it."
parameters = {
"type": "object",
"properties": {"job_id": {"type": "string"}},
"required": ["job_id"],
"additionalProperties": False,
}
def execute(self, job_id: str) -> str:
try:
item = get_job(self.user_id, UUID(job_id.strip()))
if item is None or item["task_id"] != str(self.task_id):
return "[Error] software job not found"
job, _ = request_job_cancel(self.user_id, UUID(job_id.strip()))
return json.dumps(job, ensure_ascii=False)
except (ValueError, SoftwareJobError) as exc:
return f"[Error] {exc}"