@statelyai/agent 1.0.0-beta.0 → 1.0.0

package/examples/ticTacToe.ts

@@ -2,7 +2,6 @@ import { assign, setup, assertEvent, createActor } from 'xstate';
 import { z } from 'zod';
 import { createAgent, fromDecision, fromTextStream } from '../src';
 import { openai } from '@ai-sdk/openai';
-import { defaultToolCallTemplate } from '../src/templates/defaultToolCall';
 
 const agent = createAgent({
   name: 'tic-tac-toe-bot',
@@ -24,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 = [
@@ -65,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),
@@ -82,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 = '';
@@ -173,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

@@ -12,7 +12,7 @@ const context = {
 
 const agent = createAgent({
   name: 'word',
-  model: openai('gpt-4-1106-preview'),
+  model: openai('gpt-4o'),
   events: {
     'agent.guessLetter': z.object({
       letter: z.string().min(1).max(1).describe('The letter guessed'),
@@ -36,18 +36,21 @@ const agent = createAgent({
 const wordGuesserMachine = setup({
   types: {
     context: {} as typeof context,
-    events: agent.eventTypes,
+    events: agent.types.events,
   },
   actors: {
     agent: fromDecision(agent),
     getFromTerminal,
   },
+  actions: {
+    resetContext: assign(context),
+  },
 }).createMachine({
   initial: 'providingWord',
   context,
   states: {
     providingWord: {
-      entry: assign(context),
+      entry: 'resetContext',
       invoke: {
         src: 'getFromTerminal',
         input: 'Enter a word, and an agent will try to guess it.',

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@statelyai/agent",
-  "version": "1.0.0-beta.0",
-  "description": "",
+  "version": "1.0.0",
+  "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",
@@ -9,27 +9,35 @@
     "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": [],
+  "keywords": [
+    "ai",
+    "state machine",
+    "agent",
+    "rl",
+    "reinforcement learning"
+  ],
   "author": "",
   "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.2",
+    "@types/node": "^20.14.10",
+    "@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",
+    "typescript": "^5.5.3",
     "vitest": "^1.6.0",
     "wikipedia": "^2.1.2",
     "zod": "^3.23.8"
@@ -38,11 +46,11 @@
     "access": "public"
   },
   "dependencies": {
-    "@ai-sdk/openai": "^0.0.13",
+    "@ai-sdk/openai": "^0.0.31",
     "@xstate/graph": "^2.0.0",
-    "ai": "^3.1.32",
-    "nanoid": "^5.0.7",
-    "xstate": "^5.13.2"
+    "ai": "^3.2.22",
+    "object-hash": "^3.0.0",
+    "xstate": "^5.15.0"
   },
   "packageManager": "pnpm@8.11.0"
 }

package/src/agent.test.ts

@@ -1,6 +1,8 @@
-import { test, expect } from 'vitest';
+import { test, expect, vi } from 'vitest';
 import { 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 +15,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 +52,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 +65,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 +102,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 +141,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 +210,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 +251,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 +260,100 @@ 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 GenerateTextResult<any>;
+      },
+      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 };
 });

package/src/agent.ts

@@ -7,11 +7,11 @@ import {
   Observer,
   toObserver,
 } from 'xstate';
-import { ZodEventMapping } from './schemas';
+import { ZodContextMapping, ZodEventMapping } from './schemas';
 import {
   Agent,
   AgentLogic,
-  AgentMessageHistory,
+  AgentMessage,
   AgentPlanner,
   EventsFromZodEventMapping,
   GenerateTextOptions,
@@ -20,12 +20,14 @@ import {
   ObservedState,
   AgentObservationInput,
   AgentMemoryContext,
+  AgentObservation,
+  ContextFromZodContextMapping,
 } from './types';
 import { simplePlanner } from './planners/simplePlanner';
 import { agentGenerateText, agentStreamText } from './text';
 import { agentDecide } from './decision';
 import { vercelAdapter } from './adapters/vercel';
-import { nanoid } from 'nanoid';
+import { getMachineHash, randomId } from './utils';
 
 export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
   (state, event, { emit }) => {
@@ -80,24 +82,42 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
 );
 
 export function createAgent<
+  const TContextSchema extends ZodContextMapping,
   const TEventSchemas extends ZodEventMapping,
-  TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>
+  TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
+  TContext = ContextFromZodContextMapping<TContextSchema>
 >({
   name,
   description,
   model,
   events,
-  planner = simplePlanner as AgentPlanner<Agent<TEvents>>,
+  context,
+  planner = simplePlanner as AgentPlanner<Agent<TContext, TEvents>>,
   stringify = JSON.stringify,
   getMemory,
   logic = agentLogic as AgentLogic<TEvents>,
   adapter = vercelAdapter,
   ...generateTextOptions
 }: {
+  /**
+   * The unique identifier for the agent.
+   *
+   * This should be the same across all sessions of a specific agent, as it can be
+   * used to retrieve memory for this agent.
+   *
+   * @example
+   * ```ts
+   * const agent = createAgent({
+   *  id: 'recipe-assistant',
+   *  // ...
+   * });
+   * ```
+   */
+  id?: string;
   /**
    * The name of the agent
    */
-  name: string;
+  name?: string;
   /**
    * A description of the role of the agent
    */
@@ -107,21 +127,22 @@ export function createAgent<
    * that the agent knows about.
    */
   events: TEventSchemas;
-  planner?: AgentPlanner<Agent<TEvents>>;
+  context?: TContextSchema;
+  planner?: AgentPlanner<Agent<TContext, TEvents>>;
   stringify?: typeof JSON.stringify;
   /**
    * A function that retrieves the agent's long term memory
    */
-  getMemory?: (agent: Agent<any>) => AgentLongTermMemory;
+  getMemory?: (agent: Agent<TContext, TEvents>) => AgentLongTermMemory;
   /**
    * Agent logic
    */
   logic?: AgentLogic<TEvents>;
   adapter?: AIAdapter;
-} & GenerateTextOptions): Agent<TEvents> {
-  const messageHistoryListeners: Observer<AgentMessageHistory>[] = [];
+} & GenerateTextOptions): Agent<TContext, TEvents> {
+  const messageHistoryListeners: Observer<AgentMessage>[] = [];
 
-  const agent = createActor(logic) as unknown as Agent<TEvents>;
+  const agent = createActor(logic) as unknown as Agent<TContext, TEvents>;
   agent.events = events;
   agent.model = model;
   agent.name = name;
@@ -144,7 +165,7 @@ export function createAgent<
   agent.addMessage = (messageInput) => {
     const message = {
       ...messageInput,
-      id: messageInput.id ?? nanoid(),
+      id: messageInput.id ?? randomId(),
       timestamp: messageInput.timestamp ?? Date.now(),
       sessionId: agent.sessionId,
     };
@@ -155,6 +176,7 @@ export function createAgent<
 
     return message;
   };
+  agent.getMessages = () => agent.getSnapshot().context.messages;
 
   agent.generateText = (opts) => agentGenerateText(agent, opts);
 
@@ -172,14 +194,21 @@ export function createAgent<
     });
     return feedback;
   };
+  agent.getFeedback = () => agent.getSnapshot().context.feedback;
 
   agent.addObservation = (observationInput) => {
+    const { prevState, event, state } = observationInput;
     const observation = {
-      ...observationInput,
-      id: observationInput.id ?? nanoid(),
+      prevState,
+      event,
+      state,
+      id: observationInput.id ?? randomId(),
       sessionId: agent.sessionId,
       timestamp: observationInput.timestamp ?? Date.now(),
-    };
+      machineHash: observationInput.machine
+        ? getMachineHash(observationInput.machine)
+        : undefined,
+    } satisfies AgentObservation<any>;
 
     agent.send({
       type: 'agent.observe',
@@ -188,6 +217,7 @@ export function createAgent<
 
     return observation;
   };
+  agent.getObservations = () => agent.getSnapshot().context.observations;
 
   agent.addPlan = (plan) => {
     agent.send({
@@ -195,8 +225,9 @@ export function createAgent<
       plan,
     });
   };
+  agent.getPlans = () => agent.getSnapshot().context.plans;
 
-  agent.interact = (actorRef, getInput) => {
+  agent.interact = ((actorRef, getInput) => {
     let prevState: ObservedState | undefined = undefined;
     let subscribed = true;
 
@@ -234,7 +265,8 @@ export function createAgent<
           event: inspEvent.event,
           prevState,
           state: inspEvent.snapshot as any,
-        };
+          machine: (actorRef as any).src,
+        } satisfies AgentObservationInput;
 
         await handleObservation(observationInput);
       },
@@ -246,6 +278,7 @@ export function createAgent<
         prevState: undefined,
         event: { type: '' }, // TODO: unknown events?
         state: actorRef.getSnapshot(),
+        machine: (actorRef as any).src,
       });
     }
 
@@ -254,7 +287,9 @@ export function createAgent<
         subscribed = false;
       }, // TODO: make this actually unsubscribe
     };
-  };
+  }) as typeof agent.interact;
+
+  agent.types = {} as any;
 
   agent.start();
 

package/src/decision.ts

@@ -1,17 +1,18 @@
 import { AnyMachineSnapshot, fromPromise } from 'xstate';
 import {
-  Agent,
+  AnyAgent,
   AgentDecideOptions,
   AgentDecisionLogic,
   AgentDecisionInput,
   AgentPlanner,
+  AgentPlan,
 } from './types';
 import { simplePlanner } from './planners/simplePlanner';
 
-export async function agentDecide<T extends Agent<any>>(
+export async function agentDecide<T extends AnyAgent>(
   agent: T,
   options: AgentDecideOptions
-) {
+): Promise<AgentPlan<any> | undefined> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
@@ -44,9 +45,9 @@ export async function agentDecide<T extends Agent<any>>(
 }
 
 export function fromDecision(
-  agent: Agent<any>,
+  agent: AnyAgent,
   defaultInput?: AgentDecisionInput
-) {
+): AgentDecisionLogic<any> {
   return fromPromise(async ({ input, self }) => {
     const parentRef = self._parent;
     if (!parentRef) {