@statelyai/agent 0.0.6 → 0.0.8

package/src/adapters/openai.ts

@@ -140,15 +140,20 @@ export function fromEvent<TInput>(
         .map((t) => {
           const name = t.eventType.replace(/\./g, '_');
           functionNameMapping[name] = t.eventType;
+          const eventSchema = eventSchemaMap[t.eventType];
+          const {
+            description,
+            properties: { type, ...properties },
+          } = eventSchema ?? {};
+
           return {
             type: 'function',
             function: {
               name,
-              description:
-                t.description ?? eventSchemaMap[t.eventType]?.description,
+              description: t.description ?? description,
               parameters: {
                 type: 'object',
-                properties: eventSchemaMap[t.eventType]?.properties ?? {},
+                properties: properties ?? {},
               },
             },
           } as const;

package/src/agent.ts

@@ -1,4 +1,4 @@
-import { ActorOptions, AnyStateMachine, createActor } from 'xstate';
+import { AnyStateMachine, createActor } from 'xstate';
 
 export function createAgent<T extends AnyStateMachine>(
   ...args: Parameters<typeof createActor<T>>

package/src/index.ts

@@ -1,3 +1,3 @@
-export { createSchemas } from './schemas';
+export { defineEvents as defineEvents } from './schemas';
 export { createAgent } from './agent';
 export { createOpenAIAdapter } from './adapters/openai';

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;
-  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>>;