@statelyai/agent 0.0.7 → 0.0.8

package/src/schemas.ts

@@ -1,40 +1,26 @@
 import { Values } from 'xstate';
-import {
-  ContextSchema,
-  EventSchemas,
-  ConvertToJSONSchemas,
-  createEventSchemas,
-} from './utils';
-import { FromSchema } from 'json-schema-to-ts';
+import { createZodEventSchemas } from './utils';
+import { SomeZodObject, TypeOf } from 'zod';
 
-export function createSchemas<
-  const TContextSchema extends ContextSchema,
-  const TEventSchemas extends EventSchemas
->({
-  context,
-  events,
-}: {
-  /**
-   * The JSON schema for the context object.
-   *
-   * Must be of `{ type: 'object' }`.
-   */
-  context?: TContextSchema;
-  /**
-   * An object mapping event types to each event object's JSON Schema.
-   */
-  events: TEventSchemas;
-}): {
-  context: TContextSchema | undefined;
-  events: ConvertToJSONSchemas<TEventSchemas>;
-  types: {
-    context: FromSchema<TContextSchema>;
-    events: FromSchema<Values<ConvertToJSONSchemas<TEventSchemas>>>;
+export type ZodEventTypes = {
+  // map event types to Zod types
+  [eventType: string]: SomeZodObject;
+};
+
+export function defineEvents<const TEventSchemas extends ZodEventTypes>(
+  events: TEventSchemas
+): {
+  types: Values<{
+    [K in keyof TEventSchemas]: {
+      type: K;
+    } & TypeOf<TEventSchemas[K]>;
+  }>;
+  schemas: {
+    [K in keyof TEventSchemas]: unknown;
   };
 } {
   return {
-    context,
-    events: createEventSchemas(events),
     types: {} as any,
+    schemas: createZodEventSchemas(events),
   };
 }

package/src/utils.ts

@@ -1,6 +1,9 @@
 import { AnyMachineSnapshot, AnyStateNode, Prop, Values } from 'xstate';
 import { FromSchema } from 'json-schema-to-ts';
 import { JSONSchema7 } from 'json-schema-to-ts/lib/types/definitions';
+import zodToJsonSchema, { JsonSchema7Type } from 'zod-to-json-schema';
+import { ZodEventTypes } from './schemas';
+import { z } from 'zod';
 
 export function getAllTransitions(state: AnyMachineSnapshot) {
   const nodes = state._nodes;
@@ -15,7 +18,7 @@ export type EventSchemas = {
   [key: string]: {
     description?: string;
     properties?: {
-      [key: string]: JSONSchema7;
+      [key: string]: JsonSchema7Type;
     };
   };
 };
@@ -55,5 +58,24 @@ export function createEventSchemas<T extends EventSchemas>(
   return resolvedEventSchemaMap as ConvertToJSONSchemas<T>;
 }
 
+export function createZodEventSchemas<T extends ZodEventTypes>(
+  eventSchemaMap: T
+): {
+  [K in keyof T]: unknown;
+} {
+  const resolvedEventSchemaMap = {};
+
+  for (const [eventType, zodType] of Object.entries(eventSchemaMap)) {
+    // @ts-ignore
+    resolvedEventSchemaMap[eventType] = zodToJsonSchema(
+      zodType.extend({
+        type: z.literal(eventType),
+      })
+    );
+  }
+
+  return resolvedEventSchemaMap as ConvertToJSONSchemas<T>;
+}
+
 export type InferEventsFromSchemas<T extends ConvertToJSONSchemas<any>> =
   FromSchema<Values<T>>;