feat: allow arbitrary types in workflow's StartEvent and StopEvent (r…

…un-llama#1210)

.changeset/old-vans-melt.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+Allow arbitrary types in workflow's StartEvent and StopEvent

packages/core/src/workflow/events.ts

@@ -14,5 +14,5 @@ export type EventTypes<T extends Record<string, any> = any> = new (
   data: T,
 ) => WorkflowEvent<T>;
 
-export class StartEvent extends WorkflowEvent<{ input: string }> {}
-export class StopEvent extends WorkflowEvent<{ result: string }> {}
+export class StartEvent<T = string> extends WorkflowEvent<{ input: T }> {}
+export class StopEvent<T = string> extends WorkflowEvent<{ result: T }> {}

packages/core/src/workflow/workflow.ts

@@ -131,7 +131,7 @@ export class Workflow {
     }
   }
 
-  async run(event: StartEvent | string): Promise<StopEvent> {
+  async run<T = string>(event: StartEvent<T> | string): Promise<StopEvent> {
     // Validate the workflow before running if #validate is true
     if (this.#validate) {
       this.validate();

packages/core/tests/workflow.test.ts

@@ -140,4 +140,30 @@ describe("Workflow", () => {
     expect(result.data.result).toBe("Report generated");
     expect(collectedEvents).toHaveLength(1);
   });
+
+  test("run workflow with object-based StartEvent and StopEvent", async () => {
+    const objectFlow = new Workflow({ verbose: true });
+
+    type Person = { name: string; age: number };
+
+    const processObject = vi.fn(async (_context, ev: StartEvent<Person>) => {
+      const { name, age } = ev.data.input;
+      return new StopEvent({
+        result: { greeting: `Hello ${name}, you are ${age} years old!` },
+      });
+    });
+
+    objectFlow.addStep(StartEvent<Person>, processObject);
+
+    const result = await objectFlow.run(
+      new StartEvent<Person>({
+        input: { name: "Alice", age: 30 },
+      }),
+    );
+
+    expect(processObject).toHaveBeenCalledTimes(1);
+    expect(result.data.result).toEqual({
+      greeting: "Hello Alice, you are 30 years old!",
+    });
+  });
 });