feat(llamaindex): test (#82)

.changeset/lazy-stingrays-cheer.md

@@ -0,0 +1,5 @@
+---
+"@llama-flow/llamaindex": patch
+---
+
+fix: llamaindex workflow edge cases

packages/core/src/stream/consumer.ts

@@ -1,5 +1,3 @@
-import type { WorkflowEventData } from "@llama-flow/core";
-
 /**
  * A no-op function that consumes a stream of events and does nothing with them.
  *
@@ -7,10 +5,10 @@ import type { WorkflowEventData } from "@llama-flow/core";
  * or `getContext()`, it's infinite and will never finish
  */
 export const nothing = async (
-  stream: ReadableStream<WorkflowEventData<any>>,
+  stream: ReadableStream<unknown>,
 ): Promise<void> => {
   await stream.pipeTo(
-    new WritableStream<WorkflowEventData<any>>({
+    new WritableStream<unknown>({
       write: () => {
         // no-op
       },
@@ -24,12 +22,10 @@ export const nothing = async (
  * Do not collect the raw stream from `workflow.createContext()`
  * or getContext()`, it's infinite and will never finish.
  */
-export const collect = async (
-  stream: ReadableStream<WorkflowEventData<any>>,
-): Promise<WorkflowEventData<any>[]> => {
-  const events: WorkflowEventData<any>[] = [];
+export const collect = async <T>(stream: ReadableStream<T>): Promise<T[]> => {
+  const events: T[] = [];
   await stream.pipeTo(
-    new WritableStream<WorkflowEventData<any>>({
+    new WritableStream<T>({
       write: (event) => {
         events.push(event);
       },

packages/llamaindex/src/index.ts

@@ -5,7 +5,6 @@ import {
   workflowEvent,
   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";
 
@@ -83,7 +82,7 @@ export class StopEvent<T = string> extends WorkflowEvent<T> {
 export class Workflow<ContextData, Start, Stop> {
   #workflow = withStore((data: ContextData) => data, createWorkflow());
 
-  addStep<AcceptEvents extends (typeof WorkflowEvent<any>)[]>(
+  addStep<const AcceptEvents extends (typeof WorkflowEvent<any>)[]>(
     parameters: {
       inputs: AcceptEvents;
     },
@@ -126,32 +125,76 @@ export class Workflow<ContextData, Start, Stop> {
               return contextData;
             },
           },
-          ...(events as Parameters<
-            Handler<AcceptEvents, WorkflowEvent<any> | void>
-          >),
+          ...(events.map((e) => coreEventWeakMap.get(e)!) as any),
         );
         if (result instanceof Promise) {
           return result.then((result) =>
             result instanceof WorkflowEvent
-              ? eventWeakMap.get(result.constructor)!.with(result.data)
+              ? eventDataWeakMap.get(result)!
               : undefined,
           );
         } else {
           return result instanceof WorkflowEvent
-            ? eventWeakMap.get(result.constructor)!.with(result.data)
+            ? eventDataWeakMap.get(result)!
             : undefined;
         }
       },
     );
   }
 
-  async run(start: Start, context?: ContextData): Promise<Stop> {
+  run(
+    start: Start,
+    context?: ContextData,
+  ): Promise<Stop> & AsyncIterable<WorkflowEvent<any>> {
     const { sendEvent, stream } = this.#workflow.createContext(context!);
     const startEvent = new StartEvent(start);
     const coreStartEvent = eventDataWeakMap.get(startEvent)!;
     sendEvent(coreStartEvent);
+    if (!eventWeakMap.has(StopEvent)) {
+      eventWeakMap.set(
+        StopEvent,
+        workflowEvent({
+          debugLabel: StopEvent.name,
+        }),
+      );
+    }
     const stopEvent = eventWeakMap.get(StopEvent)!;
-    const events = await collect(until(stream, stopEvent));
-    return events.at(-1)!.data;
+
+    const result = stream.pipeThrough<WorkflowEvent<any>>(
+      new TransformStream({
+        transform: (event, controller) => {
+          const ev = coreEventWeakMap.get(event)!;
+          controller.enqueue(ev);
+          if (stopEvent.include(event)) {
+            controller.terminate();
+          }
+        },
+      }),
+    );
+    Object.assign(result, {
+      then: async (resolve: any, reject: any) => {
+        try {
+          const events = await collect(result);
+          resolve(events.at(-1)!.data as Stop);
+        } catch (error) {
+          reject(error);
+        }
+      },
+      catch: async (reject: any) => {
+        try {
+          await collect(result);
+        } catch (error) {
+          reject(error);
+        }
+      },
+      finally: async (resolve: any) => {
+        try {
+          await collect(result);
+        } finally {
+          resolve();
+        }
+      },
+    });
+    return result as any;
   }
 }

packages/llamaindex/test/legacy.test.ts

@@ -0,0 +1,430 @@
+import { beforeEach, describe, expect, test, vi, type Mocked } from "vitest";
+import type { StepContext } from "../src";
+import { StartEvent, StopEvent, Workflow, WorkflowEvent } from "../src";
+
+class JokeEvent extends WorkflowEvent<{ joke: string }> {}
+
+class AnalysisEvent extends WorkflowEvent<{ analysis: string }> {}
+
+describe("workflow basic", () => {
+  let generateJoke: Mocked<
+    (context: StepContext, ev: StartEvent) => Promise<JokeEvent>
+  >;
+  let critiqueJoke: Mocked<
+    (context: StepContext, ev: JokeEvent) => Promise<StopEvent<string>>
+  >;
+
+  beforeEach(() => {
+    generateJoke = vi.fn(async (_context, _: StartEvent) => {
+      return new JokeEvent({ joke: "a joke" });
+    });
+
+    critiqueJoke = vi.fn(async (_context, _: JokeEvent) => {
+      return new StopEvent("stop");
+    });
+  });
+
+  test("workflow basic", async () => {
+    const workflow = new Workflow<
+      {
+        foo: string;
+        bar: number;
+      },
+      string,
+      string
+    >();
+    workflow.addStep(
+      {
+        inputs: [StartEvent],
+      },
+      async ({ data }, start) => {
+        expect(start).toBeInstanceOf(StartEvent);
+        expect(start.data).toBe("start");
+        expect(data.bar).toBe(42);
+        expect(data.foo).toBe("foo");
+        return new StopEvent("stopped");
+      },
+    );
+
+    const result = workflow.run("start", {
+      foo: "foo",
+      bar: 42,
+    });
+    await result;
+  });
+
+  test("run workflow", async () => {
+    const jokeFlow = new Workflow<unknown, string, string>();
+
+    jokeFlow.addStep({ inputs: [StartEvent<string>] }, generateJoke);
+    jokeFlow.addStep({ inputs: [JokeEvent] }, critiqueJoke);
+
+    const result = await jokeFlow.run("pirates");
+
+    expect(generateJoke).toHaveBeenCalledTimes(1);
+    expect(critiqueJoke).toHaveBeenCalledTimes(1);
+    expect(result).toBe("stop");
+  });
+
+  test("run workflow with multiple in-degree", async () => {
+    const jokeFlow = new Workflow<unknown, string, string>();
+
+    jokeFlow.addStep(
+      {
+        inputs: [StartEvent],
+      },
+      async (context, _) => {
+        context.sendEvent(
+          new AnalysisEvent({
+            analysis: "an analysis",
+          }),
+        );
+        return new JokeEvent({
+          joke: "a joke",
+        });
+      },
+    );
+    jokeFlow.addStep(
+      {
+        inputs: [JokeEvent, AnalysisEvent],
+      },
+      async () => {
+        return new StopEvent("The analysis is insightful and helpful.");
+      },
+    );
+
+    const result = await jokeFlow.run("pirates");
+    expect(result).toBe("The analysis is insightful and helpful.");
+  });
+
+  test("run workflow with object-based StartEvent and StopEvent", async () => {
+    const objectFlow = new Workflow<
+      unknown,
+      Person,
+      {
+        result: {
+          greeting: string;
+        };
+      }
+    >();
+
+    type Person = { name: string; age: number };
+
+    const processObject = vi.fn(async (_context, ev: StartEvent<Person>) => {
+      const { name, age } = ev.data;
+      return new StopEvent({
+        result: { greeting: `Hello ${name}, you are ${age} years old!` },
+      });
+    });
+
+    objectFlow.addStep(
+      {
+        inputs: [StartEvent<Person>],
+      },
+      processObject,
+    );
+
+    const result = await objectFlow.run({ name: "Alice", age: 30 });
+
+    expect(processObject).toHaveBeenCalledTimes(1);
+    expect(result.result).toEqual({
+      greeting: "Hello Alice, you are 30 years old!",
+    });
+  });
+
+  test("workflow with two concurrent steps", async () => {
+    const concurrentFlow = new Workflow<unknown, string, string>();
+
+    const step1 = vi.fn(async (_context, _ev: StartEvent) => {
+      await new Promise((resolve) => setTimeout(resolve, 200));
+      return new StopEvent("Step 1 completed");
+    });
+
+    const step2 = vi.fn(async (_context, _ev: StartEvent) => {
+      await new Promise((resolve) => setTimeout(resolve, 100));
+      return new StopEvent("Step 2 completed");
+    });
+
+    concurrentFlow.addStep(
+      {
+        inputs: [StartEvent<string>],
+      },
+      step1,
+    );
+    concurrentFlow.addStep(
+      {
+        inputs: [StartEvent<string>],
+      },
+      step2,
+    );
+
+    const startTime = new Date();
+    const result = await concurrentFlow.run("start");
+    const endTime = new Date();
+    const duration = endTime.getTime() - startTime.getTime();
+
+    expect(step1).toHaveBeenCalledTimes(1);
+    expect(step2).toHaveBeenCalledTimes(1);
+    expect(duration).toBeLessThan(200);
+    expect(result).toBe("Step 2 completed");
+  });
+
+  test("sendEvent", async () => {
+    const myWorkflow = new Workflow<unknown, string, string>();
+
+    class QueryEvent extends WorkflowEvent<{ query: string }> {}
+
+    class QueryResultEvent extends WorkflowEvent<{ result: string }> {}
+
+    class PendingEvent extends WorkflowEvent<void> {}
+
+    myWorkflow.addStep(
+      {
+        inputs: [StartEvent],
+      },
+      async (context: StepContext, events) => {
+        context.sendEvent(new QueryEvent({ query: "something" }));
+        return new PendingEvent();
+      },
+    );
+
+    myWorkflow.addStep(
+      {
+        inputs: [QueryEvent],
+      },
+      async (context, event) => {
+        return new QueryResultEvent({ result: "query result" });
+      },
+    );
+
+    myWorkflow.addStep(
+      {
+        inputs: [PendingEvent, QueryResultEvent],
+      },
+      async (context, ev0, ev1) => {
+        return new StopEvent(ev1.data.result);
+      },
+    );
+
+    const result = await myWorkflow.run("start");
+    expect(result).toBe("query result");
+  });
+
+  test("allow output with send event", async () => {
+    const myFlow = new Workflow<unknown, string, string>();
+    myFlow.addStep(
+      {
+        inputs: [StartEvent<string>],
+      },
+      async (context, ev) => {
+        context.sendEvent(new StopEvent(`Hello ${ev.data}!`));
+      },
+    );
+    const result = myFlow.run("world");
+    expect(await result).toBe("Hello world!");
+  });
+});
+
+describe("workflow event loop", () => {
+  test("basic", async () => {
+    const jokeFlow = new Workflow<unknown, string, string>();
+
+    jokeFlow.addStep(
+      {
+        inputs: [StartEvent<string>],
+      },
+      async (_context, ev: StartEvent) => {
+        return new StopEvent(`Hello ${ev.data}!`);
+      },
+    );
+
+    const result = await jokeFlow.run("world");
+    expect(result).toBe("Hello world!");
+  });
+
+  test("branch", async () => {
+    const myFlow = new Workflow<unknown, string, string>();
+
+    class BranchA1Event extends WorkflowEvent<{ payload: string }> {}
+
+    class BranchA2Event extends WorkflowEvent<{ payload: string }> {}
+
+    class BranchB1Event extends WorkflowEvent<{ payload: string }> {}
+
+    class BranchB2Event extends WorkflowEvent<{ payload: string }> {}
+
+    let control = false;
+
+    myFlow.addStep(
+      {
+        inputs: [StartEvent<string>],
+      },
+      async (_context, ev) => {
+        if (control) {
+          return new BranchA1Event({ payload: ev.data });
+        } else {
+          return new BranchB1Event({ payload: ev.data });
+        }
+      },
+    );
+
+    myFlow.addStep(
+      {
+        inputs: [BranchA1Event],
+      },
+      async (_context, ev) => {
+        return new BranchA2Event({ payload: ev.data.payload });
+      },
+    );
+
+    myFlow.addStep(
+      {
+        inputs: [BranchB1Event],
+      },
+      async (_context, ev) => {
+        return new BranchB2Event({ payload: ev.data.payload });
+      },
+    );
+
+    myFlow.addStep(
+      {
+        inputs: [BranchA2Event],
+      },
+      async (_context, ev) => {
+        return new StopEvent(`Branch A2: ${ev.data.payload}`);
+      },
+    );
+
+    myFlow.addStep(
+      {
+        inputs: [BranchB2Event],
+      },
+      async (_context, ev) => {
+        return new StopEvent(`Branch B2: ${ev.data.payload}`);
+      },
+    );
+
+    {
+      const result = await myFlow.run("world");
+      expect(result).toMatch(/Branch B2: world/);
+    }
+
+    control = true;
+
+    {
+      const result = await myFlow.run("world");
+      expect(result).toMatch(/Branch A2: world/);
+    }
+
+    {
+      const context = myFlow.run("world");
+      for await (const event of context) {
+        if (event instanceof BranchA2Event) {
+          expect(event.data.payload).toBe("world");
+        }
+        if (event instanceof StopEvent) {
+          expect(event.data).toMatch(/Branch A2: world/);
+        }
+      }
+    }
+  });
+
+  test("one event have multiple outputs", async () => {
+    const myFlow = new Workflow<unknown, string, string>();
+
+    class AEvent extends WorkflowEvent<{ payload: string }> {}
+
+    class BEvent extends WorkflowEvent<{ payload: string }> {}
+
+    class CEvent extends WorkflowEvent<{ payload: string }> {}
+
+    class DEvent extends WorkflowEvent<{ payload: string }> {}
+
+    myFlow.addStep(
+      {
+        inputs: [StartEvent<string>],
+      },
+      async (_context, ev) => {
+        return new StopEvent("STOP");
+      },
+    );
+
+    const fn = vi.fn(async (_context, ev: StartEvent) => {
+      return new AEvent({ payload: ev.data });
+    });
+
+    myFlow.addStep(
+      {
+        inputs: [StartEvent<string>],
+      },
+      fn,
+    );
+    myFlow.addStep(
+      {
+        inputs: [AEvent],
+      },
+      async (_context, ev: AEvent) => {
+        return new BEvent({ payload: ev.data.payload });
+      },
+    );
+    myFlow.addStep(
+      {
+        inputs: [AEvent],
+      },
+      async (_context, ev: AEvent) => {
+        return new CEvent({ payload: ev.data.payload });
+      },
+    );
+    myFlow.addStep(
+      {
+        inputs: [BEvent],
+      },
+      async (_context, ev: BEvent) => {
+        return new DEvent({ payload: ev.data.payload });
+      },
+    );
+    myFlow.addStep(
+      {
+        inputs: [CEvent],
+      },
+      async (_context, ev: CEvent) => {
+        return new DEvent({ payload: ev.data.payload });
+      },
+    );
+    myFlow.addStep(
+      {
+        inputs: [DEvent],
+      },
+      async (_context, ev: DEvent) => {
+        return new StopEvent(`Hello ${ev.data.payload}!`);
+      },
+    );
+
+    const result = await myFlow.run("world");
+    expect(result).toBe("STOP");
+    expect(fn).toHaveBeenCalledTimes(1);
+
+    // streaming events will allow to consume event even stop event is reached
+    const stream = myFlow.run("world");
+    for await (const _ of stream) {
+      /* empty */
+    }
+    expect(fn).toHaveBeenCalledTimes(2);
+  });
+
+  test("run with custom context", async () => {
+    type MyContext = { name: string };
+    const myFlow = new Workflow<MyContext, string, string>();
+    myFlow.addStep(
+      {
+        inputs: [StartEvent<string>],
+      },
+      async ({ data }, _: StartEvent) => {
+        return new StopEvent(`Hello ${data.name}!`);
+      },
+    );
+
+    const result = await myFlow.run("world", { name: "Alice" });
+    expect(result).toBe("Hello Alice!");
+  });
+});