Spaces:
Sleeping
Sleeping
| 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 | |
| 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() | |
| 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 | |
| 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) | |
| 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.") | |