@statelyai/agent 0.0.1 → 0.0.3

package/examples/weather.ts

@@ -0,0 +1,147 @@
+import OpenAI from 'openai';
+import { createAgent, fromEventChoice } from '../src';
+import { assign, createActor, fromPromise, log, setup } from 'xstate';
+import { getFromTerminal } from './helpers/helpers';
+
+async function searchTavily(
+  input: string,
+  options: {
+    maxResults?: number;
+    apiKey: string;
+  }
+) {
+  const body: Record<string, unknown> = {
+    query: input,
+    max_results: options.maxResults,
+    api_key: options.apiKey,
+  };
+
+  const response = await fetch('https://api.tavily.com/search', {
+    method: 'POST',
+    headers: {
+      'content-type': 'application/json',
+    },
+    body: JSON.stringify(body),
+  });
+  const json = await response.json();
+  if (!response.ok) {
+    throw new Error(
+      `Request failed with status code ${response.status}: ${json.error}`
+    );
+  }
+  if (!Array.isArray(json.results)) {
+    throw new Error(`Could not parse Tavily results. Please try again.`);
+  }
+  return JSON.stringify(json.results);
+}
+
+const openai = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY,
+});
+
+const agent = createAgent(openai, {
+  model: 'gpt-4-1106-preview',
+  context: {
+    location: { type: 'string' },
+    history: { type: 'array', items: { type: 'string' } },
+    count: { type: 'number' },
+  },
+  events: {
+    getWeather: {
+      description: 'Get the weather for a location',
+      properties: {
+        location: {
+          type: 'string',
+          description: 'The location to get the weather for',
+        },
+      },
+    },
+    doSomethingElse: {
+      description:
+        'Do something else, because the user did not provide a location',
+      properties: {},
+    },
+  },
+});
+
+const machine = setup({
+  types: agent.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(
+      (input: string) =>
+        `Decide what to do based on the given input, which may or may not be a location: ${input}`
+    ),
+    getFromTerminal,
+  },
+}).createMachine({
+  initial: 'getLocation',
+  context: {
+    location: '',
+    count: 0,
+    history: [],
+  },
+  states: {
+    getLocation: {
+      invoke: {
+        src: 'getFromTerminal',
+        input: 'Location?',
+        onDone: {
+          actions: assign({
+            location: ({ event }) => event.output,
+          }),
+          target: 'decide',
+        },
+      },
+      always: {
+        guard: ({ context }) => context.count >= 3,
+        target: 'stopped',
+      },
+    },
+    decide: {
+      entry: log('Deciding...'),
+      invoke: {
+        src: 'decide',
+        input: ({ context }) => context.location,
+      },
+      on: {
+        getWeather: {
+          actions: log(({ event }) => event),
+          target: 'gettingWeather',
+        },
+        doSomethingElse: 'getLocation',
+      },
+    },
+    gettingWeather: {
+      entry: log('Getting weather...'),
+      invoke: {
+        src: 'searchTavily',
+        input: ({ context }) =>
+          `Get the weather for this location: ${context.location}`,
+        onDone: {
+          actions: [
+            log(({ event }) => event.output),
+            assign({
+              count: ({ context }) => context.count + 1,
+            }),
+          ],
+          target: 'getLocation',
+        },
+      },
+    },
+    stopped: {
+      entry: log('You have used up your search quota. Goodbye!'),
+    },
+  },
+  exit: () => {
+    process.exit();
+  },
+});
+
+createActor(machine).start();

package/package.json

@@ -1,26 +1,38 @@
 {
   "name": "@statelyai/agent",
-  "version": "0.0.1",
+  "version": "0.0.3",
   "description": "",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
   "types": "dist/index.d.ts",
-  "scripts": {
-    "build": "tsup src/index.ts --format cjs,esm --dts",
-    "lint": "tsc",
-    "test": "vitest run",
-    "prepublishOnly": "tsup src/index.ts --dts"
-  },
   "keywords": [],
   "author": "",
   "license": "MIT",
   "devDependencies": {
+    "@changesets/changelog-github": "^0.5.0",
+    "@changesets/cli": "^2.27.1",
     "@types/node": "^20.10.6",
+    "dotenv": "^16.3.1",
+    "json-schema-to-ts": "^3.0.0",
+    "openai": "^4.24.1",
+    "ts-node": "^10.9.2",
     "tsup": "^8.0.1",
     "typescript": "^5.3.3"
   },
+  "publishConfig": {
+    "access": "public"
+  },
   "dependencies": {
-    "openai": "^4.24.1",
-    "xstate": "^5.3.1"
+    "xstate": "^5.5.1"
+  },
+  "packageManager": "pnpm@8.11.0",
+  "scripts": {
+    "build": "tsup src/index.ts --format cjs,esm --dts",
+    "lint": "tsc",
+    "test": "vitest run",
+    "example": "ts-node examples/helpers/runner.ts",
+    "changeset": "changeset",
+    "release": "changeset publish",
+    "version": "changeset version"
   }
-}
+}

package/readme.md

@@ -0,0 +1,41 @@
+# Stately Agent (alpha)
+
+🚧 Documentation in progress! Please see [the examples directory](https://github.com/statelyai/agent/tree/main/examples) for working examples.
+
+## Installation
+
+Install `openai`, and `@statelyai/agent`:
+
+```bash
+npm install openai @statelyai/agent
+```
+
+## Usage
+
+Work in progress. For now, see the examples:
+
+- [Joke generator](https://github.com/statelyai/agent/tree/main/examples/joke.ts)
+  - Demonstrates `agent.fromChatCompletion(...)` to generate a joke and provide a joke rating
+  - Demonstrates `agent.fromEvent(...)` to choose whether to keep generating jokes or stop
+- [Tic-tac-toe](https://github.com/statelyai/agent/tree/main/examples/ticTacToe.ts)
+  - Demonstrates `agent.fromEvent(...)` to have an agent play itself in a game of tic-tac-toe with precise events
+  - Demonstrates `agent.fromChatCompletionStream(...)` to produce a game report at the end of the game
+- [Weather](https://github.com/statelyai/agent/tree/main/examples/weather.ts)
+  - Demonstrates using [Tavily](https://tavily.com/) as an external API
+  - Demonstrates `agent.fromEvent(...)` to only use Tavily to get the weather if the user provides a valid location
+
+## Examples
+
+First, clone this repo locally. To run the examples in this repo, create a `.env` file at the root of the repo with the following contents:
+
+```bash
+OPENAI_API_KEY="your-openai-api-key"
+```
+
+Then, install the dependencies (`npm install`) and run the examples:
+
+```bash
+npm run example joke
+# or:
+# npm run example ticTacToe
+```

package/src/index.ts

@@ -2,4 +2,5 @@ export {
   fromChatCompletion,
   fromChatCompletionStream,
   fromEventChoice,
+  createAgent,
 } from './openai';

package/src/openai.ts

@@ -1,12 +1,29 @@
-import OpenAI from 'openai';
+import type OpenAI from 'openai';
 import {
   AnyEventObject,
+  ObservableActorLogic,
   Observer,
+  PromiseActorLogic,
+  Values,
   fromObservable,
   fromPromise,
+  setup,
   toObserver,
 } from 'xstate';
 import { getAllTransitions } from './utils';
+import {
+  ContextSchema,
+  EventSchemas,
+  ConvertContextToJSONSchema,
+  ConvertToJSONSchemas,
+  createEventSchemas,
+} from './utils';
+import { FromSchema } from 'json-schema-to-ts';
+import { ChatCompletionCreateParamsNonStreaming } from 'openai/resources';
+import {
+  ChatCompletionCreateParamsBase,
+  ChatCompletionCreateParamsStreaming,
+} from 'openai/resources/chat/completions';
 
 /**
  * Creates [promise actor logic](https://stately.ai/docs/promise-actors) that uses the OpenAI API to generate a completion.
@@ -17,14 +34,27 @@ import { getAllTransitions } from './utils';
  */
 export function fromChatCompletion<TInput>(
   openai: OpenAI,
+  agentSettings: CreateAgentOutput<any>,
   inputFn: (
     input: TInput
-  ) => OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming
+  ) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming
 ) {
   return fromPromise<OpenAI.Chat.Completions.ChatCompletion, TInput>(
     async ({ input }) => {
       const openAiInput = inputFn(input);
-      const response = await openai.chat.completions.create(openAiInput);
+      const params: ChatCompletionCreateParamsNonStreaming =
+        typeof openAiInput === 'string'
+          ? {
+              model: agentSettings.model,
+              messages: [
+                {
+                  role: 'user',
+                  content: openAiInput,
+                },
+              ],
+            }
+          : openAiInput;
+      const response = await openai.chat.completions.create(params);
 
       return response;
     }
@@ -39,9 +69,10 @@ export function fromChatCompletion<TInput>(
  */
 export function fromChatCompletionStream<TInput>(
   openai: OpenAI,
+  agentSettings: CreateAgentOutput<any>,
   inputFn: (
     input: TInput
-  ) => OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming
+  ) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming
 ) {
   return fromObservable<OpenAI.Chat.Completions.ChatCompletionChunk, TInput>(
     ({ input }) => {
@@ -49,8 +80,20 @@ export function fromChatCompletionStream<TInput>(
 
       (async () => {
         const openAiInput = inputFn(input);
+        const resolvedParams: ChatCompletionCreateParamsBase =
+          typeof openAiInput === 'string'
+            ? {
+                model: agentSettings.model,
+                messages: [
+                  {
+                    role: 'user',
+                    content: openAiInput,
+                  },
+                ],
+              }
+            : openAiInput;
         const stream = await openai.chat.completions.create({
-          ...openAiInput,
+          ...resolvedParams,
           stream: true,
         });
 
@@ -85,12 +128,20 @@ export function fromChatCompletionStream<TInput>(
  */
 export function fromEventChoice<TInput>(
   openai: OpenAI,
+  agentSettings: CreateAgentOutput<any>,
   inputFn: (
     input: TInput
-  ) => OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming
+  ) => string | OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming,
+  options?: {
+    /**
+     * Immediately execute sending the event to the parent actor.
+     * @default false
+     */
+    execute?: boolean;
+  }
 ) {
   return fromPromise<AnyEventObject[] | undefined, TInput>(
-    async ({ input, self }) => {
+    async ({ input, self, system }) => {
       const transitions = getAllTransitions(self._parent!.getSnapshot());
       const functionNameMapping: Record<string, string> = {};
       const tools = transitions
@@ -104,32 +155,123 @@ export function fromEventChoice<TInput>(
             type: 'function',
             function: {
               name,
-              description: t.description,
+              description:
+                t.description ??
+                agentSettings.schemas.events[t.eventType]?.description,
               parameters: {
                 type: 'object',
-                properties: t.meta?.parameters ?? {},
+                properties:
+                  agentSettings.schemas.events[t.eventType]?.properties ?? {},
               },
             },
           } 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({
-        ...openAiInput,
+        ...completionParams,
         tools,
       });
 
       const toolCalls = completion.choices[0]?.message.tool_calls;
 
       if (toolCalls) {
-        return toolCalls.map((tc) => {
+        const events = toolCalls.map((tc) => {
           return {
             type: functionNameMapping[tc.function.name],
             ...JSON.parse(tc.function.arguments),
           };
         });
+
+        if (options?.execute) {
+          events.forEach((event) => {
+            // @ts-ignore
+            system._relay(self, self._parent, event);
+          });
+        }
       }
 
-      return toolCalls ?? undefined;
+      return undefined;
     }
   );
 }
+
+interface CreateAgentOutput<
+  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']>>>;
+  };
+  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
+  >;
+}
+
+export function createAgent<
+  T extends {
+    model: ChatCompletionCreateParamsBase['model'];
+    context: ContextSchema;
+    events: EventSchemas;
+  }
+>(openai: OpenAI, settings: T): CreateAgentOutput<T> {
+  const agentSettings: CreateAgentOutput<T> = {
+    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),
+  };
+
+  return agentSettings as any;
+}

package/src/utils.ts

@@ -1,4 +1,6 @@
-import { AnyMachineSnapshot, AnyStateNode } from 'xstate';
+import { AnyMachineSnapshot, AnyStateNode, Prop, Values } from 'xstate';
+import { FromSchema } from 'json-schema-to-ts';
+import { JSONSchema7 } from 'json-schema-to-ts/lib/types/definitions';
 
 export function getAllTransitions(state: AnyMachineSnapshot) {
   const nodes = state._nodes;
@@ -8,3 +10,59 @@ export function getAllTransitions(state: AnyMachineSnapshot) {
 
   return transitions;
 }
+
+export type EventSchemas = {
+  [key: string]: {
+    description?: string;
+    properties?: {
+      [key: string]: JSONSchema7;
+    };
+  };
+};
+
+export interface ContextSchema {
+  [key: string]: JSONSchema7;
+}
+
+export type ConvertToJSONSchemas<T> = {
+  [K in keyof T]: {
+    properties: { type: { const: K } };
+    type: 'object';
+    required: Array<keyof Prop<T[K], 'properties'> | 'type'>;
+    additionalProperties: false;
+  } & T[K];
+} & {};
+
+export type ConvertContextToJSONSchema<T extends ContextSchema> = {
+  type: 'object';
+  properties: T;
+  readonly required: Array<keyof T & string>;
+  additionalProperties: false;
+};
+
+export function createEventSchemas<T extends EventSchemas>(
+  eventSchemaMap: T
+): ConvertToJSONSchemas<T> {
+  const resolvedEventSchemaMap = {};
+
+  for (const [key, schema] of Object.entries(eventSchemaMap)) {
+    // @ts-ignore
+    resolvedEventSchemaMap[key] = {
+      type: 'object',
+      required: ['type'],
+      properties: {
+        type: {
+          const: key,
+        },
+        ...schema.properties,
+      },
+      additionalProperties: false,
+      ...schema,
+    } as JSONSchema7;
+  }
+
+  return resolvedEventSchemaMap as ConvertToJSONSchemas<T>;
+}
+
+export type InferEventsFromSchemas<T extends ConvertToJSONSchemas<any>> =
+  FromSchema<Values<T>>;

package/dist/index.d.mts

@@ -1,3 +0,0 @@
-declare function helloWorld(): string;
-
-export { helloWorld };

package/dist/index.mjs

@@ -1,7 +0,0 @@
-// src/index.ts
-function helloWorld() {
-  return "Hello World!";
-}
-export {
-  helloWorld
-};