227 lines
8.3 KiB
Python
227 lines
8.3 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@v1" 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 plot job using a registered CSV, XLSX, or JSON artifact. "
|
|
"Call register_artifact first when the input is only a 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"]},
|
|
"x": {
|
|
"type": "string",
|
|
"minLength": 1,
|
|
"maxLength": 128,
|
|
"description": "Exact X column name in the input file.",
|
|
},
|
|
"y": {
|
|
"type": "array",
|
|
"minItems": 1,
|
|
"maxItems": 16,
|
|
"uniqueItems": True,
|
|
"items": {"type": "string", "minLength": 1, "maxLength": 128},
|
|
"description": "One to sixteen unique Y column names.",
|
|
},
|
|
"template": {"type": "string", "enum": ["publication_double_column"]},
|
|
"title": {"type": "string", "maxLength": 500},
|
|
"x_axis": _axis_schema,
|
|
"y_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", "x", "y"],
|
|
"additionalProperties": False,
|
|
}
|
|
_output_schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"formats": {
|
|
"type": "array",
|
|
"minItems": 1,
|
|
"uniqueItems": True,
|
|
"items": {"type": "string", "enum": ["opju", "png", "svg", "pdf"]},
|
|
},
|
|
"dpi": {"type": "integer", "minimum": 72, "maximum": 1200},
|
|
"capture_screenshots": {"type": "boolean"},
|
|
"record_video": {"type": "boolean", "enum": [False]},
|
|
},
|
|
"required": ["formats"],
|
|
"additionalProperties": False,
|
|
}
|
|
parameters = {
|
|
"type": "object",
|
|
"properties": {
|
|
"capability": {"type": "string", "enum": sorted(SUPPORTED_CAPABILITIES)},
|
|
"input_id": {
|
|
"type": "string",
|
|
"description": "Artifact UUID returned by register_artifact or another artifact-producing flow.",
|
|
},
|
|
"sheet": {
|
|
"type": "string",
|
|
"minLength": 1,
|
|
"maxLength": 128,
|
|
"description": "Optional XLSX worksheet name.",
|
|
},
|
|
"plot": _plot_schema,
|
|
"output": _output_schema,
|
|
"idempotency_key": {
|
|
"type": "string",
|
|
"description": "Stable unique key for this exact submission; omit to generate one.",
|
|
},
|
|
},
|
|
"required": ["capability", "input_id", "plot", "output"],
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
def execute(
|
|
self,
|
|
capability: str,
|
|
input_id: str = "",
|
|
plot: dict | None = None,
|
|
output: dict | None = None,
|
|
sheet: str = "",
|
|
request: dict | None = None,
|
|
idempotency_key: str = "",
|
|
) -> str:
|
|
try:
|
|
if isinstance(input_id, dict) and request is None:
|
|
request = input_id
|
|
input_id = ""
|
|
if request is not None:
|
|
if input_id or plot is not None or output is not None or sheet:
|
|
return "[Error] request cannot be combined with input_id, sheet, plot, or output"
|
|
normalized_request = request
|
|
else:
|
|
if not input_id or plot is None or output is None:
|
|
return "[Error] input_id, plot, and output are required for an Origin job"
|
|
try:
|
|
canonical_input_id = str(UUID(str(input_id)))
|
|
except ValueError:
|
|
return "[Error] input_id must be an artifact UUID; call register_artifact first"
|
|
input_spec = {"input_id": canonical_input_id}
|
|
if sheet:
|
|
input_spec["sheet"] = sheet
|
|
normalized_request = {
|
|
"schema_version": 1,
|
|
"input": input_spec,
|
|
"plot": plot,
|
|
"output": output,
|
|
}
|
|
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."
|
|
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}"
|