Merge pull request #3413 from quantified-uncertainty/resumable-workflows

Serializable workflows and other heavy refactorings

packages/ai/README.md

@@ -6,8 +6,6 @@ The example frontend that uses it is implemented in [Squiggle Hub](https://squig
 
 Note that it can take 20s-2min to run a workflow and get a response from the LLM.
 
-The key file is 'src/workflows/SquiggleWorkflow.ts'. This file contains the definition of a common workflow. It can be used with or without streaming to run the workflow with the given parameters; see `ControlledWorkflow` class API for more details.
-
 After runs are complete, the results are saved to the 'logs' folder. These are saved as Markdown files. It's recommended to use the VSCode Markdown Preview plugin or similar to view the results. In these logs, note the "expand" arrows on the left of some items - these can be clicked to expand the item and see more details.
 
 ## Use
@@ -61,3 +59,75 @@ When using `createSquiggle` and `editSquiggle` scripts, you should define the fo
 ANTHROPIC_API_KEY= // needed for using Claude models
 OPENAI_API_KEY= // needed for using OpenAI models
 ```
+
+## Internals
+
+### Concepts
+
+#### WorkflowTemplate
+
+A description of a multi-step **workflow** that would transform its **inputs** into its **outputs** by going through several **steps**.
+
+Each workflow template has a name, and all templates should be registered in `src/workflows/registry.ts`, so that we can deserialize them by name.
+
+Workflows can have inputs, which are **artifacts**.
+
+#### Workflow
+
+An instance of `WorkflowTemplate` that describes a single living workflow.
+
+Workflows incrementally run their steps and inject new steps into themselves based on the outcomes of previous steps.
+
+#### Controller loop
+
+Each `WorkflowTemplate` configures the workflow with a specific "controller loop": one or more event handlers that add new workflow steps based on events (usually `stepFinished` events) that have happened.
+
+Controller loop don't exist as objects; it's just a handle for the part of the codebase that configures the loop.
+
+The configuration happens in `configureControllerLoop` function in `WorkflowTemplate` definitions.
+
+#### Artifacts
+
+Artifacts are objects that are passed between steps. Both workflows and steps have artifacts as inputs or outputs.
+
+Each artifact has a type, which determines its "shape" (prompt, code, etc).
+
+Artifacts have unique identifiers, so that we can detect when one step is using the output of another step without explicitly connecting them.
+
+#### LLMStepTemplate
+
+Step templates describe a behavior of a single step in a workflow.
+
+Similar to `WorkflowTemplate`, each step template has a name, and all step templates should be registered in `src/steps/registry.ts`.
+
+For now, all steps are "LLM" steps. This might change in the future.
+
+#### LLMStepInstance
+
+An instance of a `LLMStepTemplate`. Step instances have specific inputs and outputs, a state ("PENDING", "DONE", or "FAILED"), and a conversation history.
+
+### Serialization formats
+
+Workflows have two different serialization formats:
+
+#### SerializedWorkflow
+
+`SerializedWorkflow` is used for storing workflows in the database. `SerializedWorkflow` can be fully reconstructed into a `Workflow` object.
+
+`SerializedWorkflow` format is based on `@quri/serializer`, which normalizes the data. The format is not optimized to be human-readable: all object references are transformed into IDs.
+
+#### ClientWorkflow
+
+`ClientWorkflow`: used for representing workflows in the frontend.
+
+`Workflow` objects include server-only code, so we can't have them on the frontend directly, and we send `ClientWorkflow` objects to the frontend.
+
+The advantage of this format is that it's simpler, and it can be incrementally updated by streaming messages as the workflow runs.
+
+### Streaming
+
+You can convert `Workflow` to a stream of JSON-encoded messages by using `workflow.runAsStream()`.
+
+Then you can decode the stream into a `ClientWorkflow` by using `decodeWorkflowFromReader`.
+
+See Squiggle Hub code for details on this.

packages/ai/package.json

@@ -24,6 +24,7 @@
     "@quri/prettier-plugin-squiggle": "workspace:*",
     "@quri/squiggle-lang": "workspace:*",
     "@quri/versioned-squiggle-components": "workspace:*",
+    "@quri/serializer": "workspace:*",
     "axios": "^1.7.2",
     "chalk": "^5.3.0",
     "clsx": "^2.1.1",

packages/ai/src/Artifact.ts

@@ -1,5 +1,5 @@
 import { Code } from "./Code.js";
-import { LLMStepInstance } from "./LLMStep.js";
+import { LLMStepInstance } from "./LLMStepInstance.js";
 
 export class BaseArtifact<T extends string, V> {
   public readonly id: string;
@@ -55,7 +55,7 @@ type ArtifactValue<T extends ArtifactKind> = Extract<
 export function makeArtifact<T extends ArtifactKind>(
   kind: T,
   value: ArtifactValue<T>,
-  createdBy: LLMStepInstance<any>
+  createdBy: LLMStepInstance<any, any>
 ): Extract<Artifact, { kind: T }> {
   // sorry for the type casting, TypeScript is not smart enough to infer the type
   switch (kind) {
@@ -78,3 +78,49 @@ export function makeArtifact<T extends ArtifactKind>(
       throw kind satisfies never;
   }
 }
+
+type ArtifactKindToSerialized<K extends Artifact["kind"]> =
+  K extends Artifact["kind"]
+    ? {
+        id: string;
+        kind: K;
+        value: ArtifactValue<K>;
+      }
+    : never;
+
+export type SerializedArtifact = ArtifactKindToSerialized<Artifact["kind"]>;
+
+export function serializeArtifact(artifact: Artifact): SerializedArtifact {
+  switch (artifact.kind) {
+    case "prompt":
+    case "source":
+      return {
+        id: artifact.id,
+        kind: artifact.kind,
+        value: artifact.value,
+      };
+    case "code":
+      // copy-pasted but type-safe
+      return {
+        id: artifact.id,
+        kind: artifact.kind,
+        value: artifact.value,
+      };
+    default:
+      throw artifact satisfies never;
+  }
+}
+
+export function deserializeArtifact(serialized: SerializedArtifact): Artifact {
+  switch (serialized.kind) {
+    case "prompt":
+      // TODO - pass in createdBy somehow
+      return new PromptArtifact(serialized.value, undefined);
+    case "source":
+      return new SourceArtifact(serialized.value, undefined);
+    case "code":
+      return new CodeArtifact(serialized.value, undefined);
+    default:
+      throw serialized satisfies never;
+  }
+}

packages/ai/src/Code.ts

@@ -5,7 +5,6 @@ import {
   result,
   simpleValueFromAny,
   simpleValueToCompactString,
-  SqError,
   SqErrorList,
   SqProject,
 } from "@quri/squiggle-lang";
@@ -27,7 +26,7 @@ export type Code =
       error: string;
       source: string;
     }
-  | { type: "runFailed"; source: string; error: SqError; project: SqProject }
+  | { type: "runFailed"; source: string; error: string }
   | {
       type: "success";
       source: string;
@@ -41,7 +40,7 @@ export function codeErrorString(code: Code): string {
   if (code.type === "formattingFailed") {
     return code.error;
   } else if (code.type === "runFailed") {
-    return code.error.toStringWithDetails();
+    return code.error;
   }
   return "";
 }
@@ -99,8 +98,7 @@ export async function codeStringToCode(code: string): Promise<Code> {
     return {
       type: "runFailed",
       source: runningCode,
-      error: run.value.error.errors[0],
-      project: run.value.project,
+      error: run.value.error.errors[0].toStringWithDetails(),
     };
   }
 

packages/ai/src/LLMClient.ts

@@ -85,12 +85,12 @@ function convertOpenAIToStandardFormat(
   };
 }
 
-export interface LlmMetrics {
+export type LlmMetrics = {
   apiCalls: number;
   inputTokens: number;
   outputTokens: number;
   llmId: LlmId;
-}
+};
 
 export function calculatePriceMultipleCalls(
   metrics: Partial<Record<LlmId, LlmMetrics>>