feat: exactly same API with llamaindex workflow v2 (#80)

.changeset/lazy-stingrays-cheer.md

@@ -0,0 +1,5 @@
+---
+"@llama-flow/llamaindex": patch
+---
+
+feat: exactly same API with llamaindex workflow v2

packages/llamaindex/src/index.ts

@@ -1,73 +1,156 @@
 import {
   createWorkflow,
-  type Handler,
-  type WorkflowEvent,
+  type WorkflowEvent as CoreWorkflowEvent,
+  type WorkflowEventData as CoreWorkflowEventData,
   workflowEvent,
-  type WorkflowEventData,
   getContext,
 } from "@llama-flow/core";
 import { until } from "@llama-flow/core/stream/until";
 import { collect } from "@llama-flow/core/stream/consumer";
 import { withStore } from "@llama-flow/core/middleware/store";
 
-export {
-  workflowEvent,
-  type WorkflowEventData,
-  type WorkflowEvent,
-  type InferWorkflowEventData,
-} from "@llama-flow/core";
+type Handler<
+  AcceptEvents extends (typeof WorkflowEvent<any>)[],
+  Result extends WorkflowEvent<any> | void,
+> = (
+  ...event: {
+    [K in keyof AcceptEvents]: InstanceType<AcceptEvents[K]>;
+  }
+) => Result | Promise<Result>;
 
-export type StepContext<T = unknown> = ReturnType<typeof getContext> & {
+export type StepContext<T = unknown> = {
+  sendEvent: (event: WorkflowEvent<any>) => void;
+  get stream(): ReadableStream<WorkflowEvent<any>>;
   get data(): T;
 };
 
 export type StepHandler<
   ContextData,
-  Inputs extends WorkflowEvent<any>[],
-  Outputs extends WorkflowEventData<any>[],
+  Inputs extends (typeof WorkflowEvent<any>)[],
+  Outputs extends WorkflowEvent<any>[],
 > = (
   context: StepContext<ContextData>,
   ...args: Parameters<Handler<Inputs, Outputs[number]>>
 ) => ReturnType<Handler<Inputs, Outputs[number]>>;
 
-export const startEvent = workflowEvent<any, "llamaindex-start">({
-  debugLabel: "llamaindex-start",
-});
+const eventWeakMap = new WeakMap<Function, CoreWorkflowEvent<any>>();
+const eventDataWeakMap = new WeakMap<
+  WorkflowEvent<any>,
+  CoreWorkflowEventData<any>
+>();
+const coreEventWeakMap = new WeakMap<
+  CoreWorkflowEventData<any>,
+  WorkflowEvent<any>
+>();
 
-export const stopEvent = workflowEvent<any, "llamaindex-stop">({
-  debugLabel: "llamaindex-stop",
-});
+export class WorkflowEvent<Data> {
+  displayName: string;
+  data: Data;
+
+  constructor(data: Data) {
+    if (!eventWeakMap.has(this.constructor)) {
+      eventWeakMap.set(
+        this.constructor,
+        workflowEvent({
+          debugLabel: this.constructor.name,
+        }),
+      );
+    }
+    this.data = data;
+    this.displayName = this.constructor.name;
+    const coreEvent = eventWeakMap.get(this.constructor)!.with(data);
+    eventDataWeakMap.set(this, coreEvent);
+    coreEventWeakMap.set(coreEvent, this);
+  }
+
+  toString() {
+    return this.displayName;
+  }
+}
+
+export class StartEvent<T = string> extends WorkflowEvent<T> {
+  constructor(data: T) {
+    super(data);
+  }
+}
+
+export class StopEvent<T = string> extends WorkflowEvent<T> {
+  constructor(data: T) {
+    super(data);
+  }
+}
 
 export class Workflow<ContextData, Start, Stop> {
   #workflow = withStore((data: ContextData) => data, createWorkflow());
 
-  addStep<AcceptEvents extends WorkflowEvent<any>[]>(
+  addStep<AcceptEvents extends (typeof WorkflowEvent<any>)[]>(
     parameters: {
       inputs: AcceptEvents;
     },
     handler: (
       context: StepContext<ContextData>,
-      ...args: Parameters<Handler<AcceptEvents, WorkflowEventData<any> | void>>
-    ) => ReturnType<Handler<AcceptEvents, WorkflowEventData<any> | void>>,
+      ...args: Parameters<Handler<AcceptEvents, WorkflowEvent<any> | void>>
+    ) => ReturnType<Handler<AcceptEvents, WorkflowEvent<any> | void>>,
   ) {
-    this.#workflow.handle(parameters.inputs, (...events) => {
-      const context = getContext();
-      const contextData = this.#workflow.getStore();
-      return handler(
-        {
-          ...context,
-          get data(): ContextData {
-            return contextData;
+    this.#workflow.handle(
+      parameters.inputs.map((i) => {
+        if (!eventWeakMap.has(i)) {
+          eventWeakMap.set(
+            i,
+            workflowEvent({
+              debugLabel: i.name,
+            }),
+          );
+        }
+        return eventWeakMap.get(i)!;
+      }),
+      (...events) => {
+        const context = getContext();
+        const contextData = this.#workflow.getStore();
+        const result = handler(
+          {
+            sendEvent: (event) => {
+              const coreEvent = eventDataWeakMap.get(event)!;
+              context.sendEvent(coreEvent);
+            },
+            get stream() {
+              return context.stream.pipeThrough<WorkflowEvent<any>>(
+                new TransformStream({
+                  transform: (event, controller) => {
+                    controller.enqueue(coreEventWeakMap.get(event)!);
+                  },
+                }),
+              );
+            },
+            get data(): ContextData {
+              return contextData;
+            },
           },
-        },
-        ...events,
-      );
-    });
+          ...(events as Parameters<
+            Handler<AcceptEvents, WorkflowEvent<any> | void>
+          >),
+        );
+        if (result instanceof Promise) {
+          return result.then((result) =>
+            result instanceof WorkflowEvent
+              ? eventWeakMap.get(result.constructor)!.with(result.data)
+              : undefined,
+          );
+        } else {
+          return result instanceof WorkflowEvent
+            ? eventWeakMap.get(result.constructor)!.with(result.data)
+            : undefined;
+        }
+      },
+    );
   }
 
   async run(start: Start, context?: ContextData): Promise<Stop> {
     const { sendEvent, stream } = this.#workflow.createContext(context!);
-    sendEvent(startEvent.with(start));
+    const startEvent = new StartEvent(start);
+    const coreStartEvent = eventDataWeakMap.get(startEvent)!;
+    sendEvent(coreStartEvent);
+    const stopEvent = eventWeakMap.get(StopEvent)!;
     const events = await collect(until(stream, stopEvent));
     return events.at(-1)!.data;
   }

packages/llamaindex/test/workflow.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "vitest";
-import { startEvent, stopEvent, Workflow, workflowEvent } from "../src";
+import { StartEvent, StopEvent, Workflow, WorkflowEvent } from "../src";
 
 describe("workflow basic", () => {
   test("basic usage", async () => {
@@ -13,13 +13,13 @@ describe("workflow basic", () => {
     >();
     workflow.addStep(
       {
-        inputs: [startEvent],
+        inputs: [StartEvent<string>],
       },
       async ({ data }, start) => {
         expect(start.data).toBe("start");
         expect(data.bar).toBe(42);
         expect(data.foo).toBe("foo");
-        return stopEvent.with("stopped");
+        return new StopEvent("stopped");
       },
     );
 
@@ -42,13 +42,13 @@ describe("workflow basic", () => {
 
     workflow.addStep(
       {
-        inputs: [startEvent],
+        inputs: [StartEvent],
       },
       async ({ data, sendEvent }, start) => {
         expect(start.data).toBe("start");
         expect(data.bar).toBe(42);
         expect(data.foo).toBe("foo");
-        sendEvent(stopEvent.with("stopped"));
+        sendEvent(new StopEvent("stopped"));
       },
     );
 
@@ -69,20 +69,20 @@ describe("workflow basic", () => {
       string
     >();
 
-    const aEvent = workflowEvent<number>();
-    const bEvent = workflowEvent<number>();
+    class aEvent extends WorkflowEvent<number> {}
+    class bEvent extends WorkflowEvent<number> {}
 
     workflow.addStep(
       {
-        inputs: [startEvent],
+        inputs: [StartEvent],
       },
       async ({ data, sendEvent }, start) => {
         expect(start.data).toBe("start");
         expect(data.bar).toBe(42);
         expect(data.foo).toBe("foo");
-        sendEvent(aEvent.with(1));
+        sendEvent(new aEvent(1));
         setTimeout(() => {
-          sendEvent(bEvent.with(2));
+          sendEvent(new bEvent(2));
         }, 100);
       },
     );
@@ -94,7 +94,7 @@ describe("workflow basic", () => {
       async ({ data }, a, b) => {
         expect(a.data).toBe(1);
         expect(b.data).toBe(2);
-        return stopEvent.with("stopped");
+        return new StopEvent("stopped");
       },
     );