refactor: workflow API

.changeset/giant-jars-grab.md

@@ -0,0 +1,8 @@
+---
+"@llamaindex/core": patch
+---
+
+feat: update workflow implementation
+
+- combine await behavior with async iterator, now there's only one queue to handle the `startEvent`.
+- each Context will have its own step map, so the step can be executed in parallel.

packages/core/src/utils/index.ts

@@ -69,6 +69,15 @@ export function stringifyJSONToMessageContent(value: JSONValue): string {
   return JSON.stringify(value, null, 2).replace(/"([^"]*)"/g, "$1");
 }
 
+export function assertExists<T>(
+  value: T | null | undefined,
+  message: string = "Value does not exist",
+): asserts value is T {
+  if (value === null || value === undefined) {
+    throw new Error(message);
+  }
+}
+
 export {
   extractDataUrlComponents,
   extractImage,

packages/core/src/workflow/context.ts

@@ -1,20 +1,81 @@
-import { type EventTypes, type WorkflowEvent } from "./events";
-import { type StepFunction, type Workflow } from "./workflow";
+import { assertExists } from "../utils";
+import {
+  type EventTypes,
+  StartEvent,
+  StopEvent,
+  WorkflowEvent,
+} from "./events";
+
+export type StepFunction<
+  Start = string,
+  T extends WorkflowEvent = WorkflowEvent,
+> = (context: Context<Start>, ev: T) => Promise<WorkflowEvent>;
+
+export type StepMap = Map<
+  StepFunction<any>,
+  { inputs: EventTypes[]; outputs: EventTypes[] | undefined }
+>;
+
+export type ReadonlyStepMap = ReadonlyMap<
+  StepFunction<any>,
+  { inputs: EventTypes[]; outputs: EventTypes[] | undefined }
+>;
+
+export type ContextParams<Start> = {
+  startEvent: StartEvent<Start>;
+  steps: StepMap;
+  timeout: number | null;
+  verbose: boolean;
+};
+
+export class Context<Start = string>
+  implements AsyncIterable<WorkflowEvent, void, void>, Promise<StopEvent>
+{
+  readonly #steps: ReadonlyStepMap;
+  // reverse map of #steps, helper for get the next step
+  readonly #eventMap: WeakMap<typeof WorkflowEvent, StepFunction<Start>>;
+
+  readonly #startEvent: StartEvent<Start>;
+  readonly #queue: WorkflowEvent[] = [];
+  readonly #globals: Map<string, any> = new Map();
 
-export class Context {
-  #workflow: Workflow;
-  #queues: Map<StepFunction, WorkflowEvent[]> = new Map();
   #eventBuffer: Map<EventTypes, WorkflowEvent[]> = new Map();
-  #globals: Map<string, any> = new Map();
-  #streamingQueue: WorkflowEvent[] = [];
-  running: boolean = true;
+
+  #running: boolean = true;
+  #timeout: number | null = null;
   #verbose: boolean = false;
 
-  constructor(params: { workflow: Workflow; verbose?: boolean }) {
-    this.#workflow = params.workflow;
+  #getStepFunction(event: WorkflowEvent): StepFunction<Start> | undefined {
+    return this.#eventMap.get(event.constructor as EventTypes);
+  }
+
+  get running(): boolean {
+    return this.#running;
+  }
+
+  constructor(params: ContextParams<Start>) {
+    this.#steps = params.steps;
+    this.#startEvent = params.startEvent;
+    // init eventMap
+    this.#eventMap = new WeakMap();
+    for (const [step, { inputs }] of this.#steps) {
+      for (const input of inputs) {
+        this.#eventMap.set(input, step);
+      }
+    }
+
+    if (typeof params.timeout === "number") {
+      this.#timeout = params.timeout;
+    }
     this.#verbose = params.verbose ?? false;
+
+    // push start event to the queue
+    const step = this.#getStepFunction(this.#startEvent);
+    assertExists(step, "No step found for start event");
+    this.#pushEvent(this.#startEvent, step);
   }
 
+  // these two API is not type safe
   set(key: string, value: any): void {
     this.#globals.set(key, value);
   }
@@ -25,7 +86,6 @@ export class Context {
     } else if (defaultValue !== undefined) {
       return defaultValue;
     }
-    throw new Error(`Key '${key}' not found in Context`);
   }
 
   collectEvents(
@@ -62,50 +122,140 @@ export class Context {
     return null;
   }
 
-  sendEvent(message: WorkflowEvent, step?: StepFunction): void {
-    const stepName = step?.name ? `step ${step.name}` : "all steps";
+  #pushEvent(event: WorkflowEvent, step: StepFunction<Start>): void {
+    const stepName = step.name ? `step ${step.name}` : "all steps";
     if (this.#verbose) {
-      console.log(`Sending event ${message} to ${stepName}`);
+      console.log(`Sending event ${event} to ${stepName}`);
     }
-    if (step === undefined) {
-      for (const queue of this.#queues.values()) {
-        queue.push(message);
-      }
-    } else {
-      if (!this.#workflow.hasStep(step)) {
-        throw new Error(`Step ${step} does not exist`);
-      }
 
-      if (!this.#queues.has(step)) {
-        this.#queues.set(step, []);
-      }
-      this.#queues.get(step)!.push(message);
+    if (!this.#steps.has(step)) {
+      throw new Error(`Step ${step} does not exist`);
     }
+
+    this.#queue.push(event);
   }
 
-  getNextEvent(step: StepFunction): WorkflowEvent | undefined {
-    const queue = this.#queues.get(step);
-    if (queue && queue.length > 0) {
-      return queue.shift();
+  // make sure it will only be called once
+  #iterator: AsyncGenerator<WorkflowEvent, void, void> | null = null;
+  #signal: AbortSignal | null = null;
+  get #iteratorSingleton(): AsyncGenerator<WorkflowEvent, void, void> {
+    if (this.#iterator === null) {
+      this.#iterator = this.#createStreamEvents();
     }
-    return undefined;
+    return this.#iterator;
   }
 
-  writeEventToStream(event: WorkflowEvent): void {
-    this.#streamingQueue.push(event);
+  [Symbol.asyncIterator](): AsyncGenerator<WorkflowEvent, void, void> {
+    return this.#iteratorSingleton;
   }
 
-  async *streamEvents(): AsyncGenerator<WorkflowEvent, void, void> {
+  async *#createStreamEvents(): AsyncGenerator<WorkflowEvent, void, void> {
     while (true) {
-      const event = this.#streamingQueue.shift();
+      const event = this.#queue.shift();
       if (event) {
         yield event;
+        // event handling
+        const step = this.#getStepFunction(event);
+        assertExists(step, `No step found for event ${event.displayName}`);
+        const nextEvent = await step.call(null, this, event);
+        if (nextEvent instanceof StopEvent) {
+          this.#running = false;
+          yield nextEvent;
+          return;
+        } else {
+          const nextStep = this.#getStepFunction(nextEvent);
+          assertExists(
+            nextStep,
+            `No step found for event ${nextEvent.displayName}`,
+          );
+          this.#pushEvent(nextEvent, nextStep);
+        }
       } else {
-        if (!this.running) {
-          break;
+        if (!this.#running) {
+          return;
         }
         await new Promise((resolve) => setTimeout(resolve, 0));
       }
     }
   }
+
+  then<TResult1, TResult2 = never>(
+    onfulfilled?:
+      | ((value: StopEvent) => TResult1 | PromiseLike<TResult1>)
+      | null
+      | undefined,
+    onrejected?:
+      | ((reason: any) => TResult2 | PromiseLike<TResult2>)
+      | null
+      | undefined,
+  ) {
+    const stopWorkflow = () => {
+      if (this.#running) {
+        this.#running = false;
+      }
+    };
+
+    if (this.#timeout !== null) {
+      const timeout = this.#timeout;
+      this.#signal = AbortSignal.timeout(timeout * 1000);
+      this.#signal.addEventListener(
+        "abort",
+        () => {
+          stopWorkflow();
+        },
+        {
+          once: true,
+        },
+      );
+    }
+
+    return new Promise<StopEvent>(async (resolve, reject) => {
+      this.#signal?.addEventListener("abort", () => {
+        reject(new Error(`Operation timed out after ${this.#timeout} seconds`));
+      });
+      try {
+        for await (const event of this.#iteratorSingleton) {
+          if (event instanceof StartEvent) {
+            if (this.#verbose) {
+              console.log(`Starting workflow with event ${event}`);
+            }
+          }
+          if (event instanceof StopEvent) {
+            resolve(event);
+          }
+        }
+        const nextValue = await this.#iteratorSingleton.next();
+        if (nextValue.done === false) {
+          reject(new Error("Workflow did not complete"));
+        }
+      } catch (err) {
+        // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
+        reject(err);
+      }
+    })
+      .then(onfulfilled)
+      .catch(onrejected);
+  }
+
+  catch<TResult = never>(
+    onrejected?:
+      | ((reason: any) => TResult | PromiseLike<TResult>)
+      | null
+      | undefined,
+  ) {
+    return this.then((v) => v, onrejected);
+  }
+
+  finally(onfinally?: (() => void) | undefined | null) {
+    return this.then(
+      () => {
+        onfinally?.();
+      },
+      () => {
+        onfinally?.();
+      },
+    ) as Promise<any>;
+  }
+
+  [Symbol.toStringTag]: string = "Context";
 }

packages/core/src/workflow/events.ts

@@ -1,12 +1,14 @@
 export class WorkflowEvent<T extends Record<string, any> = any> {
+  displayName: string;
   data: T;
 
   constructor(data: T) {
     this.data = data;
+    this.displayName = this.constructor.name;
   }
 
   toString() {
-    return `${this.constructor.name}(${JSON.stringify(this.data)})`;
+    return this.displayName;
   }
 }
 

packages/core/src/workflow/index.ts

@@ -1,3 +1,8 @@
-export * from "./context";
-export * from "./events";
-export * from "./workflow";
+export { Context } from "./context";
+export {
+  StartEvent,
+  StopEvent,
+  WorkflowEvent,
+  type EventTypes,
+} from "./events";
+export { Workflow } from "./workflow";