feat: add support for zod schemas for events (#113)

.changeset/fifty-heads-decide.md

@@ -0,0 +1,5 @@
+---
+"@llama-flow/core": patch
+---
+
+Add support for zod schemas for events

packages/core/src/util/zod.ts

@@ -8,12 +8,17 @@ import {
 export const zodEvent = <T, DebugLabel extends string>(
   schema: z.ZodType<T>,
   config?: WorkflowEventConfig<DebugLabel>,
-): WorkflowEvent<T, DebugLabel> => {
+): WorkflowEvent<T, DebugLabel> & { readonly schema: z.ZodType<T> } => {
   const event = workflowEvent<T, DebugLabel>(config);
   const originalWith = event.with;
-  event.with = (data: T) => {
-    schema.parse(data);
-    return originalWith(data);
-  };
-  return event;
+
+  return Object.assign(event, {
+    with: (data: T) => {
+      schema.parse(data);
+      return originalWith(data);
+    },
+    get schema() {
+      return schema;
+    },
+  });
 };

packages/core/tests/util/zod.spec.ts

@@ -0,0 +1,22 @@
+import { describe, expect, test } from "vitest";
+import { z } from "zod";
+import { zodEvent } from "../../src/util/zod";
+
+describe("zodEvent", () => {
+  test("can use zod schema to infer types and validate", () => {
+    const UserSchema = z.object({
+      name: z.string(),
+      age: z.number(),
+    });
+    const event = zodEvent(UserSchema);
+
+    const schema = event.schema;
+    expect(schema).toBe(UserSchema);
+
+    const validData = { name: "John", age: 30 };
+    const invalidData = { name: "John", age: "30" };
+
+    expect(() => event.with(validData)).not.toThrow();
+    expect(() => event.with(invalidData as any)).toThrow();
+  });
+});