task_flow.py

"""TaskFlowManager — CRUD, Hermes-generated flow creation, and execution.

Task Flows are structured agent pipelines that execute steps in strict order,
with variable passing between steps. Flows can be generated by Hermes from
natural language descriptions, then visually inspected before execution.
"""

import json
import re
import uuid
from datetime import datetime
from pathlib import Path
from typing import Optional

from event_buffer import EventBuffer
from hermes_client import call_hermes
from storage import load_json_dir, save_json, delete_json
from workflow_engine import execute_flow


class TaskFlowManager:
    """Manages task flow definitions and executions."""

    def __init__(
        self,
        flows_dir: Path,
        runs_dir: Path,
        agents: dict,
        hermes_url: str,
        hermes_key: str,
    ):
        self.flows_dir = flows_dir
        self.runs_dir = runs_dir
        self.agents = agents
        self.hermes_url = hermes_url
        self.hermes_key = hermes_key
        self.flows: dict[str, dict] = {}
        self.runs: dict[str, dict] = {}
        self._active_buffers: dict[str, EventBuffer] = {}

    def load(self):
        """Load flows and runs from disk."""
        self.flows = load_json_dir(self.flows_dir)
        self.runs = load_json_dir(self.runs_dir)

    def list_flows(self) -> list[dict]:
        return sorted(self.flows.values(), key=lambda f: f.get("created_at", ""), reverse=True)

    def get_flow(self, flow_id: str) -> Optional[dict]:
        return self.flows.get(flow_id)

    def create_flow(self, name: str, description: str, steps: list[dict]) -> dict:
        fid = uuid.uuid4().hex[:12]
        flow = {
            "id": fid,
            "name": name,
            "description": description,
            "created_at": datetime.now().isoformat(),
            "steps": steps,
        }
        self.flows[fid] = flow
        save_json(self.flows_dir, flow)
        return flow

    def update_flow(self, flow_id: str, **kwargs) -> Optional[dict]:
        flow = self.flows.get(flow_id)
        if not flow:
            return None
        if "name" in kwargs and kwargs["name"] is not None:
            flow["name"] = kwargs["name"]
        if "description" in kwargs and kwargs["description"] is not None:
            flow["description"] = kwargs["description"]
        if "steps" in kwargs and kwargs["steps"] is not None:
            flow["steps"] = kwargs["steps"]
        flow["updated_at"] = datetime.now().isoformat()
        save_json(self.flows_dir, flow)
        return flow

    def delete_flow(self, flow_id: str) -> bool:
        if flow_id not in self.flows:
            return False
        del self.flows[flow_id]
        delete_json(self.flows_dir, flow_id)
        return True

    async def generate_flow(self, user_description: str) -> dict:
        """Use Hermes to generate a task flow definition from natural language."""
        agent_lines = []
        for a in self.agents.values():
            preview = a.get("system_prompt", "")[:80]
            agent_lines.append(f"- {a['id']}: {a['name']} — {preview}...")
        agent_list = "\n".join(agent_lines)

        system_prompt = f"""你是一个任务流程生成器。用户会描述一个多步骤的任务，你需要生成一个结构化的 JSON 流程定义。

可用的 Agent 列表：
{agent_list}

请生成一个 JSON 对象，格式如下：
{{{{
  "name": "流程名称",
  "description": "简短描述",
  "steps": [
    {{{{
      "id": "step_1",
      "agent_id": "agent的id",
      "name": "步骤名称",
      "prompt_template": "给agent的指令模板，用 {{{{{{变量名}}}}}} 引用前序步骤的输出",
      "input_vars": ["引用的变量名列表"],
      "output_var": "本步骤输出的变量名",
      "timeout": 300
    }}}}
  ]
}}}}

规则：
1. 只输出 JSON，不要其他内容
2. 第一个步骤可以用 {{{{{{input}}}}}} 引用用户的原始输入
3. 后续步骤用前序步骤的 output_var 名引用其输出
4. 每个步骤的 agent_id 必须是上面列表中的一个
5. timeout 默认 300 秒，复杂任务可以设置更长
6. prompt_template 要写清楚具体的指令，不要笼统"""

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_description},
        ]

        response = await call_hermes(
            self.hermes_url, self.hermes_key, messages,
            temperature=0.3, max_tokens=2000,
        )

        try:
            result = json.loads(response)
        except json.JSONDecodeError:
            m = re.search(r"```(?:json)?\s*\n?(\{.*?\})\s*```", response, re.DOTALL)
            if m:
                result = json.loads(m.group(1))
            else:
                m = re.search(r"\{[\s\S]*\"steps\"[\s\S]*\}", response)
                if m:
                    result = json.loads(m.group(0))
                else:
                    raise ValueError(f"Hermes did not return valid JSON: {response[:200]}")

        if "steps" not in result:
            raise ValueError("Generated flow missing 'steps' field")
        for i, step in enumerate(result["steps"]):
            step.setdefault("id", f"step_{i+1}")
            step.setdefault("name", f"Step {i+1}")
            step.setdefault("input_vars", [])
            step.setdefault("output_var", "")
            step.setdefault("timeout", 300)
            if "agent_id" not in step:
                raise ValueError(f"Step {i+1} missing 'agent_id'")
            if "prompt_template" not in step:
                raise ValueError(f"Step {i+1} missing 'prompt_template'")

        return result

    def create_run(self, flow_id: str, input_text: str) -> dict:
        run_id = uuid.uuid4().hex[:8]
        run = {
            "id": run_id,
            "flow_id": flow_id,
            "status": "running",
            "started_at": datetime.now().isoformat(),
            "input": input_text,
            "step_results": {},
            "variables": {"input": input_text},
        }
        self.runs[run_id] = run
        save_json(self.runs_dir, run)
        return run

    def update_run(self, run_id: str, **kwargs):
        run = self.runs.get(run_id)
        if not run:
            return
        run.update(kwargs)
        save_json(self.runs_dir, run)

    async def execute(self, flow_id: str, input_text: str, event_buffer: EventBuffer, run_id: str = None) -> dict:
        """Execute a task flow. Returns the final variables dict."""
        flow = self.flows.get(flow_id)
        if not flow:
            raise ValueError(f"Flow '{flow_id}' not found")

        if run_id:
            run = self.runs.get(run_id)
            if not run:
                raise ValueError(f"Run '{run_id}' not found")
        else:
            run = self.create_run(flow_id, input_text)
            run_id = run["id"]
        self._active_buffers[run_id] = event_buffer

        try:
            variables = await execute_flow(
                flow=flow,
                input_text=input_text,
                agents=self.agents,
                hermes_url=self.hermes_url,
                hermes_key=self.hermes_key,
                event_buffer=event_buffer,
            )
            self.update_run(run_id, status="completed", completed_at=datetime.now().isoformat())
            return variables
        except Exception as e:
            self.update_run(run_id, status="failed", error=str(e)[:300], completed_at=datetime.now().isoformat())
            raise
        finally:
            self._active_buffers.pop(run_id, None)

    def get_run_buffer(self, run_id: str) -> Optional[EventBuffer]:
        return self._active_buffers.get(run_id)

    def list_runs(self, flow_id: Optional[str] = None) -> list[dict]:
        runs = list(self.runs.values())
        if flow_id:
            runs = [r for r in runs if r["flow_id"] == flow_id]
        return sorted(runs, key=lambda r: r.get("started_at", ""), reverse=True)