Update

statelyai · May 2, 2024 · b229cb3 · b229cb3
1 parent 2d63f30
commit b229cb3
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 36 deletions.
diff --git a/.changeset/cold-needles-smile.md b/.changeset/cold-needles-smile.md
@@ -5,12 +5,17 @@
 The `createAgent(…)` function now returns **agent actor logic** instead of the state machine actor logic. That means its return value can now be used directly in `actors` for a state machine, and it will perform tool calls and choose the correct event to send back to the machine:
 
 ```ts
-import { createAgent } from '@statelyai/agent';
+import { createAgent } from '../src';
 import { z } from 'zod';
-import { setup, createAcotr } from 'xstate';
+import { setup, createActor } from 'xstate';
+import OpenAI from 'openai';
 
-const agent = createAgent({
-  model: 'gpt-4-1106-preview',
+const openai = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY,
+});
+
+const agent = createAgent(openai, {
+  model: 'gpt-3.5-turbo-16k-0613',
   events: {
     'agent.thought': z.object({
       text: z.string().describe('The text of the thought'),
@@ -26,7 +31,7 @@ const machine = setup({
     thinking: {
       invoke: {
         src: 'agent',
-        input: 'Produce a random thought',
+        input: 'Think about a random topic, and then share that thought.',
       },
       on: {
         'agent.thought': {

diff --git a/examples/simple.ts b/examples/simple.ts
@@ -0,0 +1,42 @@
+import { createAgent } from '../src';
+import { z } from 'zod';
+import { setup, createActor } from 'xstate';
+import OpenAI from 'openai';
+
+const openai = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY,
+});
+
+const agent = createAgent(openai, {
+  model: 'gpt-3.5-turbo-16k-0613',
+  events: {
+    'agent.thought': z.object({
+      text: z.string().describe('The text of the thought'),
+    }),
+  },
+});
+
+const machine = setup({
+  actors: { agent },
+}).createMachine({
+  initial: 'thinking',
+  states: {
+    thinking: {
+      invoke: {
+        src: 'agent',
+        input: 'Think about a random topic, and then share that thought.',
+      },
+      on: {
+        'agent.thought': {
+          actions: ({ event }) => console.log(event.text),
+          target: 'thought',
+        },
+      },
+    },
+    thought: {
+      type: 'final',
+    },
+  },
+});
+
+const actor = createActor(machine).start();
diff --git a/src/agent.ts b/src/agent.ts
@@ -11,6 +11,26 @@ import { ZodEventTypes, EventSchemas } from './schemas';
 import { createZodEventSchemas } from './utils';
 import { TypeOf } from 'zod';
 
+type AgentLogic<TEventSchemas extends ZodEventTypes> = PromiseActorLogic<
+  void,
+  | {
+      goal: string;
+      model?: ChatCompletionCreateParamsBase['model'];
+      /**
+       * Context to include
+       */
+      context?: any;
+    }
+  | string
+> & {
+  eventTypes: Values<{
+    [K in keyof TEventSchemas]: {
+      type: K;
+    } & TypeOf<TEventSchemas[K]>;
+  }>;
+  eventSchemas: EventSchemas<keyof TEventSchemas & string>;
+};
+
 export function createAgent<const TEventSchemas extends ZodEventTypes>(
   openai: OpenAI,
   {
@@ -20,52 +40,32 @@ export function createAgent<const TEventSchemas extends ZodEventTypes>(
     model: ChatCompletionCreateParamsBase['model'];
     events?: TEventSchemas;
   }
-): PromiseActorLogic<
-  void,
-  {
-    goal: string;
-    model?: ChatCompletionCreateParamsBase['model'];
-    /**
-     * Context to include
-     */
-    context?: any;
-  }
-> & {
-  eventTypes: Values<{
-    [K in keyof TEventSchemas]: {
-      type: K;
-    } & TypeOf<TEventSchemas[K]>;
-  }>;
-  eventSchemas: EventSchemas<keyof TEventSchemas & string>;
-} {
+): AgentLogic<TEventSchemas> {
   const eventSchemas = events ? createZodEventSchemas(events) : undefined;
 
-  const logic = fromPromise<
-    void,
-    {
-      goal: string;
-      model?: ChatCompletionCreateParamsBase['model'];
-      context?: any;
-    }
-  >(async ({ input, self }) => {
+  const logic: Omit<
+    AgentLogic<TEventSchemas>,
+    'eventTypes' | 'eventSchemas'
+  > = fromPromise(async ({ input, self }) => {
     const parentRef = self._parent;
     if (!parentRef) {
       return;
     }
+    const resolvedInput = typeof input === 'string' ? { goal: input } : input;
     const state = parentRef.getSnapshot() as AnyMachineSnapshot;
-    const contextToInclude = input.context
-      ? JSON.stringify(input.context, null, 2)
+    const contextToInclude = resolvedInput.context
+      ? JSON.stringify(resolvedInput.context, null, 2)
       : 'No context provided';
 
     const toolEvents = await getToolCalls(
       openai,
       [
         `<context>\n${JSON.stringify(contextToInclude, null, 2)}\n</context>`,
-        input.goal,
+        resolvedInput.goal,
         'Only make a single tool call.',
       ].join('\n\n'),
       state,
-      input.model ?? model,
+      resolvedInput.model ?? model,
       (eventType) => eventType.startsWith('agent.'),
       eventSchemas ?? (state.machine.schemas as any)?.events
     );
@@ -79,5 +79,5 @@ export function createAgent<const TEventSchemas extends ZodEventTypes>(
 
   (logic as any).eventSchemas = eventSchemas;
 
-  return logic as any;
+  return logic as AgentLogic<TEventSchemas>;
 }