""" WorkflowRunner — the in-house orchestrator that chains small, single-decision prompts into an end-to-end workflow (stage-1 autoplan, stage-2 autogen, etc.). Why not LangChain / LangGraph / CrewAI? * LangChain adds a ~40 MB dependency tree for what is, at its core, a `for step in steps: call_llm(step); validate; retry` loop. Its retry + output-parsing abstractions would still need a shim to talk to our Enterprise Settings resolver and our YAML prompt library. * LangGraph's strength is cyclical graphs with conditional branching; our workflows are straight DAGs (mostly linear, occasionally a fan-out over scenes). The framework tax isn't justified for 7 nodes. * CrewAI is multi-agent-role oriented; we have one role ("the planner") running seven specialised prompts, not seven agents collaborating. Using CrewAI would force us to pretend each prompt is an agent, which just obscures the data flow. What this module does --------------------- Given a list of ``Step``s and a starting context dict: 1. For each step, build the prompt variables by calling ``step.build_vars(context)``. 2. Render the prompt via ``PromptLibrary``. 3. Dispatch via ``llm_router.call_prompt`` (Enterprise Settings chat model, policy-driven timeout). 4. Parse the assistant content with ``step.parse``. 5. Validate the parsed value with ``step.validate``. 6. On parse/validate failure, retry up to ``policy.retries``. 7. On final failure, consult ``policy.fallback``: 'abort' → raise StepFailure (workflow halts) other str → let ``step.on_fallback`` interpret (returns a safe default, or raises if step.build_vars can't cope). 8. Store the value under ``step.output_key`` in the context so later steps can read it. 9. Emit events via ``on_event`` hook so the frontend spinner can display real progress (REV-5 wires this to SSE). Events ------ The runner emits four event kinds: workflow_started {workflow, step_ids, started_at_ms} step_started {step_id, prompt_id, attempt=1} step_completed {step_id, prompt_id, attempt, duration_ms, preview} # first 80 chars of output step_failed {step_id, prompt_id, attempt, reason, fatal} # fatal=True only if abort workflow_completed {workflow, duration_ms, ok} ``preview`` is a short string derived from ``str(value)[:80]`` so operators can eyeball "did the LLM actually say something sensible?" without exposing full content. Design choices -------------- * **Pure Python, no framework tax.** ~180 LoC, no extra deps. * **Context is a plain dict.** Steps read/write it freely. No immutable-state ceremony; this module is intentionally small. * **No global state.** The runner holds no caches. Same library, same context, same result — easy to test. * **Parallel fan-out is deferred.** Today each Step runs sequentially (simplest, deterministic). REV-4 introduces a ``ParallelFanout`` wrapper when per-scene script generation needs it. Premature abstraction risk is not worth it now. """ from __future__ import annotations import asyncio import inspect import logging import time from dataclasses import dataclass, field from typing import ( Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, ) from ..llm_router import call_prompt from ..prompts import ( PromptLibrary, PromptLibraryError, PromptPolicy, RenderedPrompt, default_library, ) log = logging.getLogger(__name__) # ── Errors ───────────────────────────────────────────────────── class StepFailure(Exception): """Raised when a step exhausts retries and its fallback is ``abort``. Carries enough context for the route layer to produce a useful error payload. """ def __init__( self, *, step_id: str, prompt_id: str, reason: str, attempts: int, ) -> None: self.step_id = step_id self.prompt_id = prompt_id self.reason = reason self.attempts = attempts super().__init__( f"Step {step_id!r} ({prompt_id}) failed after {attempts} attempt(s): {reason}" ) # ── Data shapes ──────────────────────────────────────────────── @dataclass(frozen=True) class WorkflowEvent: """Structured telemetry emitted per step transition. The frontend SSE stream serialises these to JSON; tests assert on the ``kind`` + payload keys. Payloads are plain dicts so we don't leak dataclass types across the wire. """ kind: str # 'workflow_started' | 'step_started' | ... ts_ms: int # epoch-relative-ish, for ordering payload: Dict[str, Any] = field(default_factory=dict) @dataclass(frozen=True) class StepResult: """Outcome of one completed step. Stored on the WorkflowResult so tests can introspect the full trace. """ step_id: str prompt_id: str value: Any attempts: int used_fallback: bool duration_ms: int @dataclass class WorkflowResult: """Final context + per-step trace for a completed run. ``aborted`` is True if a StepFailure halted the workflow — in that case ``steps`` contains every step that *did* run up to and including the failing one. """ context: Dict[str, Any] steps: List[StepResult] = field(default_factory=list) events: List[WorkflowEvent] = field(default_factory=list) aborted: bool = False error: Optional[str] = None started_at_ms: int = 0 duration_ms: int = 0 # ── Step definition ──────────────────────────────────────────── # Type aliases keep the Step signature readable. VarBuilder = Callable[[Mapping[str, Any]], Mapping[str, Any]] Parser = Callable[[str], Any] Validator = Callable[[Any], Optional[str]] # returns err msg or None FallbackFn = Callable[[Mapping[str, Any], Optional[str]], Any] # (context, policy.fallback) -> value @dataclass(frozen=True) class Step: """One prompt in a workflow. Fields ------ step_id Stable ID for logs/events (often equals the prompt id's trailing segment, e.g. ``classify_mode``). prompt_id Full library id (``autoplan.classify_mode``). output_key Where in ``context`` the validated value lands. build_vars ``context -> variables dict`` for the prompt renderer. Keep it pure — no I/O. parse ``raw content -> parsed value``. Raise ValueError on parse failure; the runner treats it like a validation miss and retries. validate ``value -> None | error message``. Returning a string triggers retry (and eventually fallback). fallback ``(context, fallback_token) -> default value``. Called only when retries are exhausted AND the policy fallback is NOT ``abort``. Raise StepFailure if no safe default is producible. temperature, max_tokens Per-step generation knobs. Short enums want temp=0.0; creative prose wants 0.5ish. """ step_id: str prompt_id: str output_key: str build_vars: VarBuilder parse: Parser = field(default=lambda s: s.strip()) validate: Validator = field(default=lambda v: None) fallback: Optional[FallbackFn] = None temperature: float = 0.3 max_tokens: int = 350 # Per-step model override (Ollama model id). Empty string means # "use the server default model". Threaded into chat_ollama by # the runner so individual workflow steps can route around the # default model — used by Mature (gated) experiences to pick an # uncensored / abliterated model the operator selected in the # wizard's Step 0 picker. model: str = "" # ── Event hook type ──────────────────────────────────────────── EventHook = Callable[[WorkflowEvent], None] # ── Helper: content extractor ────────────────────────────────── def extract_content(response: Any) -> str: """Pull the assistant content out of an OpenAI-style envelope. Returns empty string on any shape we don't recognise — the caller's validator then produces the "empty output" retry. """ if not isinstance(response, dict): return "" choices = response.get("choices") if not isinstance(choices, list) or not choices: return "" first = choices[0] if not isinstance(first, dict): return "" msg = first.get("message") or {} if isinstance(msg, dict): content = msg.get("content") if isinstance(content, str): return content.strip() return "" # ── Runner ───────────────────────────────────────────────────── class WorkflowRunner: """Sequentially run a list of ``Step``s against a shared context. Stateless across runs — instantiate once per process is fine. """ def __init__( self, *, library: Optional[PromptLibrary] = None, ) -> None: self._library = library or default_library() async def run( self, *, workflow: str, steps: Sequence[Step], context: Optional[Dict[str, Any]] = None, on_event: Optional[EventHook] = None, ) -> WorkflowResult: """Execute ``steps`` in order. Returns a ``WorkflowResult`` with the final context, per-step trace, and event log. A StepFailure is caught: the workflow is marked aborted and returned (instead of raised). Callers that want the exception style should check ``result.aborted`` and re-raise. """ ctx: Dict[str, Any] = dict(context or {}) events: List[WorkflowEvent] = [] def _emit(ev: WorkflowEvent) -> None: events.append(ev) if on_event is not None: try: # Backward-compat hook bridge: # * New runner-native hooks expect one arg: # on_event(WorkflowEvent) # * Legacy planner/generator hooks expect two args: # on_event(kind, payload) # # We support both so route-level SSE hooks can be # passed through unchanged without dropping workflow # events due to signature mismatch. payload = dict(ev.payload or {}) argc: Optional[int] = None has_varargs = False try: sig = inspect.signature(on_event) params = list(sig.parameters.values()) has_varargs = any( p.kind == inspect.Parameter.VAR_POSITIONAL for p in params ) argc = len([ p for p in params if p.kind in ( inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, ) ]) except Exception: argc = None if has_varargs or (argc is not None and argc >= 2): on_event(ev.kind, payload) elif argc == 0: on_event() else: # Default path: WorkflowEvent callback. on_event(ev) except Exception: # noqa: BLE001 # Hooks must never break the workflow. log.exception("workflow event hook raised — ignored") started = _now_ms() result = WorkflowResult( context=ctx, steps=[], events=events, started_at_ms=started, ) _emit(WorkflowEvent( kind="workflow_started", ts_ms=_now_ms(), payload={ "workflow": workflow, "step_ids": [s.step_id for s in steps], "started_at_ms": started, }, )) try: for step in steps: step_result = await self._run_step(step, ctx, _emit) result.steps.append(step_result) ctx[step.output_key] = step_result.value except StepFailure as exc: result.aborted = True result.error = str(exc) log.warning("workflow_aborted workflow=%s reason=%s", workflow, exc) finally: result.duration_ms = max(0, _now_ms() - started) _emit(WorkflowEvent( kind="workflow_completed", ts_ms=_now_ms(), payload={ "workflow": workflow, "duration_ms": result.duration_ms, "ok": not result.aborted, }, )) return result # ── Per-step execution ──────────────────────────────── async def _run_step( self, step: Step, ctx: Mapping[str, Any], emit: Callable[[WorkflowEvent], None], ) -> StepResult: try: policy = self._library.policy(step.prompt_id) except PromptLibraryError as exc: emit(WorkflowEvent( kind="step_failed", ts_ms=_now_ms(), payload={ "step_id": step.step_id, "prompt_id": step.prompt_id, "attempt": 1, "reason": f"prompt-load: {exc}", "fatal": True, }, )) raise StepFailure( step_id=step.step_id, prompt_id=step.prompt_id, reason=f"prompt load failed: {exc}", attempts=1, ) from exc max_attempts = max(1, int(policy.retries) + 1) last_reason: str = "" step_started = _now_ms() for attempt in range(1, max_attempts + 1): emit(WorkflowEvent( kind="step_started", ts_ms=_now_ms(), payload={ "step_id": step.step_id, "prompt_id": step.prompt_id, "attempt": attempt, }, )) t0 = _now_ms() try: rendered = self._library.render( step.prompt_id, **dict(step.build_vars(ctx)), ) except PromptLibraryError as exc: last_reason = f"render: {exc}" emit(_fail_event(step, attempt, last_reason, fatal=False)) break # can't retry a deterministic render failure try: response = await call_prompt( rendered, policy, temperature=float(step.temperature), max_tokens=int(step.max_tokens), model_override=(step.model or None), ) except asyncio.TimeoutError: last_reason = f"timeout after {policy.timeout_s:.1f}s" emit(_fail_event(step, attempt, last_reason, fatal=False)) continue except Exception as exc: # noqa: BLE001 last_reason = f"llm: {exc.__class__.__name__}: {str(exc)[:180]}" emit(_fail_event(step, attempt, last_reason, fatal=False)) continue content = extract_content(response) if not content: last_reason = "empty response" emit(_fail_event(step, attempt, last_reason, fatal=False)) continue try: parsed = step.parse(content) except Exception as exc: # noqa: BLE001 last_reason = f"parse: {str(exc)[:180]}" emit(_fail_event(step, attempt, last_reason, fatal=False)) continue verdict = step.validate(parsed) if verdict is not None: last_reason = f"validate: {verdict}" emit(_fail_event(step, attempt, last_reason, fatal=False)) continue # Success. duration = max(0, _now_ms() - t0) emit(WorkflowEvent( kind="step_completed", ts_ms=_now_ms(), payload={ "step_id": step.step_id, "prompt_id": step.prompt_id, "attempt": attempt, "duration_ms": duration, "preview": _preview(parsed), }, )) return StepResult( step_id=step.step_id, prompt_id=step.prompt_id, value=parsed, attempts=attempt, used_fallback=False, duration_ms=max(0, _now_ms() - step_started), ) # All attempts exhausted. Consult fallback policy. fallback = (policy.fallback or "").strip().lower() if fallback == "abort" or not fallback: emit(_fail_event(step, max_attempts, last_reason, fatal=True)) raise StepFailure( step_id=step.step_id, prompt_id=step.prompt_id, reason=last_reason or "unknown failure", attempts=max_attempts, ) if step.fallback is None: # Policy said non-abort but the Step didn't wire a # fallback callable — treat as fatal so we don't # silently ship garbage. reason = f"{last_reason} (no Step.fallback configured for '{fallback}')" emit(_fail_event(step, max_attempts, reason, fatal=True)) raise StepFailure( step_id=step.step_id, prompt_id=step.prompt_id, reason=reason, attempts=max_attempts, ) try: value = step.fallback(ctx, policy.fallback) except StepFailure: raise except Exception as exc: # noqa: BLE001 reason = f"fallback raised: {exc}" emit(_fail_event(step, max_attempts, reason, fatal=True)) raise StepFailure( step_id=step.step_id, prompt_id=step.prompt_id, reason=reason, attempts=max_attempts, ) from exc emit(WorkflowEvent( kind="step_completed", ts_ms=_now_ms(), payload={ "step_id": step.step_id, "prompt_id": step.prompt_id, "attempt": max_attempts, "duration_ms": 0, "preview": _preview(value), "used_fallback": True, }, )) return StepResult( step_id=step.step_id, prompt_id=step.prompt_id, value=value, attempts=max_attempts, used_fallback=True, duration_ms=max(0, _now_ms() - step_started), ) # ── Private helpers ──────────────────────────────────────────── def _now_ms() -> int: return int(time.time() * 1000) def _preview(value: Any) -> str: try: s = str(value) except Exception: # noqa: BLE001 return "" s = s.replace("\n", " ").strip() return s if len(s) <= 80 else s[:77] + "…" def _fail_event( step: Step, attempt: int, reason: str, *, fatal: bool, ) -> WorkflowEvent: return WorkflowEvent( kind="step_failed", ts_ms=_now_ms(), payload={ "step_id": step.step_id, "prompt_id": step.prompt_id, "attempt": attempt, "reason": reason, "fatal": fatal, }, )