backend/engine.py

import json
import os
import re
import uuid
from dataclasses import dataclass, field

from openai import OpenAI

from .models import RunResponse, StartRequest, TraceEvent, Workflow, WorkflowNode
from .prompts import decision_prompt, tool_prompt, worker_prompt


MAX_STEPS = 50
MAX_TOOL_CALLS_PER_WORKER = 5


@dataclass
class RunState:
    workflow: Workflow
    current_id: str
    value: str
    trace: list[TraceEvent] = field(default_factory=list)
    steps: int = 0


class WorkflowEngine:
    def __init__(self) -> None:
        self.runs: dict[str, RunState] = {}

    def _client(self) -> OpenAI | None:
        key = os.getenv("OPENAI_API_KEY", "").strip()
        if not key:
            return None
        kwargs = {"api_key": key}
        if os.getenv("OPENAI_BASE_URL"):
            kwargs["base_url"] = os.environ["OPENAI_BASE_URL"]
        return OpenAI(**kwargs)

    def _complete(self, prompt: str) -> str:
        client = self._client()
        if not client:
            return "Demo output"
        response = client.responses.create(model=os.getenv("OPENAI_MODEL", "gpt-4.1-mini"), input=prompt)
        return response.output_text.strip()

    @staticmethod
    def _graphs(workflow: Workflow) -> tuple[dict[str, WorkflowNode], dict[str, list], dict[str, list]]:
        nodes = {node.id: node for node in workflow.nodes}
        if len(nodes) != len(workflow.nodes):
            raise ValueError("Node IDs must be unique.")
        starts = [node for node in workflow.nodes if node.data.mode == "Input"]
        if len(starts) != 1:
            raise ValueError("The workflow needs exactly one Input node.")
        flow = {node_id: [] for node_id in nodes}
        tools = {node_id: [] for node_id in nodes}
        for edge in workflow.edges:
            if edge.source not in nodes or edge.target not in nodes:
                raise ValueError("Every connection must point to an existing node.")
            if edge.data.connectionType == "tool":
                if nodes[edge.source].data.mode != "AI Worker" or nodes[edge.target].data.mode != "Simulated Tool":
                    raise ValueError("Tool connections must run from an AI Worker to a Simulated Tool.")
                tools[edge.source].append(edge)
            else:
                if nodes[edge.target].data.mode == "Simulated Tool" or nodes[edge.source].data.mode == "Simulated Tool":
                    raise ValueError("Simulated Tools use dashed tool connections, not workflow connections.")
                flow[edge.source].append(edge)
        return nodes, flow, tools

    @classmethod
    def _validate(cls, workflow: Workflow) -> tuple[dict, dict[str, list]]:
        nodes, flow, _ = cls._graphs(workflow)
        return nodes, flow

    def start(self, request: StartRequest) -> RunResponse:
        nodes, _, _ = self._graphs(request.workflow)
        start = next(node for node in nodes.values() if node.data.mode == "Input")
        value = request.input.strip() or start.data.value.strip()
        if not value:
            raise ValueError("Enter one item in the Input node.")
        run_id = str(uuid.uuid4())
        self.runs[run_id] = RunState(request.workflow, start.id, value)
        return self._advance(run_id)

    def resume(self, run_id: str, response: str) -> RunResponse:
        state = self.runs.get(run_id)
        if not state:
            raise ValueError("This run no longer exists. Start a new run.")
        nodes, flow, _ = self._graphs(state.workflow)
        node = nodes[state.current_id]
        if node.data.mode != "Human Review":
            raise ValueError("This run is not waiting for human review.")
        reply = response.strip()
        if not reply:
            raise ValueError("Write a review response before continuing.")
        state.trace.append(TraceEvent(node_id=node.id, label=node.data.label, mode=node.data.mode, input=state.value, output=reply))
        state.value = reply
        if not flow[node.id]:
            raise ValueError("Human Review needs an outgoing workflow connection.")
        state.current_id = flow[node.id][0].target
        return self._advance(run_id)

    @staticmethod
    def _parse_json(text: str) -> dict:
        cleaned = re.sub(r"^```(?:json)?\s*|\s*```$", "", text.strip(), flags=re.I)
        try:
            value = json.loads(cleaned)
        except json.JSONDecodeError as exc:
            raise ValueError(f"AI Worker returned invalid tool instructions: {text}") from exc
        if not isinstance(value, dict):
            raise ValueError("AI Worker response must be a JSON object.")
        return value

    def _run_tool(self, tool: WorkflowNode, tool_input: str) -> str:
        if not self._client():
            return f"Simulated {tool.data.label} result for: {tool_input}"
        return self._complete(tool_prompt(tool, tool_input))

    def _run_worker(self, state: RunState, worker: WorkflowNode, available_tools: list[WorkflowNode]) -> str:
        tool_results: list[dict] = []
        if not self._client():
            if available_tools:
                tool = available_tools[0]
                result = self._run_tool(tool, state.value)
                state.trace.append(TraceEvent(node_id=tool.id, label=tool.data.label, mode=tool.data.mode, input=state.value, output=result, kind="tool"))
                return f"Demo AI Worker result using {tool.data.label}: {result}"
            return f"Demo AI Worker result: {state.value}"

        tools_by_id = {tool.id: tool for tool in available_tools}
        for _ in range(MAX_TOOL_CALLS_PER_WORKER + 1):
            command = self._parse_json(self._complete(worker_prompt(worker.data.prompt, state.value, available_tools, tool_results)))
            if command.get("action") == "final":
                output = str(command.get("output", "")).strip()
                if not output:
                    raise ValueError("AI Worker returned an empty final output.")
                return output
            if command.get("action") != "call_tool":
                raise ValueError("AI Worker must choose 'call_tool' or 'final'.")
            tool_id = str(command.get("tool_id", ""))
            if tool_id not in tools_by_id:
                raise ValueError(f"AI Worker tried to call an unavailable tool: {tool_id}")
            tool_input = str(command.get("tool_input", "")).strip()
            result = self._run_tool(tools_by_id[tool_id], tool_input)
            tool_results.append({"tool_id": tool_id, "input": tool_input, "output": result})
            tool = tools_by_id[tool_id]
            state.trace.append(TraceEvent(node_id=tool.id, label=tool.data.label, mode=tool.data.mode, input=tool_input, output=result, kind="tool"))
        raise ValueError(f"AI Worker exceeded {MAX_TOOL_CALLS_PER_WORKER} tool calls.")

    def _advance(self, run_id: str) -> RunResponse:
        state = self.runs[run_id]
        nodes, flow, tool_edges = self._graphs(state.workflow)
        while state.steps < MAX_STEPS:
            state.steps += 1
            node = nodes[state.current_id]
            mode, incoming = node.data.mode, state.value

            if mode == "Human Review":
                return RunResponse(run_id=run_id, status="awaiting_review", current_node_id=node.id, message=node.data.prompt or "Review the current item. What should happen next?", result=state.value, trace=state.trace)
            if mode == "Output":
                state.trace.append(TraceEvent(node_id=node.id, label=node.data.label, mode=mode, input=incoming, output=incoming))
                del self.runs[run_id]
                return RunResponse(run_id=run_id, status="completed", current_node_id=node.id, result=incoming, trace=state.trace)
            if mode == "Simulated Tool":
                raise ValueError("A Simulated Tool cannot be a workflow step. Connect it to an AI Worker's tool handle.")

            edges = flow[node.id]
            if not edges:
                raise ValueError(f"'{node.data.label}' has no outgoing workflow connection.")
            if mode == "Input":
                output, next_id = incoming, edges[0].target
            elif mode == "Decision Rules":
                choices = [(edge.target, str(edge.label or nodes[edge.target].data.label)) for edge in edges]
                if self._client():
                    answer = self._complete(decision_prompt(node.data.prompt, incoming, choices))
                    next_id = next((target for target, _ in choices if re.search(rf"(?<![\w-]){re.escape(target)}(?![\w-])", answer)), "")
                    if not next_id:
                        raise ValueError(f"Decision Rules returned an unknown route: {answer}")
                else:
                    next_id = choices[0][0]
                output = f"Selected route: {next(label for target, label in choices if target == next_id)}\n\n{incoming}"
            elif mode == "AI Worker":
                available_tools = [nodes[edge.target] for edge in tool_edges[node.id]]
                output, next_id = self._run_worker(state, node, available_tools), edges[0].target
            else:
                raise ValueError(f"Unsupported node type: {mode}")

            state.trace.append(TraceEvent(node_id=node.id, label=node.data.label, mode=mode, input=incoming, output=output))
            state.value, state.current_id = output, next_id
        raise ValueError(f"Run stopped after {MAX_STEPS} steps. Check for a loop.")