@statelyai/agent 1.0.0-beta.1 → 1.1.2

package/examples/sandbox.ts

@@ -0,0 +1,28 @@
+import { z } from 'zod';
+import { createAgent } from '../src';
+import { openai } from '@ai-sdk/openai';
+import { createMachine } from 'xstate';
+
+const agent = createAgent({
+  model: openai('gpt-4o'),
+  events: {
+    doSomething: z.object({}).describe('Do something'),
+  },
+});
+
+async function main() {
+  const machine = createMachine({
+    on: {
+      doSomething: {},
+    },
+  });
+  const result = await agent.decide({
+    goal: 'Do not do anything',
+    state: { value: {}, context: {} },
+    machine,
+  });
+
+  console.log(result);
+}
+
+main();

package/examples/support.ts

@@ -35,7 +35,7 @@ const agent = createAgent({
 
 const machine = setup({
   types: {
-    events: agent.eventTypes,
+    events: agent.types.events,
     input: {} as string,
     context: {} as {
       customerIssue: string;

package/examples/ticTacToe.ts

@@ -23,25 +23,30 @@ const agent = createAgent({
     }),
     reset: z.object({}).describe('Reset the game to the initial state'),
   },
+  context: {
+    board: z
+      .array(z.union([z.literal(null), z.literal('x'), z.literal('o')]))
+      .describe('The 3x3 board represented as a 9-element array.'),
+    moves: z
+      .number()
+      .min(0)
+      .max(9)
+      .describe('The number of moves made in the game.'),
+    player: z
+      .union([z.literal('x'), z.literal('o')])
+      .describe('The current player (x or o)'),
+    gameReport: z.string(),
+  },
 });
 
 type Player = 'x' | 'o';
 
-interface GameContext {
-  board: (Player | null)[];
-  moves: number;
-  player: Player;
-  gameReport: string;
-  events: string[];
-}
-
 const initialContext = {
   board: Array(9).fill(null) as Array<Player | null>,
   moves: 0,
   player: 'x' as Player,
   gameReport: '',
-  events: [],
-} satisfies GameContext;
+} satisfies typeof agent.types.context;
 
 function getWinner(board: typeof initialContext.board): Player | null {
   const lines = [
@@ -64,8 +69,8 @@ function getWinner(board: typeof initialContext.board): Player | null {
 
 export const ticTacToeMachine = setup({
   types: {
-    context: {} as GameContext,
-    events: agent.eventTypes,
+    context: agent.types.context,
+    events: agent.types.events,
   },
   actors: {
     agent: fromDecision(agent),
@@ -81,16 +86,8 @@ export const ticTacToeMachine = setup({
       },
       moves: ({ context }) => context.moves + 1,
       player: ({ context }) => (context.player === 'x' ? 'o' : 'x'),
-      events: ({ context, event }) => {
-        return [...context.events, JSON.stringify(event)];
-      },
     }),
     resetGame: assign(initialContext),
-    recordEvent: assign({
-      events: ({ context, event }) => {
-        return [...context.events, JSON.stringify(event)];
-      },
-    }),
     printBoard: ({ context }) => {
       // Print the context.board in a 3 x 3 grid format
       let boardString = '';
@@ -172,7 +169,7 @@ export const ticTacToeMachine = setup({
         src: 'gameReporter',
         input: ({ context }) => ({
           context: {
-            events: context.events,
+            events: agent.getObservations().map((o) => o.event),
             board: context.board,
           },
           prompt: 'Provide a short game report analyzing the game.',

package/examples/todo.ts

@@ -36,7 +36,9 @@ const machine = setup({
       todos: Todo[];
       command: string | null;
     },
-    events: {} as typeof agent.eventTypes | { type: 'assist'; command: string },
+    events: {} as
+      | typeof agent.types.events
+      | { type: 'assist'; command: string },
   },
   actors: { agent: fromDecision(agent), getFromTerminal },
 }).createMachine({

package/examples/tutor.ts

@@ -28,7 +28,7 @@ const machine = setup({
     context: {} as {
       conversation: string[];
     },
-    events: agent.eventTypes,
+    events: agent.types.events,
   },
   actors: { agent: fromDecision(agent), getFromTerminal },
 }).createMachine({

package/examples/verify.ts

@@ -35,7 +35,7 @@ const machine = setup({
       answer: string | null;
       validation: string | null;
     },
-    events: agent.eventTypes,
+    events: agent.types.events,
   },
   actors: {
     getFromTerminal,

package/examples/weather.ts

@@ -78,7 +78,7 @@ const machine = setup({
       history: string[];
       count: number;
     },
-    events: agent.eventTypes,
+    events: agent.types.events,
   },
   actors: {
     agent: fromDecision(agent),

package/examples/wiki.ts

@@ -20,7 +20,7 @@ async function main() {
   console.log(response1.text);
 
   const response2 = await agent.generateText({
-    messages: true,
+    messages: (x) => x.select((ctx) => ctx.messages),
     prompt: 'What about the first one?',
   });
 

package/examples/word.ts

@@ -36,7 +36,7 @@ const agent = createAgent({
 const wordGuesserMachine = setup({
   types: {
     context: {} as typeof context,
-    events: agent.eventTypes,
+    events: agent.types.events,
   },
   actors: {
     agent: fromDecision(agent),

package/package.json

@@ -1,21 +1,10 @@
 {
   "name": "@statelyai/agent",
-  "version": "1.0.0-beta.1",
+  "version": "1.1.2",
   "description": "Stateful agents that make decisions based on finite-state machine models",
   "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 --noEmit",
-    "test": "vitest",
-    "test:ci": "vitest --run",
-    "example": "ts-node examples/helpers/runner.ts",
-    "prepublishOnly": "tsup src/index.ts --format cjs,esm --dts",
-    "changeset": "changeset",
-    "release": "changeset publish",
-    "version": "changeset version"
-  },
   "keywords": [
     "ai",
     "state machine",
@@ -27,17 +16,18 @@
   "license": "MIT",
   "devDependencies": {
     "@changesets/changelog-github": "^0.5.0",
-    "@changesets/cli": "^2.27.5",
+    "@changesets/cli": "^2.27.7",
     "@langchain/community": "^0.0.53",
     "@langchain/core": "^0.1.63",
     "@langchain/openai": "^0.0.28",
-    "@types/node": "^20.14.6",
+    "@types/node": "^20.14.13",
+    "@types/object-hash": "^3.0.6",
     "dotenv": "^16.4.5",
     "json-schema-to-ts": "^3.1.0",
     "ts-node": "^10.9.2",
-    "tsup": "^8.1.0",
-    "typescript": "^5.4.5",
-    "vitest": "^1.6.0",
+    "tsup": "^8.2.3",
+    "typescript": "^5.5.4",
+    "vitest": "^2.0.4",
     "wikipedia": "^2.1.2",
     "zod": "^3.23.8"
   },
@@ -45,10 +35,20 @@
     "access": "public"
   },
   "dependencies": {
-    "@ai-sdk/openai": "^0.0.31",
+    "@ai-sdk/openai": "^0.0.40",
     "@xstate/graph": "^2.0.0",
-    "ai": "^3.2.1",
-    "xstate": "^5.13.2"
+    "ai": "^3.2.40",
+    "object-hash": "^3.0.0",
+    "xstate": "^5.16.0"
   },
-  "packageManager": "pnpm@8.11.0"
-}
+  "scripts": {
+    "build": "tsup src/index.ts --format cjs,esm --dts",
+    "lint": "tsc --noEmit",
+    "test": "vitest",
+    "test:ci": "vitest --run",
+    "example": "ts-node examples/helpers/runner.ts",
+    "changeset": "changeset",
+    "release": "changeset publish",
+    "version": "changeset version"
+  }
+}

package/readme.md

@@ -1,8 +1,5 @@
 # Stately Agent
 
-> [!WARNING]  
-> Alpha software! Not ready for production yet.
-
 Stately Agent is a flexible framework for building AI agents using state machines. Stately agents go beyond normal LLM-based AI agents by:
 
 - Using state machines to guide the agent's behavior, powered by [XState](https://stately.ai/docs/xstate)
@@ -10,4 +7,4 @@ Stately Agent is a flexible framework for building AI agents using state machine
 - Enabling custom **planning** abilities for agents to achieve specific goals based on state machine logic, observations, and feedback
 - Wrapping the [Vercel AI SDK](https://sdk.vercel.ai/) to easily support multiple model providers, such as OpenAI, Anthropic, Google, Mistral, Groq, Perplexity, and more
 
-_Documentation coming soon_
+**Read the documentation: [stately.ai/docs/agents](https://stately.ai/docs/agents)**

package/src/agent.test.ts

@@ -1,6 +1,13 @@
-import { test, expect } from 'vitest';
-import { createAgent, type AIAdapter } from './';
+import { test, expect, vi } from 'vitest';
+import {
+  AgentGenerateTextResult,
+  AgentMessage,
+  createAgent,
+  type AIAdapter,
+} from './';
 import { createActor, createMachine } from 'xstate';
+import { GenerateTextResult } from 'ai';
+import { z } from 'zod';
 
 test('an agent has the expected interface', () => {
   const agent = createAgent({
@@ -13,11 +20,16 @@ test('an agent has the expected interface', () => {
   expect(agent.generateText).toBeDefined();
   expect(agent.streamText).toBeDefined();
 
-  expect(agent.addFeedback).toBeDefined();
   expect(agent.addMessage).toBeDefined();
   expect(agent.addObservation).toBeDefined();
+  expect(agent.addFeedback).toBeDefined();
   expect(agent.addPlan).toBeDefined();
 
+  expect(agent.getMessages).toBeDefined();
+  expect(agent.getObservations).toBeDefined();
+  expect(agent.getFeedback).toBeDefined();
+  expect(agent.getPlans).toBeDefined();
+
   expect(agent.interact).toBeDefined();
 });
 
@@ -45,6 +57,11 @@ test('agent.addMessage() adds to message history', () => {
       content: 'msg 1',
     })
   );
+  expect(agent.getMessages()).toContainEqual(
+    expect.objectContaining({
+      content: 'msg 1',
+    })
+  );
 
   expect(agent.select((c) => c.messages)).toContainEqual(
     expect.objectContaining({
@@ -53,6 +70,13 @@ test('agent.addMessage() adds to message history', () => {
       timestamp: expect.any(Number),
     })
   );
+  expect(agent.getMessages()).toContainEqual(
+    expect.objectContaining({
+      content: 'response 1',
+      sessionId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
 });
 
 test('agent.addFeedback() adds to feedback', () => {
@@ -83,6 +107,17 @@ test('agent.addFeedback() adds to feedback', () => {
       timestamp: expect.any(Number),
     })
   );
+  expect(agent.getFeedback()).toContainEqual(
+    expect.objectContaining({
+      attributes: {
+        score: -1,
+      },
+      goal: 'Win the game',
+      observationId: 'obs-1',
+      sessionId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
 });
 
 test('agent.addObservation() adds to observations', () => {
@@ -111,6 +146,46 @@ test('agent.addObservation() adds to observations', () => {
   );
 });
 
+test('agent.addObservation() adds to observations with machine hash', () => {
+  const agent = createAgent({
+    name: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const machine = createMachine({
+    initial: 'playing',
+    states: {
+      playing: {
+        on: {
+          play: 'lost',
+        },
+      },
+      lost: {},
+    },
+  });
+
+  const observation = agent.addObservation({
+    prevState: { value: 'playing', context: {} },
+    event: { type: 'play', position: 3 },
+    state: { value: 'lost', context: {} },
+    machine,
+  });
+
+  expect(observation.sessionId).toEqual(agent.sessionId);
+
+  expect(agent.select((c) => c.observations)).toContainEqual(
+    expect.objectContaining({
+      prevState: { value: 'playing', context: {} },
+      event: { type: 'play', position: 3 },
+      state: { value: 'lost', context: {} },
+      machineHash: expect.any(String),
+      sessionId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
+});
+
 test('agent.interact() observes machine actors (no 2nd arg)', () => {
   const machine = createMachine({
     initial: 'a',
@@ -140,6 +215,12 @@ test('agent.interact() observes machine actors (no 2nd arg)', () => {
       state: expect.objectContaining({ value: 'a' }),
     })
   );
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      prevState: undefined,
+      state: expect.objectContaining({ value: 'a' }),
+    })
+  );
 
   actor.send({ type: 'NEXT' });
 
@@ -175,7 +256,8 @@ test('Agents can use a custom adapter', async () => {
   expect(res.text).toEqual('Response');
 });
 
-test.skip('You can listen for emitted agent events', () => {
+test('You can listen for feedback events', () => {
+  const fn = vi.fn();
   const agent = createAgent({
     name: 'test',
     events: {},
@@ -183,5 +265,242 @@ test.skip('You can listen for emitted agent events', () => {
     model: {} as any,
   });
 
-  agent.on('feedback', (x) => {});
+  agent.on('feedback', fn);
+
+  agent.addFeedback({
+    attributes: {
+      score: -1,
+    },
+    goal: 'Win the game',
+    observationId: 'obs-1',
+  });
+
+  expect(fn).toHaveBeenCalled();
+});
+
+test('You can listen for plan events', async () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    name: 'test',
+    model: {} as any,
+    events: {
+      WIN: z.object({}),
+    },
+    adapter: {
+      generateText: async (arg) => {
+        const keys = Object.keys(arg.tools!);
+
+        if (keys.length !== 1) {
+          throw new Error('Expected only 1 choice');
+        }
+
+        return {
+          toolResults: [
+            {
+              result: {
+                type: keys[0],
+              },
+            },
+          ],
+        } as any as AgentGenerateTextResult;
+      },
+      streamText: {} as any,
+    },
+  });
+
+  agent.on('plan', fn);
+
+  await agent.decide({
+    goal: 'Win the game',
+    state: {
+      value: 'playing',
+      context: {},
+    },
+    machine: createMachine({
+      initial: 'playing',
+      states: {
+        playing: {
+          on: {
+            WIN: {
+              target: 'won',
+            },
+          },
+        },
+        won: {},
+      },
+    }),
+  });
+
+  expect(fn).toHaveBeenCalledWith(
+    expect.objectContaining({
+      plan: expect.objectContaining({
+        nextEvent: {
+          type: 'WIN',
+        },
+      }),
+    })
+  );
+});
+
+test('agent.types provides context and event types', () => {
+  const agent = createAgent({
+    model: {} as any,
+    events: {
+      setScore: z.object({
+        score: z.number(),
+      }),
+    },
+    context: {
+      score: z.number(),
+    },
+  });
+
+  agent.types satisfies { context: any; events: any };
+
+  agent.types.context satisfies { score: number };
+
+  // @ts-expect-error
+  agent.types.context satisfies { score: string };
 });
+
+test.each(['generateText', 'streamText'] as const)(
+  'can provide a correlation ID (%s)',
+  async (method) => {
+    const agent = createAgent({
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts: any) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    const promise = new Promise<AgentMessage>((res) => {
+      agent.onMessage((msg) => {
+        if (msg.role === 'assistant') {
+          res(msg);
+        }
+      });
+    });
+
+    await agent[method]({
+      prompt: 'hi',
+      correlationId: 'c-1',
+    });
+
+    const msg = await promise;
+
+    expect(msg.correlationId).toBe('c-1');
+    expect(msg.parentCorrelationId).toBe(undefined);
+  }
+);
+
+test.each(['generateText', 'streamText'] as const)(
+  'correlation IDs are automatically generated if not provided (%s)',
+  async (method) => {
+    const agent = createAgent({
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts: any) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    await agent[method]({
+      prompt: 'hi',
+    });
+
+    const messages = agent.getMessages();
+
+    expect(messages[0]?.correlationId).toEqual(expect.stringMatching(/.+/));
+    expect(messages[0]?.role).toBe('user');
+    expect(messages[1]?.correlationId).toEqual(expect.stringMatching(/.+/));
+    expect(messages[1]?.role).toBe('assistant');
+
+    expect(messages[0]!.correlationId).toEqual(messages[1]!.correlationId);
+  }
+);
+
+test.each(['generateText', 'streamText'] as const)(
+  'can provide a parent correlation ID (%s)',
+  async (method) => {
+    const agent = createAgent({
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts: any) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    await agent[method]({
+      prompt: 'hi',
+      correlationId: 'c-1',
+      parentCorrelationId: 'c-0',
+    });
+
+    const msg = agent.getMessages().find((msg) => msg.role === 'assistant')!;
+
+    expect(msg.correlationId).toBe('c-1');
+    expect(msg.parentCorrelationId).toBe('c-0');
+  }
+);
+
+test.each(['generateText', 'streamText'] as const)(
+  'can add feedback to a correlation (%s)',
+  async (method) => {
+    const agent = createAgent({
+      name: 'test',
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts: any) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    const res = await agent[method]({
+      prompt: 'test',
+    });
+
+    agent.addFeedback({
+      correlationId: res.correlationId,
+      reward: -1,
+    });
+
+    const message = agent.getMessages()[0]!;
+    const feedback = agent.getFeedback()[0]!;
+
+    expect(message.correlationId).toBeDefined();
+    expect(feedback.correlationId).toEqual(message.correlationId);
+  }
+);