From 2952eb7317e239585ad706f8028d65602cdc093e Mon Sep 17 00:00:00 2001
From: Vyacheslav Matyukhin <me@berekuk.ru>
Date: Sat, 16 Nov 2024 15:53:23 -0300
Subject: [PATCH 1/7] WorkflowTemplate.getInitialStep, more declarative
 approach

---
 packages/ai/src/LLMStepInstance.ts            |  2 +-
 packages/ai/src/LLMStepTemplate.ts            | 34 +++++++++++++++++++
 packages/ai/src/workflows/Workflow.ts         | 22 ++++++------
 packages/ai/src/workflows/WorkflowTemplate.ts | 16 ++++-----
 packages/ai/src/workflows/controllers.ts      | 12 ++-----
 .../createSquiggleWorkflowTemplate.ts         |  5 ++-
 .../workflows/fixSquiggleWorkflowTemplate.ts  |  4 +--
 7 files changed, 61 insertions(+), 34 deletions(-)

diff --git a/packages/ai/src/LLMStepInstance.ts b/packages/ai/src/LLMStepInstance.ts
index e72d2af3de..a5bed3cd14 100644
--- a/packages/ai/src/LLMStepInstance.ts
+++ b/packages/ai/src/LLMStepInstance.ts
@@ -56,7 +56,7 @@ export class LLMStepInstance<
 
   // These two fields are not serialized
   private logger: Logger;
-  private workflow: Workflow<WorkflowShape>;
+  public workflow: Workflow<WorkflowShape>;
 
   private constructor(
     params: StepParams<Shape> & {
diff --git a/packages/ai/src/LLMStepTemplate.ts b/packages/ai/src/LLMStepTemplate.ts
index f34105160c..2ca0d400ba 100644
--- a/packages/ai/src/LLMStepTemplate.ts
+++ b/packages/ai/src/LLMStepTemplate.ts
@@ -50,6 +50,33 @@ export type ExecuteContext<Shape extends IOShape> = {
   fail(errorType: ErrorType, message: string): void;
 };
 
+/*
+ * This is a type that's used to prepare the step for execution. `PreparedStep`
+ * will be converted by `Workflow.addStep()` to `LLMStepInstance`.
+ *
+ * Notes:
+ * 1) We can't just use LLMStepInstance here, this would cause nasty circular
+ * dependencies (sometimes even Node.js crashes! seriously).
+ * 2) We can't use a pair of template and inputs, because their generic
+ * parameters must be synchronized, and TypeScript can't express that. So we
+ * need to produce a PreparedStep from a template method.
+ *
+ * So, in total, we have three related types:
+ * - `LLMStepTemplate`
+ * - `LLMStepInstance`
+ * - `PreparedStep`
+ *
+ * (And also `StepParams` and `SerializedStep` in `LLMStepInstance`...)
+ *
+ * If this sounds messy, it is. I really tried to find other approaches, but
+ * fixing this would require giving up on `Shape` generic parameters, which has
+ * its own issues.
+ */
+export type PreparedStep<Shape extends IOShape> = {
+  template: LLMStepTemplate<Shape>;
+  inputs: Inputs<Shape>;
+};
+
 export class LLMStepTemplate<const Shape extends IOShape = IOShape> {
   constructor(
     public readonly name: string,
@@ -59,4 +86,11 @@ export class LLMStepTemplate<const Shape extends IOShape = IOShape> {
       inputs: Inputs<Shape>
     ) => Promise<void>
   ) {}
+
+  prepare(inputs: Inputs<Shape>): PreparedStep<Shape> {
+    return {
+      template: this,
+      inputs,
+    };
+  }
 }
diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts
index fa9da62d36..740b4813f5 100644
--- a/packages/ai/src/workflows/Workflow.ts
+++ b/packages/ai/src/workflows/Workflow.ts
@@ -8,7 +8,7 @@ import {
   Message,
 } from "../LLMClient.js";
 import { LLMStepInstance } from "../LLMStepInstance.js";
-import { Inputs, IOShape, LLMStepTemplate } from "../LLMStepTemplate.js";
+import { Inputs, IOShape, PreparedStep } from "../LLMStepTemplate.js";
 import { TimestampedLogEntry } from "../Logger.js";
 import { LlmId } from "../modelConfigs.js";
 import {
@@ -142,7 +142,9 @@ export class Workflow<Shape extends IOShape = IOShape> {
   private configure() {
     // we configure the controller loop first, so it has a chance to react to its initial step
     this.template.configureControllerLoop(this, this.inputs);
-    this.template.configureInitialSteps(this, this.inputs);
+
+    const initialStep = this.template.getInitialStep(this);
+    this.addStep(initialStep);
   }
 
   // Run workflow to the ReadableStream, appropriate for streaming in Next.js routes
@@ -190,14 +192,14 @@ export class Workflow<Shape extends IOShape = IOShape> {
     return this.steps[index - 1];
   }
 
-  addStep<S extends IOShape>(
-    template: LLMStepTemplate<S>,
-    inputs: Inputs<S>
+  private addStep<S extends IOShape>(
+    prepatedStep: PreparedStep<S>
   ): LLMStepInstance<S, Shape> {
-    // sorry for "any"; contravariance issues
+    // `any` is necessary because of countervariance issues.
+    // But that's not important because `PreparedStep` was already strictly typed.
     const step: LLMStepInstance<any, Shape> = LLMStepInstance.create({
-      template,
-      inputs,
+      template: prepatedStep.template,
+      inputs: prepatedStep.inputs,
       workflow: this,
     });
 
@@ -219,10 +221,10 @@ export class Workflow<Shape extends IOShape = IOShape> {
       const result = produce(step, new WorkflowGuardHelpers(this, step));
       switch (result.kind) {
         case "repeat":
-          this.addStep(step.template, step.inputs);
+          this.addStep(step.template.prepare(step.inputs));
           break;
         case "step":
-          this.addStep(result.step, result.inputs);
+          this.addStep(result.step.prepare(result.inputs));
           break;
         case "finish":
           // no new steps to add
diff --git a/packages/ai/src/workflows/WorkflowTemplate.ts b/packages/ai/src/workflows/WorkflowTemplate.ts
index 50c8be8ea4..27a00ff67d 100644
--- a/packages/ai/src/workflows/WorkflowTemplate.ts
+++ b/packages/ai/src/workflows/WorkflowTemplate.ts
@@ -1,5 +1,5 @@
 import { LLMStepInstance } from "../LLMStepInstance.js";
-import { Inputs, IOShape } from "../LLMStepTemplate.js";
+import { Inputs, IOShape, PreparedStep } from "../LLMStepTemplate.js";
 import { type LlmConfig, Workflow } from "./Workflow.js";
 
 export type WorkflowInstanceParams<Shape extends IOShape> = {
@@ -19,26 +19,24 @@ export type WorkflowInstanceParams<Shape extends IOShape> = {
  * It works similarly to LLMStepTemplate, but for workflows.
  */
 export class WorkflowTemplate<const Shape extends IOShape> {
-  public readonly name: string;
+  readonly name: string;
 
-  public configureControllerLoop: (
-    workflow: Workflow<Shape>,
-    inputs: Inputs<Shape>
-  ) => void;
-  public configureInitialSteps: (
+  configureControllerLoop: (
     workflow: Workflow<Shape>,
     inputs: Inputs<Shape>
   ) => void;
 
+  getInitialStep: (workflow: Workflow<Shape>) => PreparedStep<any>;
+
   // TODO - shape parameter
   constructor(params: {
     name: string;
     // TODO - do we need two separate functions? we always call them together
     configureControllerLoop: WorkflowTemplate<Shape>["configureControllerLoop"];
-    configureInitialSteps: WorkflowTemplate<Shape>["configureInitialSteps"];
+    getInitialStep: WorkflowTemplate<Shape>["getInitialStep"];
   }) {
     this.name = params.name;
-    this.configureInitialSteps = params.configureInitialSteps;
+    this.getInitialStep = params.getInitialStep;
     this.configureControllerLoop = params.configureControllerLoop;
   }
 
diff --git a/packages/ai/src/workflows/controllers.ts b/packages/ai/src/workflows/controllers.ts
index b83aa00b4e..2c3370ca84 100644
--- a/packages/ai/src/workflows/controllers.ts
+++ b/packages/ai/src/workflows/controllers.ts
@@ -14,11 +14,6 @@ import {
 
 const MAX_MINOR_ERRORS = 5;
 
-type config = {
-  maxNumericSteps: number;
-  maxStyleGuideSteps: number;
-};
-
 // Error Messages
 const ERROR_MESSAGES = {
   FAILED_STEP: (stepName: string, errorType: string) =>
@@ -32,8 +27,7 @@ const ERROR_MESSAGES = {
 // Helper function to handle failed states
 function handleFailedState<Shape extends IOShape>(
   step: LLMStepInstance<IOShape, Shape>,
-  h: WorkflowGuardHelpers<Shape>,
-  config: config
+  h: WorkflowGuardHelpers<Shape>
 ): NextStepAction | undefined {
   const state = step.getState();
 
@@ -66,7 +60,7 @@ export function fixAdjustRetryLoop<Shape extends IOShape>(
     maxStyleGuideSteps: workflow.llmConfig.styleGuideSteps,
   };
   workflow.addLinearRule((step, h) => {
-    function getNextIntendedState<Shape extends IOShape>(
+    function getNextIntendedState(
       intendedStep: "AdjustToFeedback" | "MatchStyleGuide",
       code: CodeArtifact
     ): NextStepAction {
@@ -85,7 +79,7 @@ export function fixAdjustRetryLoop<Shape extends IOShape>(
     }
 
     // process bad states
-    const failedState = handleFailedState(step, h, config);
+    const failedState = handleFailedState(step, h);
     if (failedState) return failedState;
 
     function fixCodeOrAdjustToFeedback(
diff --git a/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts b/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
index 89cbbf6c81..752da8ac09 100644
--- a/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
+++ b/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
@@ -19,7 +19,6 @@ export const createSquiggleWorkflowTemplate = new WorkflowTemplate<{
   configureControllerLoop(workflow, inputs) {
     fixAdjustRetryLoop(workflow, inputs.prompt);
   },
-  configureInitialSteps(workflow, inputs) {
-    workflow.addStep(generateCodeStep, { prompt: inputs.prompt });
-  },
+  getInitialStep: (workflow) =>
+    generateCodeStep.prepare({ prompt: workflow.inputs.prompt }),
 });
diff --git a/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts b/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
index 1f4a4022b3..4777a9debf 100644
--- a/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
+++ b/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
@@ -15,7 +15,7 @@ export const fixSquiggleWorkflowTemplate = new WorkflowTemplate<{
     // (but it's better to just refactor steps to make the prompt optional, somehow)
     fixAdjustRetryLoop(workflow, new PromptArtifact(""));
   },
-  configureInitialSteps(workflow, inputs) {
-    workflow.addStep(runAndFormatCodeStep, { source: inputs.source });
+  getInitialStep(workflow) {
+    return runAndFormatCodeStep.prepare({ source: workflow.inputs.source });
   },
 });

From 2c39050e494fb86547b5fed0569f98b23e890ae6 Mon Sep 17 00:00:00 2001
From: Vyacheslav Matyukhin <me@berekuk.ru>
Date: Sat, 16 Nov 2024 16:20:43 -0300
Subject: [PATCH 2/7] rename/refactor: StepTransitionRule

---
 packages/ai/src/workflows/Workflow.ts         | 55 ++++++++++---------
 packages/ai/src/workflows/WorkflowTemplate.ts | 12 ++--
 packages/ai/src/workflows/controllers.ts      |  8 +--
 .../createSquiggleWorkflowTemplate.ts         |  5 +-
 .../workflows/fixSquiggleWorkflowTemplate.ts  |  9 ++-
 5 files changed, 43 insertions(+), 46 deletions(-)

diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts
index 740b4813f5..f811696147 100644
--- a/packages/ai/src/workflows/Workflow.ts
+++ b/packages/ai/src/workflows/Workflow.ts
@@ -96,6 +96,11 @@ export type WorkflowEventListener<
   Shape extends IOShape,
 > = (event: WorkflowEvent<T, Shape>) => void;
 
+export type StepTransitionRule<Shape extends IOShape> = (
+  step: LLMStepInstance<IOShape, Shape>,
+  helpers: WorkflowGuardHelpers<Shape>
+) => NextStepAction;
+
 /**
  * This class is responsible for managing the steps in a workflow.
  *
@@ -140,9 +145,31 @@ export class Workflow<Shape extends IOShape = IOShape> {
   }
 
   private configure() {
-    // we configure the controller loop first, so it has a chance to react to its initial step
-    this.template.configureControllerLoop(this, this.inputs);
+    // configure the controller loop
+    {
+      const transitionRule = this.template.getTransitionRule(this);
+      this.addEventListener("stepFinished", ({ data: { step } }) => {
+        const result = transitionRule(
+          step,
+          new WorkflowGuardHelpers(this, step)
+        );
+        switch (result.kind) {
+          case "repeat":
+            this.addStep(step.template.prepare(step.inputs));
+            break;
+          case "step":
+            this.addStep(result.step.prepare(result.inputs));
+            break;
+          case "finish":
+            // no new steps to add
+            break;
+          case "fatal":
+            throw new Error(result.message);
+        }
+      });
+    }
 
+    // add the first step
     const initialStep = this.template.getInitialStep(this);
     this.addStep(initialStep);
   }
@@ -211,30 +238,6 @@ export class Workflow<Shape extends IOShape = IOShape> {
     return step;
   }
 
-  addLinearRule(
-    produce: (
-      step: LLMStepInstance<IOShape, Shape>,
-      helpers: WorkflowGuardHelpers<Shape>
-    ) => NextStepAction
-  ) {
-    this.addEventListener("stepFinished", ({ data: { step } }) => {
-      const result = produce(step, new WorkflowGuardHelpers(this, step));
-      switch (result.kind) {
-        case "repeat":
-          this.addStep(step.template.prepare(step.inputs));
-          break;
-        case "step":
-          this.addStep(result.step.prepare(result.inputs));
-          break;
-        case "finish":
-          // no new steps to add
-          break;
-        case "fatal":
-          throw new Error(result.message);
-      }
-    });
-  }
-
   private async runNextStep(): Promise<void> {
     const step = this.getCurrentStep();
 
diff --git a/packages/ai/src/workflows/WorkflowTemplate.ts b/packages/ai/src/workflows/WorkflowTemplate.ts
index 27a00ff67d..5aa8637d06 100644
--- a/packages/ai/src/workflows/WorkflowTemplate.ts
+++ b/packages/ai/src/workflows/WorkflowTemplate.ts
@@ -1,6 +1,6 @@
 import { LLMStepInstance } from "../LLMStepInstance.js";
 import { Inputs, IOShape, PreparedStep } from "../LLMStepTemplate.js";
-import { type LlmConfig, Workflow } from "./Workflow.js";
+import { type LlmConfig, StepTransitionRule, Workflow } from "./Workflow.js";
 
 export type WorkflowInstanceParams<Shape extends IOShape> = {
   id: string;
@@ -21,23 +21,19 @@ export type WorkflowInstanceParams<Shape extends IOShape> = {
 export class WorkflowTemplate<const Shape extends IOShape> {
   readonly name: string;
 
-  configureControllerLoop: (
-    workflow: Workflow<Shape>,
-    inputs: Inputs<Shape>
-  ) => void;
-
+  getTransitionRule: (workflow: Workflow<Shape>) => StepTransitionRule<Shape>;
   getInitialStep: (workflow: Workflow<Shape>) => PreparedStep<any>;
 
   // TODO - shape parameter
   constructor(params: {
     name: string;
     // TODO - do we need two separate functions? we always call them together
-    configureControllerLoop: WorkflowTemplate<Shape>["configureControllerLoop"];
+    getTransitionRule: WorkflowTemplate<Shape>["getTransitionRule"];
     getInitialStep: WorkflowTemplate<Shape>["getInitialStep"];
   }) {
     this.name = params.name;
     this.getInitialStep = params.getInitialStep;
-    this.configureControllerLoop = params.configureControllerLoop;
+    this.getTransitionRule = params.getTransitionRule;
   }
 
   instantiate(
diff --git a/packages/ai/src/workflows/controllers.ts b/packages/ai/src/workflows/controllers.ts
index 2c3370ca84..cf5d4f93bb 100644
--- a/packages/ai/src/workflows/controllers.ts
+++ b/packages/ai/src/workflows/controllers.ts
@@ -6,7 +6,7 @@ import { fixCodeUntilItRunsStep } from "../steps/fixCodeUntilItRunsStep.js";
 import { generateCodeStep } from "../steps/generateCodeStep.js";
 import { matchStyleGuideStep } from "../steps/matchStyleGuideStep.js";
 import { runAndFormatCodeStep } from "../steps/runAndFormatCodeStep.js";
-import { Workflow } from "./Workflow.js";
+import { StepTransitionRule, Workflow } from "./Workflow.js";
 import {
   NextStepAction,
   WorkflowGuardHelpers,
@@ -54,12 +54,12 @@ function handleFailedState<Shape extends IOShape>(
 export function fixAdjustRetryLoop<Shape extends IOShape>(
   workflow: Workflow<Shape>,
   prompt: PromptArtifact
-) {
+): StepTransitionRule<Shape> {
   const config = {
     maxNumericSteps: workflow.llmConfig.numericSteps,
     maxStyleGuideSteps: workflow.llmConfig.styleGuideSteps,
   };
-  workflow.addLinearRule((step, h) => {
+  return (step, h) => {
     function getNextIntendedState(
       intendedStep: "AdjustToFeedback" | "MatchStyleGuide",
       code: CodeArtifact
@@ -143,5 +143,5 @@ export function fixAdjustRetryLoop<Shape extends IOShape>(
     }
 
     return h.fatal("Unknown step");
-  });
+  };
 }
diff --git a/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts b/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
index 752da8ac09..f5bcb40918 100644
--- a/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
+++ b/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
@@ -16,9 +16,8 @@ export const createSquiggleWorkflowTemplate = new WorkflowTemplate<{
   outputs: Record<string, never>;
 }>({
   name: "CreateSquiggle",
-  configureControllerLoop(workflow, inputs) {
-    fixAdjustRetryLoop(workflow, inputs.prompt);
-  },
+  getTransitionRule: (workflow) =>
+    fixAdjustRetryLoop(workflow, workflow.inputs.prompt),
   getInitialStep: (workflow) =>
     generateCodeStep.prepare({ prompt: workflow.inputs.prompt }),
 });
diff --git a/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts b/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
index 4777a9debf..795186fc7f 100644
--- a/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
+++ b/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
@@ -10,12 +10,11 @@ export const fixSquiggleWorkflowTemplate = new WorkflowTemplate<{
   outputs: Record<string, never>;
 }>({
   name: "FixSquiggle",
-  configureControllerLoop(workflow) {
+  getTransitionRule: (workflow) => {
     // TODO - cache the prompt artifact once? maybe even as a global variable
     // (but it's better to just refactor steps to make the prompt optional, somehow)
-    fixAdjustRetryLoop(workflow, new PromptArtifact(""));
-  },
-  getInitialStep(workflow) {
-    return runAndFormatCodeStep.prepare({ source: workflow.inputs.source });
+    return fixAdjustRetryLoop(workflow, new PromptArtifact(""));
   },
+  getInitialStep: (workflow) =>
+    runAndFormatCodeStep.prepare({ source: workflow.inputs.source }),
 });

From e1070f96f65c4ccddd6ea6c80b47d9b71d3aeb94 Mon Sep 17 00:00:00 2001
From: Vyacheslav Matyukhin <me@berekuk.ru>
Date: Sat, 16 Nov 2024 16:50:34 -0300
Subject: [PATCH 3/7] various cleanups

---
 packages/ai/src/LLMStepInstance.ts            | 10 +++
 packages/ai/src/generateSummary.ts            | 16 ----
 packages/ai/src/workflows/Workflow.ts         | 83 +++++++------------
 .../ai/src/workflows/WorkflowGuardHelpers.ts  | 16 ++--
 packages/ai/src/workflows/WorkflowTemplate.ts |  7 +-
 .../createSquiggleWorkflowTemplate.ts         |  4 +-
 .../workflows/fixSquiggleWorkflowTemplate.ts  |  4 +-
 packages/ai/src/workflows/streaming.ts        |  2 -
 8 files changed, 58 insertions(+), 84 deletions(-)

diff --git a/packages/ai/src/LLMStepInstance.ts b/packages/ai/src/LLMStepInstance.ts
index a5bed3cd14..29ffcb5703 100644
--- a/packages/ai/src/LLMStepInstance.ts
+++ b/packages/ai/src/LLMStepInstance.ts
@@ -211,6 +211,16 @@ export class LLMStepInstance<
     return this.state.kind === "DONE";
   }
 
+  getPreviousStep() {
+    const steps = this.workflow.getSteps();
+    const index = steps.indexOf(this as LLMStepInstance<any, WorkflowShape>);
+    if (index < 1) {
+      // first step or not found
+      return undefined;
+    }
+    return steps[index - 1];
+  }
+
   // private methods
 
   private setOutput<K extends Extract<keyof Shape["outputs"], string>>(
diff --git a/packages/ai/src/generateSummary.ts b/packages/ai/src/generateSummary.ts
index f890a4592b..549b9c955e 100644
--- a/packages/ai/src/generateSummary.ts
+++ b/packages/ai/src/generateSummary.ts
@@ -1,6 +1,3 @@
-import fs from "fs";
-import path from "path";
-
 import { Artifact, ArtifactKind } from "./Artifact.js";
 import { Code } from "./Code.js";
 import { calculatePriceMultipleCalls } from "./LLMClient.js";
@@ -227,16 +224,3 @@ ${code.source}
 `;
   }
 }
-
-export function saveSummaryToFile(summary: string): void {
-  const logDir = path.join(process.cwd(), "logs");
-  if (!fs.existsSync(logDir)) {
-    fs.mkdirSync(logDir, { recursive: true });
-  }
-
-  const timestamp = new Date().toISOString().replace(/:/g, "-");
-  const logFile = path.join(logDir, `squiggle_summary_${timestamp}.md`);
-
-  fs.writeFileSync(logFile, summary);
-  console.log(`Summary saved to ${logFile}`);
-}
diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts
index f811696147..703c76cd78 100644
--- a/packages/ai/src/workflows/Workflow.ts
+++ b/packages/ai/src/workflows/Workflow.ts
@@ -112,12 +112,12 @@ export class Workflow<Shape extends IOShape = IOShape> {
   public id: string;
   public readonly template: WorkflowTemplate<Shape>;
   public readonly inputs: Inputs<Shape>;
-  private started: boolean = false;
 
   public llmConfig: LlmConfig;
   public startTime: number;
 
   private steps: LLMStepInstance<IOShape, Shape>[];
+  private transitionRule: StepTransitionRule<Shape>;
 
   public llmClient: LLMClient;
 
@@ -135,40 +135,17 @@ export class Workflow<Shape extends IOShape = IOShape> {
       params.openaiApiKey,
       params.anthropicApiKey
     );
+
+    this.transitionRule = params.template.getTransitionRule(this);
   }
 
   private startOrThrow() {
-    if (this.started) {
+    // This function just inserts the first step.
+    // Previously we had a `started` flag, but that wasn't very useful.
+    // But if we ever implement resumable workflows, we'll need to change this.
+    if (this.steps.length) {
       throw new Error("Workflow already started");
     }
-    this.started = true;
-  }
-
-  private configure() {
-    // configure the controller loop
-    {
-      const transitionRule = this.template.getTransitionRule(this);
-      this.addEventListener("stepFinished", ({ data: { step } }) => {
-        const result = transitionRule(
-          step,
-          new WorkflowGuardHelpers(this, step)
-        );
-        switch (result.kind) {
-          case "repeat":
-            this.addStep(step.template.prepare(step.inputs));
-            break;
-          case "step":
-            this.addStep(result.step.prepare(result.inputs));
-            break;
-          case "finish":
-            // no new steps to add
-            break;
-          case "fatal":
-            throw new Error(result.message);
-        }
-      });
-    }
-
     // add the first step
     const initialStep = this.template.getInitialStep(this);
     this.addStep(initialStep);
@@ -186,10 +163,6 @@ export class Workflow<Shape extends IOShape = IOShape> {
         // handlers, but before we add any steps.
         this.dispatchEvent({ type: "workflowStarted" });
 
-        // Important! `configure` should be called after all event listeners are
-        // set up. We want to capture `stepAdded` events.
-        this.configure();
-
         await this.runUntilComplete();
         controller.close();
       },
@@ -201,24 +174,12 @@ export class Workflow<Shape extends IOShape = IOShape> {
   // Run workflow without streaming, only capture the final result
   async runToResult(): Promise<ClientWorkflowResult> {
     this.startOrThrow();
-    this.configure();
 
     await this.runUntilComplete();
 
-    // saveSummaryToFile(generateSummary(workflow));
     return this.getFinalResult();
   }
 
-  getPreviousStep(
-    step: LLMStepInstance<any, Shape>
-  ): LLMStepInstance<any, Shape> | undefined {
-    const index = this.steps.indexOf(step);
-    if (index < 1) {
-      return undefined;
-    }
-    return this.steps[index - 1];
-  }
-
   private addStep<S extends IOShape>(
     prepatedStep: PreparedStep<S>
   ): LLMStepInstance<S, Shape> {
@@ -241,7 +202,7 @@ export class Workflow<Shape extends IOShape = IOShape> {
   private async runNextStep(): Promise<void> {
     const step = this.getCurrentStep();
 
-    if (!step) {
+    if (!step || step.getState().kind !== "PENDING") {
       return;
     }
 
@@ -256,9 +217,31 @@ export class Workflow<Shape extends IOShape = IOShape> {
       type: "stepFinished",
       payload: { step },
     });
+
+    // apply the transition rule, produce the next step in PENDING state
+    // this code is inlined in this method, which guarantees that we always have one pending step
+    // (until the transition rule decides to finish)
+    const result = this.transitionRule(step, new WorkflowGuardHelpers(step));
+    switch (result.kind) {
+      case "repeat":
+        this.addStep(step.template.prepare(step.inputs));
+        break;
+      case "step":
+        this.addStep(result.step.prepare(result.inputs));
+        break;
+      case "finish":
+        // no new steps to add
+        break;
+      case "fatal":
+        throw new Error(result.message);
+    }
   }
 
   async runUntilComplete() {
+    if (!this.steps.length) {
+      throw new Error("Workflow not started");
+    }
+
     while (!this.isProcessComplete()) {
       await this.runNextStep();
     }
@@ -319,10 +302,8 @@ export class Workflow<Shape extends IOShape = IOShape> {
           if (output.value.type === "success") {
             return { step, code: output.value.source };
           }
-          // Otherwise store the first step with any code
-          if (!stepWithAnyCode) {
-            stepWithAnyCode = { step, code: output.value.source };
-          }
+          // Store the first step with any code
+          stepWithAnyCode ??= { step, code: output.value.source };
         }
       }
     }
diff --git a/packages/ai/src/workflows/WorkflowGuardHelpers.ts b/packages/ai/src/workflows/WorkflowGuardHelpers.ts
index 32e82c36ee..18981cc9db 100644
--- a/packages/ai/src/workflows/WorkflowGuardHelpers.ts
+++ b/packages/ai/src/workflows/WorkflowGuardHelpers.ts
@@ -31,10 +31,10 @@ addRule({
 Where totalRepeats counts the number of template runs in the entire history, while receptRepeats counts the number of immediate repeats in the step -> parent -> parent chain. (I'm stubborn and still trying to make it parallelism-agnostic, so thinking in terms of graphs, not lists). 
 */
 export class WorkflowGuardHelpers<Shape extends IOShape> {
-  constructor(
-    private readonly workflow: Workflow<Shape>,
-    private readonly _step: LLMStepInstance<any, Shape>
-  ) {}
+  private workflow: Workflow<Shape>;
+  constructor(private readonly _step: LLMStepInstance<any, Shape>) {
+    this.workflow = _step.workflow;
+  }
 
   totalRepeats(template: LLMStepTemplate<any>) {
     return this.workflow.getSteps().filter((step) => step.instanceOf(template))
@@ -43,20 +43,20 @@ export class WorkflowGuardHelpers<Shape extends IOShape> {
 
   recentRepeats(template: LLMStepTemplate<any>) {
     let count = -1;
-    let step = this.workflow.getPreviousStep(this._step);
+    let step = this._step.getPreviousStep();
     while (step?.instanceOf(template)) {
       count++;
-      step = this.workflow.getPreviousStep(step);
+      step = step.getPreviousStep();
     }
     return count;
   }
 
   recentFailedRepeats(template: LLMStepTemplate<any>) {
     let count = -1;
-    let step = this.workflow.getPreviousStep(this._step);
+    let step = this._step.getPreviousStep();
     while (step?.instanceOf(template) && step.getState().kind === "FAILED") {
       count++;
-      step = this.workflow.getPreviousStep(step);
+      step = step.getPreviousStep();
     }
     return count;
   }
diff --git a/packages/ai/src/workflows/WorkflowTemplate.ts b/packages/ai/src/workflows/WorkflowTemplate.ts
index 5aa8637d06..ceccaacdb1 100644
--- a/packages/ai/src/workflows/WorkflowTemplate.ts
+++ b/packages/ai/src/workflows/WorkflowTemplate.ts
@@ -21,15 +21,16 @@ export type WorkflowInstanceParams<Shape extends IOShape> = {
 export class WorkflowTemplate<const Shape extends IOShape> {
   readonly name: string;
 
-  getTransitionRule: (workflow: Workflow<Shape>) => StepTransitionRule<Shape>;
   getInitialStep: (workflow: Workflow<Shape>) => PreparedStep<any>;
+  getTransitionRule: (workflow: Workflow<Shape>) => StepTransitionRule<Shape>;
 
   // TODO - shape parameter
   constructor(params: {
     name: string;
-    // TODO - do we need two separate functions? we always call them together
-    getTransitionRule: WorkflowTemplate<Shape>["getTransitionRule"];
+    // This function will be called to obtain the first step of the workflow.
     getInitialStep: WorkflowTemplate<Shape>["getInitialStep"];
+    // This function will be called to obtain the next step of the workflow based on the current step.
+    getTransitionRule: WorkflowTemplate<Shape>["getTransitionRule"];
   }) {
     this.name = params.name;
     this.getInitialStep = params.getInitialStep;
diff --git a/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts b/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
index f5bcb40918..cc7d80ca1b 100644
--- a/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
+++ b/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
@@ -16,8 +16,8 @@ export const createSquiggleWorkflowTemplate = new WorkflowTemplate<{
   outputs: Record<string, never>;
 }>({
   name: "CreateSquiggle",
-  getTransitionRule: (workflow) =>
-    fixAdjustRetryLoop(workflow, workflow.inputs.prompt),
   getInitialStep: (workflow) =>
     generateCodeStep.prepare({ prompt: workflow.inputs.prompt }),
+  getTransitionRule: (workflow) =>
+    fixAdjustRetryLoop(workflow, workflow.inputs.prompt),
 });
diff --git a/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts b/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
index 795186fc7f..7648085678 100644
--- a/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
+++ b/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
@@ -10,11 +10,11 @@ export const fixSquiggleWorkflowTemplate = new WorkflowTemplate<{
   outputs: Record<string, never>;
 }>({
   name: "FixSquiggle",
+  getInitialStep: (workflow) =>
+    runAndFormatCodeStep.prepare({ source: workflow.inputs.source }),
   getTransitionRule: (workflow) => {
     // TODO - cache the prompt artifact once? maybe even as a global variable
     // (but it's better to just refactor steps to make the prompt optional, somehow)
     return fixAdjustRetryLoop(workflow, new PromptArtifact(""));
   },
-  getInitialStep: (workflow) =>
-    runAndFormatCodeStep.prepare({ source: workflow.inputs.source }),
 });
diff --git a/packages/ai/src/workflows/streaming.ts b/packages/ai/src/workflows/streaming.ts
index 230780be0b..a64b0424ce 100644
--- a/packages/ai/src/workflows/streaming.ts
+++ b/packages/ai/src/workflows/streaming.ts
@@ -134,8 +134,6 @@ export function addStreamingListeners<Shape extends IOShape>(
     });
   });
   workflow.addEventListener("allStepsFinished", (event) => {
-    // saveSummaryToFile(generateSummary(prompt, workflow));
-
     send({
       kind: "finalResult",
       content: event.workflow.getFinalResult(),

From 28f017ac5c7a12e52fbc01714ba4f26297bf44d0 Mon Sep 17 00:00:00 2001
From: Vyacheslav Matyukhin <me@berekuk.ru>
Date: Tue, 19 Nov 2024 15:36:16 -0300
Subject: [PATCH 4/7] simplified steps API; optional outputs; fix logs for
 finished workflows

---
 packages/ai/src/LLMStepInstance.ts            | 183 ++++++++++--------
 packages/ai/src/LLMStepTemplate.ts            |  40 ++--
 packages/ai/src/generateSummary.ts            |  13 +-
 packages/ai/src/steps/adjustToFeedbackStep.ts |  14 +-
 .../ai/src/steps/fixCodeUntilItRunsStep.ts    |  23 ++-
 packages/ai/src/steps/generateCodeStep.ts     |  19 +-
 packages/ai/src/steps/matchStyleGuideStep.ts  |  14 +-
 packages/ai/src/steps/runAndFormatCodeStep.ts |   4 +-
 packages/ai/src/workflows/Workflow.ts         |   7 +-
 packages/ai/src/workflows/controllers.ts      |  43 ++--
 .../createSquiggleWorkflowTemplate.ts         |   4 +-
 .../workflows/fixSquiggleWorkflowTemplate.ts  |   4 +-
 packages/ai/src/workflows/streaming.ts        |  34 ++--
 packages/hub/src/server/ai/storage.ts         |  10 +-
 14 files changed, 232 insertions(+), 180 deletions(-)

diff --git a/packages/ai/src/LLMStepInstance.ts b/packages/ai/src/LLMStepInstance.ts
index 29ffcb5703..fda0aae38b 100644
--- a/packages/ai/src/LLMStepInstance.ts
+++ b/packages/ai/src/LLMStepInstance.ts
@@ -1,4 +1,4 @@
-import { BaseArtifact, makeArtifact } from "./Artifact.js";
+import { ArtifactKind, BaseArtifact, makeArtifact } from "./Artifact.js";
 import {
   calculatePriceMultipleCalls,
   LlmMetrics,
@@ -10,6 +10,7 @@ import {
   Inputs,
   IOShape,
   LLMStepTemplate,
+  OutputKind,
   Outputs,
   StepState,
 } from "./LLMStepTemplate.js";
@@ -27,14 +28,22 @@ export type StepParams<Shape extends IOShape> = {
   id: string;
   sequentialId: number;
   template: LLMStepTemplate<Shape>;
-  state: StepState;
+  state: StepState<Shape>;
   inputs: Inputs<Shape>;
-  outputs: Partial<Outputs<Shape>>;
   startTime: number;
   conversationMessages: Message[];
   llmMetricsList: LlmMetrics[];
 };
 
+class FailError extends Error {
+  constructor(
+    public type: ErrorType,
+    message: string
+  ) {
+    super(message);
+  }
+}
+
 export class LLMStepInstance<
   const Shape extends IOShape = IOShape,
   const WorkflowShape extends IOShape = IOShape,
@@ -43,10 +52,8 @@ export class LLMStepInstance<
   public sequentialId: StepParams<Shape>["sequentialId"];
   public readonly template: StepParams<Shape>["template"];
 
-  private state: StepParams<Shape>["state"];
-
   // must be public for `instanceOf` type guard to work
-  public readonly _outputs: StepParams<Shape>["outputs"];
+  private _state: StepParams<Shape>["state"];
 
   public readonly inputs: StepParams<Shape>["inputs"];
 
@@ -70,8 +77,7 @@ export class LLMStepInstance<
     this.conversationMessages = params.conversationMessages;
 
     this.startTime = params.startTime;
-    this.state = params.state;
-    this._outputs = params.outputs;
+    this._state = params.state;
 
     this.template = params.template;
     this.inputs = params.inputs;
@@ -93,7 +99,6 @@ export class LLMStepInstance<
       llmMetricsList: [],
       startTime: Date.now(),
       state: { kind: "PENDING" },
-      outputs: {},
       ...params,
     });
   }
@@ -113,7 +118,7 @@ export class LLMStepInstance<
   }
 
   async _run() {
-    if (this.state.kind !== "PENDING") {
+    if (this._state.kind !== "PENDING") {
       return;
     }
 
@@ -123,27 +128,49 @@ export class LLMStepInstance<
       return;
     }
 
-    const executeContext: ExecuteContext<Shape> = {
-      setOutput: (key, value) => this.setOutput(key, value),
+    const executeContext: ExecuteContext = {
       log: (log) => this.log(log),
       queryLLM: (promptPair) => this.queryLLM(promptPair),
-      fail: (errorType, message) => this.fail(errorType, message),
+      fail: (errorType, message) => {
+        throw new FailError(errorType, message);
+      },
     };
 
     try {
-      await this.template.execute(executeContext, this.inputs);
-    } catch (error) {
-      this.fail(
-        "MINOR",
-        error instanceof Error ? error.message : String(error)
-      );
-      return;
-    }
-
-    const hasFailed = (this.state as StepState).kind === "FAILED";
+      const result = await this.template.execute(executeContext, this.inputs);
+
+      const outputs = {} as Outputs<Shape>; // will be filled in below
+      Object.keys(result).forEach((key: keyof typeof result) => {
+        const value = result[key];
+        if (value instanceof BaseArtifact) {
+          // already existing artifact - probably passed through from another step
+          outputs[key] = value as any;
+        } else if (!value) {
+          outputs[key] = undefined as any;
+        } else {
+          const outputKind = this.template.shape.outputs[key] as OutputKind;
+          const artifactKind = outputKind.endsWith("?")
+            ? (outputKind.slice(0, -1) as ArtifactKind)
+            : (outputKind as ArtifactKind);
+
+          outputs[key] = makeArtifact(artifactKind, value as any, this) as any;
+        }
+      });
 
-    if (!hasFailed) {
-      this.state = { kind: "DONE", durationMs: this.calculateDuration() };
+      this._state = {
+        kind: "DONE",
+        durationMs: this.calculateDuration(),
+        outputs,
+      };
+    } catch (error) {
+      if (error instanceof FailError) {
+        this.fail(error.type, error.message);
+      } else {
+        this.fail(
+          "MINOR", // TODO - critical?
+          error instanceof Error ? error.message : String(error)
+        );
+      }
     }
   }
 
@@ -155,8 +182,8 @@ export class LLMStepInstance<
 
     await this._run();
 
-    const completionMessage = `Step "${this.template.name}" completed with status: ${this.state.kind}${
-      this.state.kind !== "PENDING" && `, in ${this.state.durationMs / 1000}s`
+    const completionMessage = `Step "${this.template.name}" completed with status: ${this._state.kind}${
+      this._state.kind !== "PENDING" && `, in ${this._state.durationMs / 1000}s`
     }`;
 
     this.log({
@@ -166,15 +193,11 @@ export class LLMStepInstance<
   }
 
   getState() {
-    return this.state;
+    return this._state;
   }
 
   getDuration() {
-    return this.state.kind === "PENDING" ? 0 : this.state.durationMs;
-  }
-
-  getOutputs(): this["_outputs"] {
-    return this._outputs;
+    return this._state.kind === "PENDING" ? 0 : this._state.durationMs;
   }
 
   getInputs() {
@@ -208,7 +231,7 @@ export class LLMStepInstance<
   }
 
   isDone() {
-    return this.state.kind === "DONE";
+    return this._state.kind === "DONE";
   }
 
   getPreviousStep() {
@@ -223,29 +246,6 @@ export class LLMStepInstance<
 
   // private methods
 
-  private setOutput<K extends Extract<keyof Shape["outputs"], string>>(
-    key: K,
-    value: Outputs<Shape>[K] | Outputs<Shape>[K]["value"]
-  ): void {
-    if (key in this._outputs) {
-      this.fail(
-        "CRITICAL",
-        `Output ${key} is already set. This is a bug with the workflow code.`
-      );
-      return;
-    }
-
-    if (value instanceof BaseArtifact) {
-      // already existing artifact - probably passed through from another step
-      this._outputs[key] = value;
-    } else {
-      const kind = this.template.shape.outputs[
-        key
-      ] as Outputs<Shape>[K]["kind"];
-      this._outputs[key] = makeArtifact(kind, value as any, this) as any;
-    }
-  }
-
   private log(log: LogEntry): void {
     this.logger.log(log, {
       workflowId: this.workflow.id,
@@ -255,7 +255,7 @@ export class LLMStepInstance<
 
   private fail(errorType: ErrorType, message: string) {
     this.log({ type: "error", message });
-    this.state = {
+    this._state = {
       kind: "FAILED",
       durationMs: this.calculateDuration(),
       errorType,
@@ -330,14 +330,13 @@ export class LLMStepInstance<
   // Serialization/deserialization
 
   // StepParams don't contain the workflow reference, to to avoid circular dependencies
-  toParams(): StepParams<any> {
+  toParams(): StepParams<Shape> {
     return {
       id: this.id,
       sequentialId: this.sequentialId,
       template: this.template,
-      state: this.state,
+      state: this._state,
       inputs: this.inputs,
-      outputs: this._outputs,
       startTime: this.startTime,
       conversationMessages: this.conversationMessages,
       llmMetricsList: this.llmMetricsList,
@@ -352,7 +351,13 @@ export class LLMStepInstance<
   }
 
   static deserialize(
-    { templateName, inputIds, outputIds, ...params }: SerializedStep,
+    {
+      templateName,
+      inputIds,
+      outputIds: legacyOutputIds,
+      state: serializedState,
+      ...params
+    }: SerializedStep,
     visitor: AiDeserializationVisitor
   ): StepParams<any> {
     const template: LLMStepTemplate<any> = getStepTemplateByName(templateName);
@@ -362,18 +367,24 @@ export class LLMStepInstance<
         visitor.artifact(inputId),
       ])
     );
-    const outputs = Object.fromEntries(
-      Object.entries(outputIds).map(([name, outputId]) => [
-        name,
-        visitor.artifact(outputId),
-      ])
-    );
+    const state: StepState<any> =
+      serializedState.kind === "DONE"
+        ? {
+            ...serializedState,
+            // TODO - remove outputIds
+            outputs: Object.fromEntries(
+              Object.entries(legacyOutputIds ?? serializedState.outputIds).map(
+                ([name, outputId]) => [name, visitor.artifact(outputId)]
+              )
+            ),
+          }
+        : serializedState;
 
     return {
       ...params,
       template,
       inputs,
-      outputs,
+      state,
     };
   }
 }
@@ -381,12 +392,24 @@ export class LLMStepInstance<
 export function serializeStepParams(
   params: StepParams<IOShape>,
   visitor: AiSerializationVisitor
-) {
+): SerializedStep {
   return {
     id: params.id,
     sequentialId: params.sequentialId,
     templateName: params.template.name,
-    state: params.state,
+    state:
+      params.state.kind === "DONE"
+        ? {
+            ...params.state,
+            outputIds: Object.fromEntries(
+              Object.entries(params.state.outputs)
+                .map(([key, output]) =>
+                  output ? [key, visitor.artifact(output)] : undefined
+                )
+                .filter((x) => x !== undefined)
+            ),
+          }
+        : params.state,
     startTime: params.startTime,
     conversationMessages: params.conversationMessages,
     llmMetricsList: params.llmMetricsList,
@@ -396,21 +419,21 @@ export function serializeStepParams(
         visitor.artifact(input),
       ])
     ),
-    outputIds: Object.fromEntries(
-      Object.entries(params.outputs)
-        .map(([key, output]) =>
-          output ? [key, visitor.artifact(output)] : undefined
-        )
-        .filter((x) => x !== undefined)
-    ),
   };
 }
 
+type SerializedState<Shape extends IOShape> =
+  | (Omit<Extract<StepState<Shape>, { kind: "DONE" }>, "outputs"> & {
+      outputIds: Record<string, number>;
+    })
+  | Exclude<StepState<Shape>, { kind: "DONE" }>;
+
 export type SerializedStep = Omit<
   StepParams<IOShape>,
-  "inputs" | "outputs" | "template"
+  "inputs" | "outputs" | "template" | "state"
 > & {
   templateName: string;
   inputIds: Record<string, number>;
-  outputIds: Record<string, number>;
+  outputIds?: Record<string, number>;
+  state: SerializedState<IOShape>;
 };
diff --git a/packages/ai/src/LLMStepTemplate.ts b/packages/ai/src/LLMStepTemplate.ts
index 2ca0d400ba..60df22de26 100644
--- a/packages/ai/src/LLMStepTemplate.ts
+++ b/packages/ai/src/LLMStepTemplate.ts
@@ -4,13 +4,14 @@ import { PromptPair } from "./prompts.js";
 
 export type ErrorType = "CRITICAL" | "MINOR";
 
-export type StepState =
+export type StepState<Shape extends IOShape> =
   | {
       kind: "PENDING";
     }
   | {
       kind: "DONE";
       durationMs: number;
+      outputs: Outputs<Shape>;
     }
   | {
       kind: "FAILED";
@@ -19,9 +20,11 @@ export type StepState =
       message: string;
     };
 
+export type OutputKind = ArtifactKind | `${ArtifactKind}?`;
+
 export type IOShape<
   I extends Record<string, ArtifactKind> = Record<string, ArtifactKind>,
-  O extends Record<string, ArtifactKind> = Record<string, ArtifactKind>,
+  O extends Record<string, OutputKind> = Record<string, OutputKind>,
 > = {
   inputs: I;
   outputs: O;
@@ -31,23 +34,32 @@ export type Inputs<Shape extends IOShape<any, any>> = {
   [K in keyof Shape["inputs"]]: Extract<Artifact, { kind: Shape["inputs"][K] }>;
 };
 
+// Possible output value based on output kind.
+// For example, if the step specifies `code: "code"`, then the value can be
+// a code artifact or a string (the code itself).
+// If the step specifies `code: "code?"`, then the value can be any of those or it can be undefined.
+type AllowedOutputValue<Kind extends OutputKind> = Kind extends ArtifactKind
+  ? Extract<Artifact, { kind: Kind }>
+  : Kind extends `${infer ArtifactKind}?`
+    ? Extract<Artifact, { kind: ArtifactKind }> | undefined
+    : never;
+
 export type Outputs<Shape extends IOShape<any, any>> = {
-  [K in keyof Shape["outputs"]]: Extract<
-    Artifact,
-    { kind: Shape["outputs"][K] }
-  >;
+  [K in keyof Shape["outputs"]]: AllowedOutputValue<Shape["outputs"][K]>;
+};
+
+type StepExecuteResult<Shape extends IOShape> = {
+  [K in keyof Shape["outputs"]]:
+    | AllowedOutputValue<Shape["outputs"][K]>
+    | NonNullable<AllowedOutputValue<Shape["outputs"][K]>>["value"];
 };
 
 // ExecuteContext is the context that's available to the step implementation.
 // We intentionally don't pass the reference to the step implementation, so that steps won't mess with their internal state.
-export type ExecuteContext<Shape extends IOShape> = {
-  setOutput<K extends Extract<keyof Shape["outputs"], string>>(
-    key: K,
-    value: Outputs<Shape>[K] | Outputs<Shape>[K]["value"] // can be either the artifact or the value inside the artifact
-  ): void;
+export type ExecuteContext = {
   queryLLM(promptPair: PromptPair): Promise<string | null>;
   log(log: LogEntry): void;
-  fail(errorType: ErrorType, message: string): void;
+  fail(errorType: ErrorType, message: string): never;
 };
 
 /*
@@ -82,9 +94,9 @@ export class LLMStepTemplate<const Shape extends IOShape = IOShape> {
     public readonly name: string,
     public readonly shape: Shape,
     public readonly execute: (
-      context: ExecuteContext<Shape>,
+      context: ExecuteContext,
       inputs: Inputs<Shape>
-    ) => Promise<void>
+    ) => Promise<StepExecuteResult<Shape>>
   ) {}
 
   prepare(inputs: Inputs<Shape>): PreparedStep<Shape> {
diff --git a/packages/ai/src/generateSummary.ts b/packages/ai/src/generateSummary.ts
index 549b9c955e..99cd91e3e6 100644
--- a/packages/ai/src/generateSummary.ts
+++ b/packages/ai/src/generateSummary.ts
@@ -96,10 +96,15 @@ function generateDetailedStepLogs<Shape extends IOShape>(
       detailedLogs += getFullArtifact(key, artifact);
     }
 
-    detailedLogs += "### Outputs:\n";
-    for (const [key, artifact] of Object.entries(step.getOutputs())) {
-      if (!artifact) continue;
-      detailedLogs += getFullArtifact(key, artifact);
+    {
+      const state = step.getState();
+      if (state.kind === "DONE") {
+        detailedLogs += "### Outputs:\n";
+        for (const [key, artifact] of Object.entries(state.outputs)) {
+          if (!artifact) continue;
+          detailedLogs += getFullArtifact(key, artifact);
+        }
+      }
     }
 
     detailedLogs += "### Logs:\n";
diff --git a/packages/ai/src/steps/adjustToFeedbackStep.ts b/packages/ai/src/steps/adjustToFeedbackStep.ts
index 7b778b5a6b..0c53a170d6 100644
--- a/packages/ai/src/steps/adjustToFeedbackStep.ts
+++ b/packages/ai/src/steps/adjustToFeedbackStep.ts
@@ -59,7 +59,7 @@ export const adjustToFeedbackStep = new LLMStepTemplate(
       code: "code",
     },
     outputs: {
-      code: "code",
+      code: "code?",
     },
   },
   async (context, { prompt, code }) => {
@@ -73,7 +73,7 @@ export const adjustToFeedbackStep = new LLMStepTemplate(
 
     if (!completion) {
       // failed
-      return;
+      return context.fail("CRITICAL", "LLM failed to provide a response");
     }
 
     // handle adjustment response
@@ -87,7 +87,7 @@ export const adjustToFeedbackStep = new LLMStepTemplate(
         type: "info",
         message: "LLM determined no adjustment is needed",
       });
-      return;
+      return { code: undefined };
     }
 
     if (
@@ -101,19 +101,17 @@ export const adjustToFeedbackStep = new LLMStepTemplate(
           message: "FAIL: " + diffResponse.value,
         });
         // try again
-        context.setOutput("code", code);
-        return;
+        return { code };
       }
 
       const adjustedCode = await codeStringToCode(diffResponse.value);
-      context.setOutput("code", adjustedCode);
-      return;
+      return { code: adjustedCode };
     } else {
       context.log({
         type: "info",
         message: "No adjustments provided, considering process complete",
       });
-      return;
+      return { code: undefined };
     }
   }
 );
diff --git a/packages/ai/src/steps/fixCodeUntilItRunsStep.ts b/packages/ai/src/steps/fixCodeUntilItRunsStep.ts
index 0689b504bc..b70705afe6 100644
--- a/packages/ai/src/steps/fixCodeUntilItRunsStep.ts
+++ b/packages/ai/src/steps/fixCodeUntilItRunsStep.ts
@@ -86,16 +86,19 @@ export const fixCodeUntilItRunsStep = new LLMStepTemplate(
     const promptPair = editExistingSquiggleCodePrompt(code.value);
 
     const completion = await context.queryLLM(promptPair);
-    if (completion) {
-      const newCodeResult = await diffCompletionContentToCode(
-        completion,
-        code.value
-      );
-      if (newCodeResult.ok) {
-        context.setOutput("code", newCodeResult.value);
-      } else {
-        context.fail("MINOR", newCodeResult.value);
-      }
+    if (!completion) {
+      // failed
+      return context.fail("CRITICAL", "LLM failed to provide a response");
+    }
+
+    const newCodeResult = await diffCompletionContentToCode(
+      completion,
+      code.value
+    );
+    if (newCodeResult.ok) {
+      return { code: newCodeResult.value };
+    } else {
+      return context.fail("MINOR", newCodeResult.value);
     }
   }
 );
diff --git a/packages/ai/src/steps/generateCodeStep.ts b/packages/ai/src/steps/generateCodeStep.ts
index fde584d932..d57d1e2f63 100644
--- a/packages/ai/src/steps/generateCodeStep.ts
+++ b/packages/ai/src/steps/generateCodeStep.ts
@@ -138,16 +138,15 @@ export const generateCodeStep = new LLMStepTemplate(
     const promptPair = generateNewSquiggleCodePrompt(prompt.value);
     const completion = await context.queryLLM(promptPair);
 
-    if (completion) {
-      const state = await generationCompletionContentToCode(completion);
-      if (state.ok) {
-        context.setOutput("code", state.value);
-      } else {
-        context.log({
-          type: "error",
-          message: state.value,
-        });
-      }
+    if (!completion) {
+      return context.fail("MINOR", "No completion");
+    }
+
+    const state = await generationCompletionContentToCode(completion);
+    if (state.ok) {
+      return { code: state.value };
+    } else {
+      return context.fail("MINOR", state.value);
     }
   }
 );
diff --git a/packages/ai/src/steps/matchStyleGuideStep.ts b/packages/ai/src/steps/matchStyleGuideStep.ts
index 0be376e6c8..d6f515c840 100644
--- a/packages/ai/src/steps/matchStyleGuideStep.ts
+++ b/packages/ai/src/steps/matchStyleGuideStep.ts
@@ -67,7 +67,7 @@ export const matchStyleGuideStep = new LLMStepTemplate(
       code: "code",
     },
     outputs: {
-      code: "code",
+      code: "code?",
     },
   },
   async (context, { prompt, code }) => {
@@ -81,7 +81,7 @@ export const matchStyleGuideStep = new LLMStepTemplate(
 
     if (!completion) {
       // failed
-      return;
+      return context.fail("MINOR", "No completion");
     }
 
     // handle adjustment response
@@ -95,7 +95,7 @@ export const matchStyleGuideStep = new LLMStepTemplate(
         type: "info",
         message: "LLM determined no adjustment is needed",
       });
-      return;
+      return { code: undefined };
     }
 
     if (
@@ -109,19 +109,17 @@ export const matchStyleGuideStep = new LLMStepTemplate(
           message: "FAIL: " + diffResponse.value,
         });
         // try again
-        context.setOutput("code", code);
-        return;
+        return { code };
       }
 
       const adjustedCode = await codeStringToCode(diffResponse.value);
-      context.setOutput("code", adjustedCode);
-      return;
+      return { code: adjustedCode };
     } else {
       context.log({
         type: "info",
         message: "No adjustments provided, considering process complete",
       });
-      return;
+      return { code: undefined };
     }
   }
 );
diff --git a/packages/ai/src/steps/runAndFormatCodeStep.ts b/packages/ai/src/steps/runAndFormatCodeStep.ts
index deae894a65..0c5f3ec1fd 100644
--- a/packages/ai/src/steps/runAndFormatCodeStep.ts
+++ b/packages/ai/src/steps/runAndFormatCodeStep.ts
@@ -7,8 +7,8 @@ export const runAndFormatCodeStep = new LLMStepTemplate(
     inputs: { source: "source" },
     outputs: { code: "code" },
   },
-  async (context, { source }) => {
+  async (_, { source }) => {
     const code = await codeStringToCode(source.value);
-    context.setOutput("code", code);
+    return { code };
   }
 );
diff --git a/packages/ai/src/workflows/Workflow.ts b/packages/ai/src/workflows/Workflow.ts
index 703c76cd78..db30ac33d3 100644
--- a/packages/ai/src/workflows/Workflow.ts
+++ b/packages/ai/src/workflows/Workflow.ts
@@ -294,9 +294,12 @@ export class Workflow<Shape extends IOShape = IOShape> {
     // Single pass through steps from most recent to oldest
     for (let i = this.steps.length - 1; i >= 0; i--) {
       const step = this.steps[i];
-      const outputs = step.getOutputs();
+      const stepState = step.getState();
+      if (stepState.kind !== "DONE") {
+        continue;
+      }
 
-      for (const output of Object.values(outputs)) {
+      for (const output of Object.values(stepState.outputs)) {
         if (output?.kind === "code") {
           // If we find successful code, return immediately
           if (output.value.type === "success") {
diff --git a/packages/ai/src/workflows/controllers.ts b/packages/ai/src/workflows/controllers.ts
index cf5d4f93bb..ad8cc6d2d3 100644
--- a/packages/ai/src/workflows/controllers.ts
+++ b/packages/ai/src/workflows/controllers.ts
@@ -1,6 +1,6 @@
 import { CodeArtifact, PromptArtifact } from "../Artifact.js";
 import { LLMStepInstance } from "../LLMStepInstance.js";
-import { IOShape } from "../LLMStepTemplate.js";
+import { IOShape, StepState } from "../LLMStepTemplate.js";
 import { adjustToFeedbackStep } from "../steps/adjustToFeedbackStep.js";
 import { fixCodeUntilItRunsStep } from "../steps/fixCodeUntilItRunsStep.js";
 import { generateCodeStep } from "../steps/generateCodeStep.js";
@@ -44,14 +44,21 @@ function handleFailedState<Shape extends IOShape>(
     );
   }
 
+  return undefined;
+}
+
+// assumes that the step is done
+function getOutputs<Shape extends IOShape, WorkflowShape extends IOShape>(
+  step: LLMStepInstance<Shape, WorkflowShape>
+): Extract<StepState<Shape>, { kind: "DONE" }>["outputs"] {
+  const state = step.getState();
   if (state.kind !== "DONE") {
-    return h.fatal(ERROR_MESSAGES.NOT_DONE(step.template.name, state.kind));
+    throw new Error(ERROR_MESSAGES.NOT_DONE(step.template.name, state.kind));
   }
-
-  return undefined;
+  return state.outputs;
 }
 
-export function fixAdjustRetryLoop<Shape extends IOShape>(
+export function getDefaultTransitionRule<Shape extends IOShape>(
   workflow: Workflow<Shape>,
   prompt: PromptArtifact
 ): StepTransitionRule<Shape> {
@@ -59,11 +66,12 @@ export function fixAdjustRetryLoop<Shape extends IOShape>(
     maxNumericSteps: workflow.llmConfig.numericSteps,
     maxStyleGuideSteps: workflow.llmConfig.styleGuideSteps,
   };
+
   return (step, h) => {
-    function getNextIntendedState(
+    const getNextIntendedState = (
       intendedStep: "AdjustToFeedback" | "MatchStyleGuide",
       code: CodeArtifact
-    ): NextStepAction {
+    ): NextStepAction => {
       const nextSteps = [
         // Only check AdjustToFeedback if it's intended and has runs remaining
         intendedStep === "AdjustToFeedback" &&
@@ -76,45 +84,40 @@ export function fixAdjustRetryLoop<Shape extends IOShape>(
           : null,
       ].filter((r) => !!r);
       return nextSteps.length ? nextSteps[0] : h.finish();
-    }
+    };
 
     // process bad states
     const failedState = handleFailedState(step, h);
     if (failedState) return failedState;
 
-    function fixCodeOrAdjustToFeedback(
-      code: CodeArtifact | undefined
-    ): NextStepAction {
-      if (!code) {
-        return h.fatal(ERROR_MESSAGES.NO_CODE);
-      }
+    const fixCodeOrAdjustToFeedback = (code: CodeArtifact): NextStepAction => {
       if (code.value.type !== "success") {
         return h.step(fixCodeUntilItRunsStep, { code });
       }
       return getNextIntendedState("AdjustToFeedback", code);
-    }
+    };
 
     // generateCodeStep
     if (step.instanceOf(generateCodeStep)) {
-      const { code } = step.getOutputs();
+      const { code } = getOutputs(step);
       return fixCodeOrAdjustToFeedback(code);
     }
 
     // runAndFormatCodeStep
     if (step.instanceOf(runAndFormatCodeStep)) {
-      const { code } = step.getOutputs();
+      const { code } = getOutputs(step);
       return fixCodeOrAdjustToFeedback(code);
     }
 
     // fixCodeUntilItRunsStep
     if (step.instanceOf(fixCodeUntilItRunsStep)) {
-      const { code } = step.getOutputs();
+      const { code } = getOutputs(step);
       return fixCodeOrAdjustToFeedback(code);
     }
 
     // adjustToFeedbackStep
     if (step.instanceOf(adjustToFeedbackStep)) {
-      const { code } = step.getOutputs();
+      const { code } = getOutputs(step);
 
       // no code means no need for adjustment, apply style guide
       if (!code) {
@@ -130,7 +133,7 @@ export function fixAdjustRetryLoop<Shape extends IOShape>(
 
     // matchStyleGuideStep
     if (step.instanceOf(matchStyleGuideStep)) {
-      const { code } = step.getOutputs();
+      const { code } = getOutputs(step);
       if (!code) {
         return h.finish();
       }
diff --git a/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts b/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
index cc7d80ca1b..48532b29e7 100644
--- a/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
+++ b/packages/ai/src/workflows/createSquiggleWorkflowTemplate.ts
@@ -1,5 +1,5 @@
 import { generateCodeStep } from "../steps/generateCodeStep.js";
-import { fixAdjustRetryLoop } from "./controllers.js";
+import { getDefaultTransitionRule } from "./controllers.js";
 import { WorkflowTemplate } from "./WorkflowTemplate.js";
 
 /**
@@ -19,5 +19,5 @@ export const createSquiggleWorkflowTemplate = new WorkflowTemplate<{
   getInitialStep: (workflow) =>
     generateCodeStep.prepare({ prompt: workflow.inputs.prompt }),
   getTransitionRule: (workflow) =>
-    fixAdjustRetryLoop(workflow, workflow.inputs.prompt),
+    getDefaultTransitionRule(workflow, workflow.inputs.prompt),
 });
diff --git a/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts b/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
index 7648085678..b768e9cccd 100644
--- a/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
+++ b/packages/ai/src/workflows/fixSquiggleWorkflowTemplate.ts
@@ -1,6 +1,6 @@
 import { PromptArtifact } from "../Artifact.js";
 import { runAndFormatCodeStep } from "../steps/runAndFormatCodeStep.js";
-import { fixAdjustRetryLoop } from "./controllers.js";
+import { getDefaultTransitionRule } from "./controllers.js";
 import { WorkflowTemplate } from "./WorkflowTemplate.js";
 
 export const fixSquiggleWorkflowTemplate = new WorkflowTemplate<{
@@ -15,6 +15,6 @@ export const fixSquiggleWorkflowTemplate = new WorkflowTemplate<{
   getTransitionRule: (workflow) => {
     // TODO - cache the prompt artifact once? maybe even as a global variable
     // (but it's better to just refactor steps to make the prompt optional, somehow)
-    return fixAdjustRetryLoop(workflow, new PromptArtifact(""));
+    return getDefaultTransitionRule(workflow, new PromptArtifact(""));
   },
 });
diff --git a/packages/ai/src/workflows/streaming.ts b/packages/ai/src/workflows/streaming.ts
index a64b0424ce..2076399a45 100644
--- a/packages/ai/src/workflows/streaming.ts
+++ b/packages/ai/src/workflows/streaming.ts
@@ -45,6 +45,22 @@ export function artifactToClientArtifact(value: Artifact): ClientArtifact {
   }
 }
 
+function getClientOutputs<Shape extends IOShape, WorkflowShape extends IOShape>(
+  step: LLMStepInstance<Shape, WorkflowShape>
+): Record<string, ClientArtifact> {
+  const stepState = step.getState();
+  return stepState.kind === "DONE"
+    ? Object.fromEntries(
+        Object.entries(stepState.outputs)
+          .filter(
+            (pair): pair is [string, NonNullable<(typeof pair)[1]>] =>
+              pair[1] !== undefined
+          )
+          .map(([key, value]) => [key, artifactToClientArtifact(value)])
+      )
+    : {};
+}
+
 export function stepToClientStep(step: LLMStepInstance): ClientStep {
   return {
     id: step.id,
@@ -56,14 +72,7 @@ export function stepToClientStep(step: LLMStepInstance): ClientStep {
         artifactToClientArtifact(value),
       ])
     ),
-    outputs: Object.fromEntries(
-      Object.entries(step.getOutputs())
-        .filter(
-          (pair): pair is [string, NonNullable<(typeof pair)[1]>] =>
-            pair[1] !== undefined
-        )
-        .map(([key, value]) => [key, artifactToClientArtifact(value)])
-    ),
+    outputs: getClientOutputs(step),
     messages: step.getConversationMessages(),
   };
 }
@@ -121,14 +130,7 @@ export function addStreamingListeners<Shape extends IOShape>(
       content: {
         id: event.data.step.id,
         state: event.data.step.getState().kind,
-        outputs: Object.fromEntries(
-          Object.entries(event.data.step.getOutputs())
-            .filter(
-              (pair): pair is [string, NonNullable<(typeof pair)[1]>] =>
-                pair[1] !== undefined
-            )
-            .map(([key, value]) => [key, artifactToClientArtifact(value)])
-        ),
+        outputs: getClientOutputs(event.data.step),
         messages: event.data.step.getConversationMessages(),
       },
     });
diff --git a/packages/hub/src/server/ai/storage.ts b/packages/hub/src/server/ai/storage.ts
index de7e6cb71f..b85050db48 100644
--- a/packages/hub/src/server/ai/storage.ts
+++ b/packages/hub/src/server/ai/storage.ts
@@ -14,8 +14,14 @@ export function decodeDbWorkflowToClientWorkflow(
     switch (row.format) {
       case "V1_0":
         return decodeV1_0JsonToClientWorkflow(row.workflow);
-      case "V2_0":
-        return decodeV2_0JsonToClientWorkflow(row.workflow);
+      case "V2_0": {
+        const clientWorkflow = decodeV2_0JsonToClientWorkflow(row.workflow);
+        // serialized workflow doesn't include full logs, but we store them in the database
+        if (clientWorkflow.status === "finished") {
+          clientWorkflow.result.logSummary = row.markdown;
+        }
+        return clientWorkflow;
+      }
       default:
         throw new Error(
           `Unknown workflow format: ${row.format satisfies never}`

From 202104b46e21db5c04518409ddf5f1f4fd539630 Mon Sep 17 00:00:00 2001
From: Vyacheslav Matyukhin <me@berekuk.ru>
Date: Tue, 19 Nov 2024 15:47:25 -0300
Subject: [PATCH 5/7] docs

---
 packages/ai/src/LLMStepInstance.ts            |  2 ++
 packages/ai/src/LLMStepTemplate.ts            | 17 +++++++++++++++++
 packages/hub/src/server/ai/analytics/index.ts | 11 +++++++++--
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/packages/ai/src/LLMStepInstance.ts b/packages/ai/src/LLMStepInstance.ts
index fda0aae38b..d4162b9995 100644
--- a/packages/ai/src/LLMStepInstance.ts
+++ b/packages/ai/src/LLMStepInstance.ts
@@ -140,6 +140,8 @@ export class LLMStepInstance<
       const result = await this.template.execute(executeContext, this.inputs);
 
       const outputs = {} as Outputs<Shape>; // will be filled in below
+
+      // This code is unfortunately full of `as any`, caused by `Object.keys` limitations in TypeScript.
       Object.keys(result).forEach((key: keyof typeof result) => {
         const value = result[key];
         if (value instanceof BaseArtifact) {
diff --git a/packages/ai/src/LLMStepTemplate.ts b/packages/ai/src/LLMStepTemplate.ts
index 60df22de26..c8754c6667 100644
--- a/packages/ai/src/LLMStepTemplate.ts
+++ b/packages/ai/src/LLMStepTemplate.ts
@@ -93,6 +93,23 @@ export class LLMStepTemplate<const Shape extends IOShape = IOShape> {
   constructor(
     public readonly name: string,
     public readonly shape: Shape,
+    /**
+     * This function is the main API for implementing steps.
+     *
+     * It takes the context and inputs, and returns the outputs.
+     *
+     * The returned value must match the template's shape. For example, if the
+     * shape has `{ outputs: { result: 'code' }}`, then the result must be an
+     * object with a `result` property that's either a code artifact, or a value
+     * of such artifact.
+     *
+     * It's possible to use optional outputs, but only if the shape has `?` in
+     * the output kind. In this case, you must return `{ result: undefined }`.
+     *
+     * See the individual step implementations in `./steps/*.ts` for examples.
+     *
+     * If the step has failed, you can call `context.fail()` to stop execution.
+     */
     public readonly execute: (
       context: ExecuteContext,
       inputs: Inputs<Shape>
diff --git a/packages/hub/src/server/ai/analytics/index.ts b/packages/hub/src/server/ai/analytics/index.ts
index be6017b9f4..9e592cd4df 100644
--- a/packages/hub/src/server/ai/analytics/index.ts
+++ b/packages/hub/src/server/ai/analytics/index.ts
@@ -68,7 +68,10 @@ export async function getTypeStats() {
             (typeStats[stepName]["FAILED"] ?? 0) + 1;
           continue;
         }
-        const code = step.getOutputs()["code"];
+        if (state.kind !== "DONE") {
+          continue;
+        }
+        const code = state.outputs["code"];
         if (code && code instanceof CodeArtifact) {
           typeStats[stepName][code.value.type] =
             (typeStats[stepName][code.value.type] ?? 0) + 1;
@@ -92,7 +95,11 @@ export async function getCodeErrors() {
   const errors: StepError[] = [];
   for (const workflow of getModernWorkflows(rows)) {
     for (const step of workflow.getSteps()) {
-      const code = step.getOutputs()["code"];
+      const state = step.getState();
+      if (state.kind !== "DONE") {
+        continue;
+      }
+      const code = state.outputs["code"];
       if (
         code &&
         code instanceof CodeArtifact &&

From fd85dac4a9738e7b1615ed4558b6b7aa2637b9a8 Mon Sep 17 00:00:00 2001
From: Vyacheslav Matyukhin <me@berekuk.ru>
Date: Tue, 19 Nov 2024 15:50:22 -0300
Subject: [PATCH 6/7] clean up deserialization code

---
 packages/ai/src/LLMStepInstance.ts | 30 ++++++++++++++++++------------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/packages/ai/src/LLMStepInstance.ts b/packages/ai/src/LLMStepInstance.ts
index d4162b9995..c638460f5f 100644
--- a/packages/ai/src/LLMStepInstance.ts
+++ b/packages/ai/src/LLMStepInstance.ts
@@ -369,18 +369,22 @@ export class LLMStepInstance<
         visitor.artifact(inputId),
       ])
     );
-    const state: StepState<any> =
-      serializedState.kind === "DONE"
-        ? {
-            ...serializedState,
-            // TODO - remove outputIds
-            outputs: Object.fromEntries(
-              Object.entries(legacyOutputIds ?? serializedState.outputIds).map(
-                ([name, outputId]) => [name, visitor.artifact(outputId)]
-              )
-            ),
-          }
-        : serializedState;
+
+    let state: StepState<any>;
+
+    if (serializedState.kind === "DONE") {
+      const { outputIds, ...serializedStateWithoutOutputs } = serializedState;
+      state = {
+        ...serializedStateWithoutOutputs,
+        outputs: Object.fromEntries(
+          Object.entries(legacyOutputIds ?? outputIds).map(
+            ([name, outputId]) => [name, visitor.artifact(outputId)]
+          )
+        ),
+      };
+    } else {
+      state = serializedState;
+    }
 
     return {
       ...params,
@@ -436,6 +440,8 @@ export type SerializedStep = Omit<
 > & {
   templateName: string;
   inputIds: Record<string, number>;
+  // Legacy - in the modern format we store outputs on SerializedState, but we still support this field for old workflows.
+  // Can be removed if we deserialize and serialize again all workflows in the existing database.
   outputIds?: Record<string, number>;
   state: SerializedState<IOShape>;
 };

From af2d43b826e611a9db456a4ce9e2ec448ae4b222 Mon Sep 17 00:00:00 2001
From: Vyacheslav Matyukhin <me@berekuk.ru>
Date: Tue, 19 Nov 2024 15:54:00 -0300
Subject: [PATCH 7/7] more comments

---
 packages/ai/src/LLMStepTemplate.ts | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/packages/ai/src/LLMStepTemplate.ts b/packages/ai/src/LLMStepTemplate.ts
index c8754c6667..ab8a2e4d76 100644
--- a/packages/ai/src/LLMStepTemplate.ts
+++ b/packages/ai/src/LLMStepTemplate.ts
@@ -48,10 +48,11 @@ export type Outputs<Shape extends IOShape<any, any>> = {
   [K in keyof Shape["outputs"]]: AllowedOutputValue<Shape["outputs"][K]>;
 };
 
+// What the step implementation returns.
 type StepExecuteResult<Shape extends IOShape> = {
-  [K in keyof Shape["outputs"]]:
-    | AllowedOutputValue<Shape["outputs"][K]>
-    | NonNullable<AllowedOutputValue<Shape["outputs"][K]>>["value"];
+  [K in keyof Shape["outputs"]]:  // for any output field...
+    | AllowedOutputValue<Shape["outputs"][K]> // ...return either the artifact object...
+    | NonNullable<AllowedOutputValue<Shape["outputs"][K]>>["value"]; // or its value, as a shorthand
 };
 
 // ExecuteContext is the context that's available to the step implementation.