126 lines
4.6 KiB
Python
126 lines
4.6 KiB
Python
"""Drive one job by asking the state center where it is, then acting."""
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from .runners import mock as mock_runners
|
|
from .runtime.shell import execute_node
|
|
|
|
HUMAN_NODES = {"SOURCE_ADD", "M5_MANUAL_CONFIRM", "M5_RECEIPT_NOTICE"}
|
|
|
|
|
|
class JobAgent:
|
|
def __init__(
|
|
self,
|
|
client: Any,
|
|
job_id: str,
|
|
*,
|
|
runner_mode: str = "mock",
|
|
auto_human: bool = False,
|
|
auto_retry: bool = False,
|
|
operator: str = "mock-agent",
|
|
):
|
|
self.client = client
|
|
self.job_id = job_id
|
|
self.runner_mode = runner_mode
|
|
self.auto_human = auto_human
|
|
self.auto_retry = auto_retry
|
|
self.operator = operator
|
|
|
|
def step(self) -> dict[str, Any]:
|
|
detail = self.client.get_job(self.job_id)
|
|
job = detail["job"]
|
|
tasks = detail.get("tasks") or []
|
|
status = job.get("status")
|
|
node = job.get("current_node_code")
|
|
|
|
if status == "SUCCESS":
|
|
return {"stopped": "SUCCESS", "next_node": None}
|
|
if status == "FAILED_RETRYABLE":
|
|
if not self.auto_retry:
|
|
return {"stopped": "FAILED_RETRYABLE"}
|
|
return self.client.retry_job(self.job_id)
|
|
if status == "WAITING_USER" or node in HUMAN_NODES:
|
|
if not self.auto_human:
|
|
return {"stopped": "WAITING_USER"}
|
|
return self._submit_human(job, tasks)
|
|
if node == "PRODUCT_MATCH":
|
|
return self.client.match_run(self.job_id)
|
|
if node == "CONTRACT_APPROVE_WAIT":
|
|
return self.client.check_contract_result(self.job_id, True, operator=self.operator)
|
|
if node == "ARCHIVE_JOB":
|
|
return {"stopped": "SUCCESS", "next_node": None}
|
|
return execute_node(
|
|
self.client,
|
|
self.job_id,
|
|
node,
|
|
lambda: self._run(node, job),
|
|
attempt_no=int(job.get("attempt_no") or 1),
|
|
operator=self.operator,
|
|
)
|
|
|
|
def _run(self, node_code: str, job: dict[str, Any]) -> dict[str, Any]:
|
|
if self.runner_mode != "mock":
|
|
raise RuntimeError(f"unsupported runner_mode {self.runner_mode}")
|
|
return mock_runners.run(node_code, job, self.client)
|
|
|
|
def _submit_human(self, job: dict[str, Any], tasks: list[dict[str, Any]]) -> dict[str, Any]:
|
|
node = job["current_node_code"]
|
|
task = _open_task(tasks, node)
|
|
if task is None and hasattr(self.client, "list_tasks"):
|
|
listed = self.client.list_tasks(status="OPEN")
|
|
task = _open_task(listed.get("items") or [], node, self.job_id)
|
|
if task is None:
|
|
return {"stopped": "WAITING_USER", "reason": "no_open_task"}
|
|
|
|
action = "CONFIRM_AND_CONTINUE"
|
|
payload: dict[str, Any] = {}
|
|
if node == "PRODUCT_MATCH":
|
|
action = "SELECT_CANDIDATE"
|
|
payload = {
|
|
"productMatchId": task.get("product_match_id"),
|
|
"goodsId": "G-001",
|
|
"goodsCode": "YP-001",
|
|
"baseUnit": "盒",
|
|
"saveAsAlias": True,
|
|
}
|
|
elif node == "M5_RECEIPT_NOTICE":
|
|
payload = {"receiptNoticeNo": f"SH{uuid.uuid4().hex[:8].upper()}"}
|
|
|
|
return self.client.submit_task(
|
|
task["task_id"],
|
|
action,
|
|
payload,
|
|
submit_token=uuid.uuid4().hex,
|
|
operator=self.operator,
|
|
)
|
|
|
|
def run(self, *, max_steps: int = 40) -> list[dict[str, Any]]:
|
|
history: list[dict[str, Any]] = []
|
|
for _ in range(max_steps):
|
|
result = self.step()
|
|
history.append(result)
|
|
if result.get("stopped") in {"SUCCESS", "WAITING_USER", "FAILED_RETRYABLE", "FAILED_MANUAL"}:
|
|
return history
|
|
if result.get("waiting_user"):
|
|
continue
|
|
if result.get("next_node") is None and result.get("ok"):
|
|
detail = self.client.get_job(self.job_id)
|
|
if detail["job"].get("status") == "SUCCESS":
|
|
history.append({"stopped": "SUCCESS", "next_node": None})
|
|
return history
|
|
raise RuntimeError(f"job {self.job_id} exceeded {max_steps} agent steps")
|
|
|
|
|
|
def _open_task(tasks: list[dict[str, Any]], node_code: str, job_id: str | None = None):
|
|
for task in tasks:
|
|
if task.get("node_code") != node_code:
|
|
continue
|
|
if task.get("status") not in (None, "OPEN"):
|
|
continue
|
|
if job_id and task.get("job_id") not in (None, job_id):
|
|
continue
|
|
return task
|
|
return None
|