@statelyai/agent 0.0.3 → 0.0.5

package/.vscode/launch.json

@@ -1,17 +1,17 @@
 {
-  // Use IntelliSense to learn about possible attributes.
-  // Hover to view descriptions of existing attributes.
   // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
   "version": "0.2.0",
   "configurations": [
     {
       "type": "node",
       "request": "launch",
-      "name": "Launch Program",
-      "skipFiles": ["<node_internals>/**"],
-      "program": "${file}",
-      "preLaunchTask": "tsc: build - tsconfig.json",
-      "outFiles": ["${workspaceFolder}/**/*.js"]
+      "name": "Debug Current Test File",
+      "autoAttachChildProcesses": true,
+      "skipFiles": ["<node_internals>/**", "**/node_modules/**"],
+      "program": "${workspaceRoot}/node_modules/vitest/vitest.mjs",
+      "args": ["run", "${relativeFile}"],
+      "smartStep": true,
+      "console": "integratedTerminal"
     }
   ]
 }

package/CHANGELOG.md

@@ -1,5 +1,45 @@
 # @statelyai/agent
 
+## 0.0.5
+
+### Patch Changes
+
+- [#9](https://github.com/statelyai/agent/pull/9) [`d8e7b67`](https://github.com/statelyai/agent/commit/d8e7b673f6d265f37b2096b25d75310845860271) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Add `adapter.fromTool(…)`, which creates an actor that chooses agent logic based on a input.
+
+  ```ts
+  const actor = adapter.fromTool(() => "Draw me a picture of a donut", {
+    // tools
+    makeIllustration: {
+      description: "Makes an illustration",
+      run: async (input) => {
+        /* ... */
+      },
+      inputSchema: {
+        /* ... */
+      },
+    },
+    getWeather: {
+      description: "Gets the weather",
+      run: async (input) => {
+        /* ... */
+      },
+      inputSchema: {
+        /* ... */
+      },
+    },
+  });
+
+  //...
+  ```
+
+## 0.0.4
+
+### Patch Changes
+
+- [#5](https://github.com/statelyai/agent/pull/5) [`ae473d7`](https://github.com/statelyai/agent/commit/ae473d73399a15ac3199d77d00eb44a0ea5626db) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Simplify API (WIP)
+
+- [#5](https://github.com/statelyai/agent/pull/5) [`687bed8`](https://github.com/statelyai/agent/commit/687bed87f29bd1d13447cc53b5154da0fe6fdcab) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Add `createSchemas`, `createOpenAIAdapter`, and change `createAgent`
+
 ## 0.0.3
 
 ### Patch Changes

package/dist/index.d.ts

@@ -1,9 +1,10 @@
-import OpenAI from 'openai';
-import { Prop, PromiseActorLogic, ObservableActorLogic, AnyEventObject, Values } from 'xstate';
+import * as xstate from 'xstate';
+import { Prop, Values, AnyStateMachine, createActor, PromiseActorLogic, AnyEventObject, ObservableActorLogic } from 'xstate';
 import { JSONSchema7 } from 'json-schema-to-ts/lib/types/definitions';
 import { FromSchema } from 'json-schema-to-ts';
-import { ChatCompletionCreateParamsNonStreaming } from 'openai/resources';
-import { ChatCompletionCreateParamsBase, ChatCompletionCreateParamsStreaming } from 'openai/resources/chat/completions';
+import OpenAI from 'openai';
+import { ChatCompletionCreateParamsBase } from 'openai/resources/chat/completions';
+import { ChatCompletionCreateParamsNonStreaming, ChatCompletionCreateParamsStreaming } from 'openai/resources';
 
 type EventSchemas = {
     [key: string]: {
@@ -35,54 +36,59 @@ type ConvertContextToJSONSchema<T extends ContextSchema> = {
     additionalProperties: false;
 };
 
-/**
- * Creates [promise actor logic](https://stately.ai/docs/promise-actors) that uses the OpenAI API to generate a completion.
- *
- * @param openai The OpenAI instance.
- * @param inputFn A function that maps arbitrary input to OpenAI chat completion input.
- *
- */
-declare function fromChatCompletion<TInput>(openai: OpenAI, agentSettings: CreateAgentOutput<any>, inputFn: (input: TInput) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming): PromiseActorLogic<OpenAI.Chat.Completions.ChatCompletion, TInput>;
-/**
- * Creates [observable actor logic](https://stately.ai/docs/observable-actors) that uses the OpenAI API to generate a completion stream.
- *
- * @param openai The OpenAI instance to use.
- * @param inputFn A function that maps arbitrary input to OpenAI chat completion input.
- */
-declare function fromChatCompletionStream<TInput>(openai: OpenAI, agentSettings: CreateAgentOutput<any>, inputFn: (input: TInput) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming): ObservableActorLogic<OpenAI.Chat.Completions.ChatCompletionChunk, TInput>;
-/**
- * Creates [promise actor logic](https://stately.ai/docs/promise-actors) that passes the next possible transitions as functions to [OpenAI tool calls](https://platform.openai.com/docs/guides/function-calling) and returns an array of potential next events.
- *
- * @param openai The OpenAI instance to use.
- * @param inputFn A function that maps arbitrary input to OpenAI chat completion input.
- */
-declare function fromEventChoice<TInput>(openai: OpenAI, agentSettings: CreateAgentOutput<any>, inputFn: (input: TInput) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming, options?: {
-    /**
-     * Immediately execute sending the event to the parent actor.
-     * @default false
-     */
-    execute?: boolean;
-}): PromiseActorLogic<AnyEventObject[] | undefined, TInput>;
-interface CreateAgentOutput<T extends {
-    model: ChatCompletionCreateParamsBase['model'];
-    context: ContextSchema;
-    events: EventSchemas;
-}> {
-    model: T['model'];
-    schemas: T;
+declare function createSchemas<TContextSchema extends ContextSchema, TEventSchemas extends EventSchemas>({ context, events, }: {
+    context: TContextSchema;
+    events: TEventSchemas;
+}): {
+    context: ConvertContextToJSONSchema<TContextSchema>;
+    events: ConvertToJSONSchemas<TEventSchemas>;
     types: {
-        context: FromSchema<ConvertContextToJSONSchema<T['context']>>;
-        events: FromSchema<Values<ConvertToJSONSchemas<T['events']>>>;
+        context: FromSchema<ConvertContextToJSONSchema<TContextSchema>>;
+        events: FromSchema<Values<ConvertToJSONSchemas<TEventSchemas>>>;
     };
-    fromEvent: <TInput>(inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming) => PromiseActorLogic<FromSchema<Values<ConvertToJSONSchemas<T['events']>>>[] | undefined, TInput>;
-    fromEventChoice: <TInput>(inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming) => PromiseActorLogic<FromSchema<Values<ConvertToJSONSchemas<T['events']>>>[] | undefined, TInput>;
-    fromChatCompletion: <TInput>(inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming) => PromiseActorLogic<OpenAI.Chat.Completions.ChatCompletion, TInput>;
-    fromChatCompletionStream: <TInput>(inputFn: (input: TInput) => string | ChatCompletionCreateParamsStreaming) => ObservableActorLogic<OpenAI.Chat.Completions.ChatCompletionChunk, TInput>;
+};
+
+declare function createAgent<T extends AnyStateMachine>(...args: Parameters<typeof createActor<T>>): xstate.Actor<T>;
+
+interface StatelyAgentAdapter {
+    model: string;
+    fromEvent: <TInput>(inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming, options?: {
+        /**
+         * Immediately execute sending the event to the parent actor.
+         * @default true
+         */
+        execute?: boolean;
+    }) => PromiseActorLogic<AnyEventObject[] | undefined, TInput>;
+    /**
+     * Creates promise actor logic that resolves with a chat completion.
+     */
+    fromChat: <TInput>(inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming) => PromiseActorLogic<OpenAI.Chat.Completions.ChatCompletion, TInput>;
+    /**
+     * Creates observable actor logic that emits a chat completion stream.
+     */
+    fromChatStream: <TInput>(inputFn: (input: TInput) => string | ChatCompletionCreateParamsStreaming) => ObservableActorLogic<OpenAI.Chat.Completions.ChatCompletionChunk, TInput>;
+    fromTool: <TInput>(inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming, tools: {
+        [key: string]: Tool<any, any>;
+    }, options?: {
+        /**
+         * Immediately execute sending the event to the parent actor.
+         * @default true
+         */
+        execute?: boolean;
+    }) => PromiseActorLogic<{
+        result: any;
+        tool: string;
+        toolCall: OpenAI.Chat.Completions.ChatCompletionMessageToolCall;
+    } | undefined, TInput>;
+}
+interface Tool<TInput, TOutput> {
+    description: string;
+    inputSchema: any;
+    run: (input: TInput) => TOutput;
 }
-declare function createAgent<T extends {
+
+declare function createOpenAIAdapter<T extends {
     model: ChatCompletionCreateParamsBase['model'];
-    context: ContextSchema;
-    events: EventSchemas;
-}>(openai: OpenAI, settings: T): CreateAgentOutput<T>;
+}>(openai: OpenAI, settings: T): StatelyAgentAdapter;
 
-export { createAgent, fromChatCompletion, fromChatCompletionStream, fromEventChoice };
+export { createAgent, createOpenAIAdapter, createSchemas };

package/dist/index.js

@@ -21,15 +21,11 @@ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: tru
 var src_exports = {};
 __export(src_exports, {
   createAgent: () => createAgent,
-  fromChatCompletion: () => fromChatCompletion,
-  fromChatCompletionStream: () => fromChatCompletionStream,
-  fromEventChoice: () => fromEventChoice
+  createOpenAIAdapter: () => createOpenAIAdapter,
+  createSchemas: () => createSchemas
 });
 module.exports = __toCommonJS(src_exports);
 
-// src/openai.ts
-var import_xstate = require("xstate");
-
 // src/utils.ts
 function getAllTransitions(state) {
   const nodes = state._nodes;
@@ -55,9 +51,34 @@ function createEventSchemas(eventSchemaMap) {
   return resolvedEventSchemaMap;
 }
 
-// src/openai.ts
+// src/schemas.ts
+function createSchemas({
+  context,
+  events
+}) {
+  return {
+    context: {
+      type: "object",
+      properties: context,
+      additionalProperties: false,
+      required: Object.keys(context)
+    },
+    events: createEventSchemas(events),
+    types: {}
+  };
+}
+
+// src/agent.ts
+var import_xstate = require("xstate");
+function createAgent(...args) {
+  const [machine, options] = args;
+  return (0, import_xstate.createActor)(machine, options);
+}
+
+// src/adapters/openai.ts
+var import_xstate2 = require("xstate");
 function fromChatCompletion(openai, agentSettings, inputFn) {
-  return (0, import_xstate.fromPromise)(
+  return (0, import_xstate2.fromPromise)(
     async ({ input }) => {
       const openAiInput = inputFn(input);
       const params = typeof openAiInput === "string" ? {
@@ -74,8 +95,8 @@ function fromChatCompletion(openai, agentSettings, inputFn) {
     }
   );
 }
-function fromChatCompletionStream(openai, agentSettings, inputFn) {
-  return (0, import_xstate.fromObservable)(
+function fromChatStream(openai, agentSettings, inputFn) {
+  return (0, import_xstate2.fromObservable)(
     ({ input }) => {
       const observers = /* @__PURE__ */ new Set();
       (async () => {
@@ -101,7 +122,7 @@ function fromChatCompletionStream(openai, agentSettings, inputFn) {
       })();
       return {
         subscribe: (...args) => {
-          const observer = (0, import_xstate.toObserver)(...args);
+          const observer = (0, import_xstate2.toObserver)(...args);
           observers.add(observer);
           return {
             unsubscribe: () => {
@@ -113,9 +134,15 @@ function fromChatCompletionStream(openai, agentSettings, inputFn) {
     }
   );
 }
-function fromEventChoice(openai, agentSettings, inputFn, options) {
-  return (0, import_xstate.fromPromise)(
+function fromEvent(openai, agentSettings, inputFn, options) {
+  return (0, import_xstate2.fromPromise)(
     async ({ input, self, system }) => {
+      const parentSnapshot = self._parent?.getSnapshot();
+      if (!parentSnapshot || !(0, import_xstate2.isMachineSnapshot)(parentSnapshot)) {
+        return void 0;
+      }
+      const schemas = parentSnapshot.machine.schemas;
+      const eventSchemaMap = schemas.events ?? {};
       const transitions = getAllTransitions(self._parent.getSnapshot());
       const functionNameMapping = {};
       const tools = transitions.filter((t) => {
@@ -127,10 +154,10 @@ function fromEventChoice(openai, agentSettings, inputFn, options) {
           type: "function",
           function: {
             name,
-            description: t.description ?? agentSettings.schemas.events[t.eventType]?.description,
+            description: t.description ?? eventSchemaMap[t.eventType]?.description,
             parameters: {
               type: "object",
-              properties: agentSettings.schemas.events[t.eventType]?.properties ?? {}
+              properties: eventSchemaMap[t.eventType]?.properties ?? {}
             }
           }
         };
@@ -167,33 +194,66 @@ function fromEventChoice(openai, agentSettings, inputFn, options) {
     }
   );
 }
-function createAgent(openai, settings) {
+function fromTool(openai, agentSettings, tools, inputFn) {
+  return (0, import_xstate2.fromPromise)(async ({ input, self, system }) => {
+    const functionNameMapping = {};
+    const resolvedTools = Object.entries(tools).map(([key, value]) => {
+      return {
+        type: "function",
+        function: {
+          name: key,
+          description: value.description,
+          parameters: value.inputSchema
+        }
+      };
+    });
+    const openAiInput = inputFn(input);
+    const completionParams = typeof openAiInput === "string" ? {
+      model: agentSettings.model,
+      messages: [
+        {
+          role: "user",
+          content: openAiInput
+        }
+      ]
+    } : openAiInput;
+    const completion = await openai.chat.completions.create({
+      ...completionParams,
+      tools: resolvedTools
+    });
+    const toolCalls = completion.choices[0]?.message.tool_calls;
+    if (toolCalls?.length) {
+      const toolCall = toolCalls[0];
+      const tool = tools[toolCall.function.name];
+      const args = JSON.parse(toolCall.function.arguments);
+      if (tool) {
+        const result = await tool.run(args);
+        return {
+          toolCall,
+          tool: toolCall.function.name,
+          result
+        };
+      }
+    }
+    return void 0;
+  });
+}
+function createOpenAIAdapter(openai, settings) {
   const agentSettings = {
     model: settings.model,
-    schemas: {
-      context: {
-        type: "object",
-        properties: settings.context,
-        additionalProperties: false
-      },
-      events: createEventSchemas(settings.events)
-    },
-    types: {},
     fromEvent: (input) => (
-      // @ts-ignore
-      fromEventChoice(openai, agentSettings, input, { execute: true })
+      // @ts-ignore infinitely deep
+      fromEvent(openai, agentSettings, input, { execute: true })
     ),
-    // @ts-ignore infinitely deep
-    fromEventChoice: (input) => fromEventChoice(openai, agentSettings, input),
-    fromChatCompletion: (input) => fromChatCompletion(openai, agentSettings, input),
-    fromChatCompletionStream: (input) => fromChatCompletionStream(openai, agentSettings, input)
+    fromChat: (input) => fromChatCompletion(openai, agentSettings, input),
+    fromChatStream: (input) => fromChatStream(openai, agentSettings, input),
+    fromTool: (input, tools) => fromTool(openai, agentSettings, tools, input)
   };
   return agentSettings;
 }
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
   createAgent,
-  fromChatCompletion,
-  fromChatCompletionStream,
-  fromEventChoice
+  createOpenAIAdapter,
+  createSchemas
 });

package/examples/helpers/helpers.ts

@@ -3,7 +3,7 @@ import { fromPromise } from 'xstate';
 export const getFromTerminal = fromPromise<string, string>(
   async ({ input }) => {
     const topic = await new Promise<string>((res) => {
-      console.log(input);
+      console.log(input + '\n');
       const listener = (data: Buffer) => {
         const result = data.toString().trim();
         process.stdin.off('data', listener);

package/examples/joke.ts

@@ -1,21 +1,13 @@
 import OpenAI from 'openai';
-import {
-  assign,
-  createActor,
-  fromCallback,
-  fromPromise,
-  log,
-  setup,
-} from 'xstate';
-import { createAgent } from '../src';
+import { assign, fromCallback, fromPromise, log, setup } from 'xstate';
+import { createAgent, createOpenAIAdapter, createSchemas } from '../src';
 import { loadingAnimation } from './helpers/loader';
 
 const openai = new OpenAI({
   apiKey: process.env.OPENAI_API_KEY,
 });
 
-const agent = createAgent(openai, {
-  model: 'gpt-3.5-turbo-1106',
+const schemas = createSchemas({
   context: {
     topic: { type: 'string' },
     jokes: {
@@ -23,18 +15,35 @@ const agent = createAgent(openai, {
       items: {
         type: 'string',
       },
-      desire: { type: ['string', 'null'] },
-      lastRating: { type: ['string', 'null'] },
+    },
+    desire: { type: ['string', 'null'] as const },
+    lastRating: { type: ['string', 'null'] as const },
+  },
+  events: {
+    askForTopic: {
+      type: 'object',
+      properties: {
+        topic: {
+          type: 'string',
+        },
+      },
+    },
+    endJokes: {
+      type: 'object',
+      properties: {},
     },
   },
-  events: {},
 });
 
-const getJokeCompletion = agent.fromChatCompletion(
+const adapter = createOpenAIAdapter(openai, {
+  model: 'gpt-3.5-turbo-1106',
+});
+
+const getJokeCompletion = adapter.fromChat(
   (topic: string) => `Tell me a joke about ${topic}.`
 );
 
-const rateJoke = agent.fromChatCompletion(
+const rateJoke = adapter.fromChat(
   (joke: string) => `Rate this joke on a scale of 1 to 10: ${joke}`
 );
 
@@ -52,7 +61,7 @@ const getTopic = fromPromise(async () => {
   return topic;
 });
 
-const decide = agent.fromEvent(
+const decide = adapter.fromEvent(
   (lastRating: string) =>
     `Choose what to do next, given the previous rating of the joke: ${lastRating}`
 );
@@ -96,15 +105,8 @@ const loader = fromCallback(({ input }: { input: string }) => {
 });
 
 const jokeMachine = setup({
-  types: {
-    context: {} as {
-      topic: string;
-      jokes: string[];
-      desire: string | null;
-      lastRating: string | null;
-    },
-    input: {} as { topic: string },
-  },
+  schemas,
+  types: schemas.types,
   actors: {
     getJokeCompletion,
     getTopic,
@@ -146,7 +148,7 @@ const jokeMachine = setup({
                     event.output.choices[0]!.message.content!
                   ),
               }),
-              log((x) => x.context.jokes.at(-1)),
+              log((x) => `\n` + x.context.jokes.at(-1)),
             ],
             target: 'rateJoke',
           },
@@ -168,7 +170,7 @@ const jokeMachine = setup({
                 lastRating: ({ event }) =>
                   event.output.choices[0]!.message.content!,
               }),
-              log(({ context }) => context.lastRating),
+              log(({ context }) => '\n' + context.lastRating),
             ],
             target: 'decide',
           },
@@ -210,5 +212,5 @@ const jokeMachine = setup({
   },
 });
 
-const actor = createActor(jokeMachine);
-actor.start();
+const agent = createAgent(jokeMachine);
+agent.start();

package/examples/ticTacToe.ts

@@ -1,6 +1,6 @@
-import { assign, setup, assertEvent, createActor } from 'xstate';
+import { assign, setup, assertEvent } from 'xstate';
 import OpenAI from 'openai';
-import { createAgent } from '../src/openai';
+import { createOpenAIAdapter, createSchemas, createAgent } from '../src';
 
 const openai = new OpenAI({
   apiKey: process.env.OPENAI_API_KEY,
@@ -8,8 +8,7 @@ const openai = new OpenAI({
 
 type Player = 'x' | 'o';
 
-const agent = createAgent(openai, {
-  model: 'gpt-4-1106-preview',
+const schemas = createSchemas({
   context: {
     board: {
       type: 'array',
@@ -69,16 +68,20 @@ const agent = createAgent(openai, {
   },
 });
 
+const adapter = createOpenAIAdapter(openai, {
+  model: 'gpt-4-1106-preview',
+});
+
 const initialContext = {
   board: Array(9).fill(null) as Array<Player | null>,
   moves: 0,
   player: 'x' as Player,
   gameReport: '',
   events: [],
-} satisfies typeof agent.types.context;
+} satisfies typeof schemas.types.context;
 
-const bot = agent.fromEvent(
-  ({ context }: { context: typeof agent.types.context }) => `
+const bot = adapter.fromEvent(
+  ({ context }: { context: typeof schemas.types.context }) => `
 You are playing a game of tic tac toe. This is the current game state. The 3x3 board is represented by a 9-element array. The first element is the top-left cell, the second element is the top-middle cell, the third element is the top-right cell, the fourth element is the middle-left cell, and so on. The value of each cell is either null, x, or o. The value of null means that the cell is empty. The value of x means that the cell is occupied by an x. The value of o means that the cell is occupied by an o.
 
 ${JSON.stringify(context, null, 2)}
@@ -86,11 +89,11 @@ ${JSON.stringify(context, null, 2)}
 Execute the single best next move to try to win the game. Do not play on an existing cell.`
 );
 
-const gameReporter = agent.fromChatCompletionStream(
+const gameReporter = adapter.fromChatStream(
   ({
     context,
   }: {
-    context: typeof agent.types.context;
+    context: typeof schemas.types.context;
   }) => `Here is the game board:
 
 ${JSON.stringify(context.board, null, 2)}
@@ -124,7 +127,8 @@ function getWinner(board: typeof initialContext.board): Player | null {
 }
 
 export const ticTacToeMachine = setup({
-  types: agent.types,
+  schemas,
+  types: schemas.types,
   actors: {
     bot,
     gameReporter,
@@ -248,14 +252,8 @@ export const ticTacToeMachine = setup({
   },
 });
 
-const actor = createActor(ticTacToeMachine, {
-  inspect: (e) => {
-    if (e.type === '@xstate.event') {
-      console.log(e.event);
-    }
-  },
-});
-actor.subscribe((s) => {
+const agent = createAgent(ticTacToeMachine);
+agent.subscribe((s) => {
   console.log(s.value, s.context);
 });
-actor.start();
+agent.start();

package/examples/weather.ts

@@ -1,6 +1,6 @@
 import OpenAI from 'openai';
-import { createAgent, fromEventChoice } from '../src';
-import { assign, createActor, fromPromise, log, setup } from 'xstate';
+import { createAgent, createOpenAIAdapter, createSchemas } from '../src';
+import { assign, fromPromise, log, setup } from 'xstate';
 import { getFromTerminal } from './helpers/helpers';
 
 async function searchTavily(
@@ -23,6 +23,7 @@ async function searchTavily(
     },
     body: JSON.stringify(body),
   });
+
   const json = await response.json();
   if (!response.ok) {
     throw new Error(
@@ -39,8 +40,7 @@ const openai = new OpenAI({
   apiKey: process.env.OPENAI_API_KEY,
 });
 
-const agent = createAgent(openai, {
-  model: 'gpt-4-1106-preview',
+const schemas = createSchemas({
   context: {
     location: { type: 'string' },
     history: { type: 'array', items: { type: 'string' } },
@@ -64,17 +64,27 @@ const agent = createAgent(openai, {
   },
 });
 
+const adapter = createOpenAIAdapter(openai, {
+  model: 'gpt-4-1106-preview',
+});
+
+const getWeather = fromPromise(async ({ input }: { input: string }) => {
+  const results = await searchTavily(
+    `Get the weather for this location: ${input}`,
+    {
+      maxResults: 5,
+      apiKey: process.env.TAVILY_API_KEY!,
+    }
+  );
+  return results;
+});
+
 const machine = setup({
-  types: agent.types,
+  schemas,
+  types: schemas.types,
   actors: {
-    searchTavily: fromPromise(async ({ input }: { input: string }) => {
-      const results = await searchTavily(input, {
-        maxResults: 5,
-        apiKey: process.env.TAVILY_API_KEY!,
-      });
-      return results;
-    }),
-    decide: agent.fromEvent(
+    getWeather,
+    decide: adapter.fromEvent(
       (input: string) =>
         `Decide what to do based on the given input, which may or may not be a location: ${input}`
     ),
@@ -121,9 +131,8 @@ const machine = setup({
     gettingWeather: {
       entry: log('Getting weather...'),
       invoke: {
-        src: 'searchTavily',
-        input: ({ context }) =>
-          `Get the weather for this location: ${context.location}`,
+        src: 'getWeather',
+        input: ({ context }) => context.location,
         onDone: {
           actions: [
             log(({ event }) => event.output),
@@ -144,4 +153,8 @@ const machine = setup({
   },
 });
 
-createActor(machine).start();
+createAgent(machine, {
+  input: {
+    location: 'New York',
+  },
+}).start();

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@statelyai/agent",
-  "version": "0.0.3",
+  "version": "0.0.5",
   "description": "",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -17,13 +17,14 @@
     "openai": "^4.24.1",
     "ts-node": "^10.9.2",
     "tsup": "^8.0.1",
-    "typescript": "^5.3.3"
+    "typescript": "^5.3.3",
+    "vitest": "^1.2.2"
   },
   "publishConfig": {
     "access": "public"
   },
   "dependencies": {
-    "xstate": "^5.5.1"
+    "xstate": "^5.6.0"
   },
   "packageManager": "pnpm@8.11.0",
   "scripts": {

package/src/adapter.test.ts

@@ -0,0 +1,217 @@
+import { test, expect } from 'vitest';
+import { createOpenAIAdapter, createTool } from './adapters/openai';
+import OpenAI from 'openai';
+import { createActor, toPromise } from 'xstate';
+
+test('fromTool - weather or illustration', async () => {
+  const openAi = new OpenAI({
+    apiKey: process.env.OPENAI_API_KEY,
+  });
+
+  const adapter = createOpenAIAdapter(openAi, {
+    model: 'gpt-3.5-turbo',
+  });
+
+  const toolChoice = adapter.fromTool(() => 'Create an image of a donut', {
+    makeIllustration: {
+      description: 'Make an illustration',
+      run: async () => 'Illustration',
+      inputSchema: {
+        type: 'object',
+        properties: {
+          name: {
+            type: 'string',
+            description: 'The name of the illustration',
+          },
+        },
+        required: ['name'],
+      },
+    },
+    getWeather: {
+      description: 'Get the weather for a location',
+      run: async () => 'Weather',
+      inputSchema: {
+        type: 'object',
+        properties: {
+          location: {
+            type: 'object',
+            properties: {
+              city: {
+                type: 'string',
+                description: 'The name of the city',
+              },
+              state: {
+                type: 'string',
+                description: 'The name of the state',
+              },
+            },
+            required: ['city', 'state'],
+          },
+        },
+        required: ['location'],
+      },
+    },
+  });
+
+  const actor = createActor(toolChoice);
+
+  actor.start();
+
+  const res = await toPromise(actor);
+
+  expect(res?.result).toBe('Illustration');
+});
+
+test('fromTool - GitHub PR description inserter', async () => {
+  const openAi = new OpenAI({
+    apiKey: process.env.OPENAI_API_KEY,
+  });
+
+  const adapter = createOpenAIAdapter(openAi, {
+    model: 'gpt-3.5-turbo-16k-0613',
+  });
+
+  const toolChoice = adapter.fromTool(
+    (input: string) =>
+      `Create a GitHub PR description for the following: ${input}`,
+    {
+      fetchGitHubPR: {
+        description: 'Fetch a GitHub PR',
+        run: async (input: string) => {
+          return {
+            title: 'Title',
+            body: input,
+          };
+        },
+        inputSchema: {
+          type: 'object',
+          properties: {
+            repo: {
+              type: 'string',
+              description: 'The name of the repo',
+            },
+            number: {
+              type: 'number',
+              description: 'The number of the PR',
+            },
+          },
+          required: ['repo', 'number'],
+        },
+      },
+      createPullRequestDescription: {
+        description: 'Create a GitHub PR description',
+        run: () => 'Description',
+        inputSchema: {
+          type: 'object',
+          properties: {
+            title: {
+              type: 'string',
+              description: 'The title of the PR',
+            },
+            body: {
+              type: 'string',
+              description: 'The body of the PR',
+            },
+          },
+          required: ['title', 'body'],
+        },
+      },
+    }
+  );
+
+  const actor = createActor(toolChoice, {
+    input:
+      // 'Get the details from this: https://github.com/microsoft/TypeScript/pull/47198',
+      'Make a summary of this PR: (some code here)',
+  });
+
+  actor.start();
+
+  const res = await toPromise(actor);
+
+  expect(res?.tool).toEqual('createPullRequestDescription');
+  expect(res?.result).toEqual('Description');
+});
+
+test('fromTool - joke creator or rater', async () => {
+  const openAi = new OpenAI({
+    apiKey: process.env.OPENAI_API_KEY,
+  });
+
+  const adapter = createOpenAIAdapter(openAi, {
+    model: 'gpt-4-1106-preview',
+  });
+
+  const rateJoke = createTool({
+    description: 'Rate a joke',
+    inputSchema: {
+      type: 'object',
+      properties: {
+        joke: {
+          type: 'string',
+          description: 'The joke to rate',
+        },
+      },
+    },
+    run: async ({ topic }: { topic: string }) => {
+      return `Here is a joke about ${topic}`;
+    },
+  });
+
+  const createJoke = createTool({
+    description: 'Create a joke',
+    inputSchema: {
+      type: 'object',
+      properties: {
+        category: {
+          type: 'string',
+          description: 'The category of the joke',
+        },
+      },
+      required: ['category'],
+    },
+    run: async () => {
+      return 'Some joke';
+    },
+  });
+
+  const toolChoice = adapter.fromTool(
+    (input: string) => `
+The user provided this input:
+
+<input>
+${input}
+</input>
+
+Determine what to do:
+- If the input is asking for a joke, create a joke,
+- But if the input is providing a joke, then rate the joke.
+    `,
+    {
+      rateJoke,
+      createJoke,
+    }
+  );
+
+  const actor = createActor(toolChoice, {
+    // input: 'Why did the chicken cross the road? To get to the other side!',
+    input: 'Tell me a joke about chickens',
+  });
+
+  actor.start();
+
+  const res = await toPromise(actor);
+
+  expect(res?.tool).toEqual('createJoke');
+  expect(res?.result).toEqual('Some joke');
+
+  const actor2 = createActor(toolChoice, {
+    input:
+      'Check this joke out: Why did the chicken cross the road? To get to the other side!',
+  });
+
+  actor2.start();
+
+  const res2 = await toPromise(actor2);
+  expect(res2?.tool).toEqual('rateJoke');
+});

package/src/{openai.ts → adapters/openai.ts}

@@ -4,26 +4,18 @@ import {
   ObservableActorLogic,
   Observer,
   PromiseActorLogic,
-  Values,
   fromObservable,
   fromPromise,
-  setup,
+  isMachineSnapshot,
   toObserver,
 } from 'xstate';
-import { getAllTransitions } from './utils';
-import {
-  ContextSchema,
-  EventSchemas,
-  ConvertContextToJSONSchema,
-  ConvertToJSONSchemas,
-  createEventSchemas,
-} from './utils';
-import { FromSchema } from 'json-schema-to-ts';
+import { getAllTransitions } from '../utils';
 import { ChatCompletionCreateParamsNonStreaming } from 'openai/resources';
 import {
   ChatCompletionCreateParamsBase,
   ChatCompletionCreateParamsStreaming,
 } from 'openai/resources/chat/completions';
+import { StatelyAgentAdapter, Tool } from '../types';
 
 /**
  * Creates [promise actor logic](https://stately.ai/docs/promise-actors) that uses the OpenAI API to generate a completion.
@@ -34,7 +26,7 @@ import {
  */
 export function fromChatCompletion<TInput>(
   openai: OpenAI,
-  agentSettings: CreateAgentOutput<any>,
+  agentSettings: OpenAIAdapterOutput<any>,
   inputFn: (
     input: TInput
   ) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming
@@ -67,9 +59,9 @@ export function fromChatCompletion<TInput>(
  * @param openai The OpenAI instance to use.
  * @param inputFn A function that maps arbitrary input to OpenAI chat completion input.
  */
-export function fromChatCompletionStream<TInput>(
+export function fromChatStream<TInput>(
   openai: OpenAI,
-  agentSettings: CreateAgentOutput<any>,
+  agentSettings: OpenAIAdapterOutput<any>,
   inputFn: (
     input: TInput
   ) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming
@@ -126,9 +118,9 @@ export function fromChatCompletionStream<TInput>(
  * @param openai The OpenAI instance to use.
  * @param inputFn A function that maps arbitrary input to OpenAI chat completion input.
  */
-export function fromEventChoice<TInput>(
+export function fromEvent<TInput>(
   openai: OpenAI,
-  agentSettings: CreateAgentOutput<any>,
+  agentSettings: OpenAIAdapterOutput<any>,
   inputFn: (
     input: TInput
   ) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming,
@@ -142,6 +134,15 @@ export function fromEventChoice<TInput>(
 ) {
   return fromPromise<AnyEventObject[] | undefined, TInput>(
     async ({ input, self, system }) => {
+      const parentSnapshot = self._parent?.getSnapshot();
+
+      if (!parentSnapshot || !isMachineSnapshot(parentSnapshot)) {
+        return undefined;
+      }
+
+      const schemas = parentSnapshot.machine.schemas as any;
+      const eventSchemaMap = schemas.events ?? {};
+
       const transitions = getAllTransitions(self._parent!.getSnapshot());
       const functionNameMapping: Record<string, string> = {};
       const tools = transitions
@@ -156,12 +157,10 @@ export function fromEventChoice<TInput>(
             function: {
               name,
               description:
-                t.description ??
-                agentSettings.schemas.events[t.eventType]?.description,
+                t.description ?? eventSchemaMap[t.eventType]?.description,
               parameters: {
                 type: 'object',
-                properties:
-                  agentSettings.schemas.events[t.eventType]?.properties ?? {},
+                properties: eventSchemaMap[t.eventType]?.properties ?? {},
               },
             },
           } as const;
@@ -208,35 +207,124 @@ export function fromEventChoice<TInput>(
   );
 }
 
-interface CreateAgentOutput<
+export function createTool<TInput, T>({
+  description,
+  inputSchema,
+  run,
+}: Tool<TInput, T>): Tool<TInput, T> {
+  return {
+    description,
+    inputSchema,
+    run,
+  };
+}
+
+/**
+ * Creates [promise actor logic](https://stately.ai/docs/promise-actors) that passes the next possible transitions as functions to [OpenAI tool calls](https://platform.openai.com/docs/guides/function-calling) and returns an array of potential next events.
+ *
+ * @param openai The OpenAI instance to use.
+ * @param inputFn A function that maps arbitrary input to OpenAI chat completion input.
+ */
+export function fromTool<TInput>(
+  openai: OpenAI,
+  agentSettings: StatelyAgentAdapter,
+  tools: {
+    [key: string]: Tool<any, any>;
+  },
+  inputFn: (
+    input: TInput
+  ) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming
+) {
+  return fromPromise<
+    | {
+        result: any;
+        tool: string;
+        toolCall: OpenAI.Chat.Completions.ChatCompletionMessageToolCall;
+      }
+    | undefined,
+    TInput
+  >(async ({ input, self, system }) => {
+    const functionNameMapping: Record<string, string> = {};
+    const resolvedTools = Object.entries(tools).map(([key, value]) => {
+      return {
+        type: 'function',
+        function: {
+          name: key,
+          description: value.description,
+          parameters: value.inputSchema,
+        },
+      } as const;
+    });
+
+    const openAiInput = inputFn(input);
+    const completionParams: ChatCompletionCreateParamsNonStreaming =
+      typeof openAiInput === 'string'
+        ? {
+            model: agentSettings.model,
+            messages: [
+              {
+                role: 'user',
+                content: openAiInput,
+              },
+            ],
+          }
+        : openAiInput;
+    const completion = await openai.chat.completions.create({
+      ...completionParams,
+      tools: resolvedTools,
+    });
+
+    const toolCalls = completion.choices[0]?.message.tool_calls;
+
+    if (toolCalls?.length) {
+      const toolCall = toolCalls[0]!;
+      const tool = tools[toolCall.function.name];
+      const args = JSON.parse(toolCall.function.arguments);
+
+      if (tool) {
+        const result = await tool.run(args);
+
+        return {
+          toolCall,
+          tool: toolCall.function.name,
+          result,
+        };
+      }
+    }
+
+    return undefined;
+  });
+}
+
+interface OpenAIAdapterOutput<
   T extends {
     model: ChatCompletionCreateParamsBase['model'];
-    context: ContextSchema;
-    events: EventSchemas;
   }
 > {
   model: T['model'];
-  schemas: T;
-  types: {
-    context: FromSchema<ConvertContextToJSONSchema<T['context']>>;
-    events: FromSchema<Values<ConvertToJSONSchemas<T['events']>>>;
-  };
+  /**
+   * Determines which event to send to the parent state machine actor based on the prompt.
+   */
   fromEvent: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming
-  ) => PromiseActorLogic<
-    FromSchema<Values<ConvertToJSONSchemas<T['events']>>>[] | undefined,
-    TInput
-  >;
-  fromEventChoice: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming
-  ) => PromiseActorLogic<
-    FromSchema<Values<ConvertToJSONSchemas<T['events']>>>[] | undefined,
-    TInput
-  >;
-  fromChatCompletion: <TInput>(
+    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming,
+    options?: {
+      /**
+       * Immediately execute sending the event to the parent actor.
+       * @default true
+       */
+      execute?: boolean;
+    }
+  ) => PromiseActorLogic<AnyEventObject[] | undefined, TInput>;
+  /**
+   * Creates promise actor logic that resolves with a chat completion.
+   */
+  fromChat: <TInput>(
     inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming
   ) => PromiseActorLogic<OpenAI.Chat.Completions.ChatCompletion, TInput>;
-  fromChatCompletionStream: <TInput>(
+  /**
+   * Creates observable actor logic that emits a chat completion stream.
+   */
+  fromChatStream: <TInput>(
     inputFn: (input: TInput) => string | ChatCompletionCreateParamsStreaming
   ) => ObservableActorLogic<
     OpenAI.Chat.Completions.ChatCompletionChunk,
@@ -244,34 +332,20 @@ interface CreateAgentOutput<
   >;
 }
 
-export function createAgent<
+export function createOpenAIAdapter<
   T extends {
     model: ChatCompletionCreateParamsBase['model'];
-    context: ContextSchema;
-    events: EventSchemas;
   }
->(openai: OpenAI, settings: T): CreateAgentOutput<T> {
-  const agentSettings: CreateAgentOutput<T> = {
+>(openai: OpenAI, settings: T): StatelyAgentAdapter {
+  const agentSettings: StatelyAgentAdapter = {
     model: settings.model,
-    schemas: {
-      context: {
-        type: 'object',
-        properties: settings.context,
-        additionalProperties: false,
-      },
-      events: createEventSchemas(settings.events),
-    } as any,
-    types: {} as any,
     fromEvent: (input) =>
-      // @ts-ignore
-      fromEventChoice(openai, agentSettings, input, { execute: true }),
-    // @ts-ignore infinitely deep
-    fromEventChoice: (input) => fromEventChoice(openai, agentSettings, input),
-    fromChatCompletion: (input) =>
-      fromChatCompletion(openai, agentSettings, input),
-    fromChatCompletionStream: (input) =>
-      fromChatCompletionStream(openai, agentSettings, input),
+      // @ts-ignore infinitely deep
+      fromEvent(openai, agentSettings, input, { execute: true }) as any,
+    fromChat: (input) => fromChatCompletion(openai, agentSettings, input),
+    fromChatStream: (input) => fromChatStream(openai, agentSettings, input),
+    fromTool: (input, tools) => fromTool(openai, agentSettings, tools, input),
   };
 
-  return agentSettings as any;
+  return agentSettings;
 }

package/src/agent.ts

@@ -0,0 +1,8 @@
+import { ActorOptions, AnyStateMachine, createActor } from 'xstate';
+
+export function createAgent<T extends AnyStateMachine>(
+  ...args: Parameters<typeof createActor<T>>
+) {
+  const [machine, options] = args;
+  return createActor(machine, options);
+}

package/src/index.ts

@@ -1,6 +1,3 @@
-export {
-  fromChatCompletion,
-  fromChatCompletionStream,
-  fromEventChoice,
-  createAgent,
-} from './openai';
+export { createSchemas } from './schemas';
+export { createAgent } from './agent';
+export { createOpenAIAdapter } from './adapters/openai';

package/src/schemas.ts

@@ -0,0 +1,38 @@
+import { Values } from 'xstate';
+import {
+  ContextSchema,
+  EventSchemas,
+  ConvertContextToJSONSchema,
+  ConvertToJSONSchemas,
+  createEventSchemas,
+} from './utils';
+import { FromSchema } from 'json-schema-to-ts';
+
+export function createSchemas<
+  TContextSchema extends ContextSchema,
+  TEventSchemas extends EventSchemas
+>({
+  context,
+  events,
+}: {
+  context: TContextSchema;
+  events: TEventSchemas;
+}): {
+  context: ConvertContextToJSONSchema<TContextSchema>;
+  events: ConvertToJSONSchemas<TEventSchemas>;
+  types: {
+    context: FromSchema<ConvertContextToJSONSchema<TContextSchema>>;
+    events: FromSchema<Values<ConvertToJSONSchemas<TEventSchemas>>>;
+  };
+} {
+  return {
+    context: {
+      type: 'object',
+      properties: context,
+      additionalProperties: false,
+      required: Object.keys(context),
+    },
+    events: createEventSchemas(events),
+    types: {} as any,
+  };
+}

package/src/types.ts

@@ -0,0 +1,69 @@
+import OpenAI from 'openai';
+import {
+  ChatCompletionCreateParamsNonStreaming,
+  ChatCompletionCreateParamsStreaming,
+} from 'openai/resources';
+import {
+  AnyActorLogic,
+  AnyActorRef,
+  AnyEventObject,
+  ObservableActorLogic,
+  PromiseActorLogic,
+} from 'xstate';
+
+export interface StatelyAgentAdapter {
+  model: string;
+  fromEvent: <TInput>(
+    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming,
+    options?: {
+      /**
+       * Immediately execute sending the event to the parent actor.
+       * @default true
+       */
+      execute?: boolean;
+    }
+  ) => PromiseActorLogic<AnyEventObject[] | undefined, TInput>;
+  /**
+   * Creates promise actor logic that resolves with a chat completion.
+   */
+  fromChat: <TInput>(
+    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming
+  ) => PromiseActorLogic<OpenAI.Chat.Completions.ChatCompletion, TInput>;
+  /**
+   * Creates observable actor logic that emits a chat completion stream.
+   */
+  fromChatStream: <TInput>(
+    inputFn: (input: TInput) => string | ChatCompletionCreateParamsStreaming
+  ) => ObservableActorLogic<
+    OpenAI.Chat.Completions.ChatCompletionChunk,
+    TInput
+  >;
+
+  fromTool: <TInput>(
+    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming,
+    tools: {
+      [key: string]: Tool<any, any>;
+    },
+    options?: {
+      /**
+       * Immediately execute sending the event to the parent actor.
+       * @default true
+       */
+      execute?: boolean;
+    }
+  ) => PromiseActorLogic<
+    | {
+        result: any;
+        tool: string;
+        toolCall: OpenAI.Chat.Completions.ChatCompletionMessageToolCall;
+      }
+    | undefined,
+    TInput
+  >;
+}
+
+export interface Tool<TInput, TOutput> {
+  description: string;
+  inputSchema: any;
+  run: (input: TInput) => TOutput;
+}