@statelyai/agent 2.0.0-next.0 → 2.0.0-next.2

package/examples/verify.ts

@@ -1,11 +1,11 @@
 import { assign, createActor, setup, log } from 'xstate';
 import { fromTerminal } from './helpers/helpers';
-import { createAgent, fromDecision } from '../src';
+import { createAgent, EventsFromAgent, fromDecision } from '../src';
 import { z } from 'zod';
 import { openai } from '@ai-sdk/openai';
 
 const agent = createAgent({
-  name: 'verifier',
+  id: 'verifier',
   model: openai('gpt-3.5-turbo-16k-0613'),
   events: {
     'agent.validateAnswer': z.object({
@@ -35,7 +35,7 @@ const machine = setup({
       answer: string | null;
       validation: string | null;
     },
-    events: agent.types.events,
+    events: {} as EventsFromAgent<typeof agent>,
   },
   actors: {
     getFromTerminal: fromTerminal,

package/examples/weather-agent.ts

@@ -0,0 +1,141 @@
+import { z } from 'zod';
+import { createAgent } from '../src';
+import { assign, createActor, fromPromise, setup } from 'xstate';
+// import { anthropic } from '@ai-sdk/anthropic';
+import { fromTerminal } from './helpers/helpers';
+import { openai } from '@ai-sdk/openai';
+
+// Create the weather agent
+const agent = createAgent({
+  id: 'weather-agent',
+  model: openai('gpt-4o'),
+  events: {
+    'weather.check': z.object({
+      location: z.string().describe('The location to check weather for'),
+    }),
+    'weather.report': z.object({
+      temperature: z.string(),
+      conditions: z.string(),
+    }),
+    'agent.respond': z.object({
+      response: z.string(),
+    }),
+  },
+});
+
+// Create the weather tool/service
+const getWeather = (location: string) => {
+  if (
+    location.toLowerCase().includes('sf') ||
+    location.toLowerCase().includes('san francisco')
+  ) {
+    return {
+      temperature: '60 degrees',
+      conditions: 'foggy',
+    };
+  }
+  return {
+    temperature: '90 degrees',
+    conditions: 'sunny',
+  };
+};
+
+// Create the state machine
+const machine = setup({
+  types: {
+    context: {} as {
+      input: string | null;
+      location: string | null;
+      weather: any;
+    },
+  },
+  actors: {
+    getFromTerminal: fromTerminal,
+    getWeather: fromPromise(async ({ input }: { input: string }) => {
+      return getWeather(input);
+    }),
+  },
+}).createMachine({
+  context: {
+    input: null,
+    location: null,
+    weather: null,
+  },
+  initial: 'user',
+  states: {
+    user: {
+      invoke: {
+        src: 'getFromTerminal',
+        input: 'Ask me about the weather!',
+        onDone: {
+          actions: assign({
+            input: ({ event }) => event.output,
+          }),
+          target: 'processing',
+        },
+      },
+    },
+    processing: {
+      entry: () => console.log('Processing...'),
+      on: {
+        'weather.check': {
+          actions: assign({
+            input: ({ event }) => event.location,
+          }),
+          target: 'checking',
+        },
+        'agent.respond': {
+          actions: ({ event }) => console.log(event.response),
+          target: 'responded',
+        },
+      },
+    },
+    checking: {
+      entry: ({ event }) => console.log('Checking weather...', event),
+      invoke: {
+        src: 'getWeather',
+        input: ({ context }) => context.input!,
+        onDone: {
+          actions: assign({
+            weather: ({ event }) => event.output,
+          }),
+          target: 'responding',
+        },
+      },
+    },
+    responding: {
+      on: {
+        'agent.respond': {
+          actions: ({ event }) => console.log(event.response),
+          target: 'responded',
+        },
+      },
+    },
+    responded: {
+      after: {
+        1000: { target: 'user' },
+      },
+    },
+  },
+});
+
+// Create and start the actor
+const actor = createActor(machine).start();
+
+agent.interact(actor, (obs) => {
+  if (obs.state.matches('processing')) {
+    return {
+      goal: 'Determine if user is asking about weather and for which location. If so, get the weather. Otherwise, respond to the user.',
+      context: obs.state.context,
+      messages: agent.getMessages(),
+    };
+  }
+
+  if (obs.state.matches('responding')) {
+    return {
+      goal: 'Provide a natural response about the weather in ${context.location}',
+      context: obs.state.context,
+      messages: agent.getMessages(),
+    };
+  }
+});

package/examples/weather.ts

@@ -1,4 +1,4 @@
-import { createAgent, fromDecision } from '../src';
+import { createAgent, EventsFromAgent, fromDecision } from '../src';
 import { assign, createActor, fromPromise, log, setup } from 'xstate';
 import { fromTerminal } from './helpers/helpers';
 import { z } from 'zod';
@@ -49,7 +49,7 @@ const getWeather = fromPromise(async ({ input }: { input: string }) => {
 });
 
 const agent = createAgent({
-  name: 'weather',
+  id: 'weather',
   model: openai('gpt-4o-mini'),
   events: {
     'agent.getWeather': z.object({
@@ -75,10 +75,10 @@ const machine = setup({
   types: {
     context: {} as {
       location: string;
-      history: string[];
       count: number;
+      result: string | null;
     },
-    events: agent.types.events,
+    events: {} as EventsFromAgent<typeof agent>,
   },
   actors: {
     agent: fromDecision(agent),
@@ -90,7 +90,7 @@ const machine = setup({
   context: {
     location: '',
     count: 0,
-    history: [],
+    result: null,
   },
   states: {
     getLocation: {
@@ -111,15 +111,6 @@ const machine = setup({
     },
     decide: {
       entry: log('Deciding...'),
-      invoke: {
-        src: 'agent',
-        input: ({ context }) => ({
-          context: {
-            location: context.location,
-          },
-          goal: `Decide what to do based on the given location, which may or may not be a location`,
-        }),
-      },
       on: {
         'agent.getWeather': {
           actions: log(({ event }) => event),
@@ -138,6 +129,7 @@ const machine = setup({
             log(({ event }) => event.output),
             assign({
               count: ({ context }) => context.count + 1,
+              result: ({ event }) => event.output,
             }),
           ],
           target: 'reportWeather',
@@ -145,13 +137,6 @@ const machine = setup({
       },
     },
     reportWeather: {
-      invoke: {
-        src: 'agent',
-        input: ({ context }) => ({
-          goal: 'Report the weather', // TODO
-          context,
-        }),
-      },
       on: {
         'agent.reportWeather': {
           actions: log(({ event }) => event),
@@ -169,7 +154,22 @@ const machine = setup({
 });
 
 const actor = createActor(machine);
-actor.subscribe((s) => {
-  console.log(s.value);
-});
 actor.start();
+
+agent.interact(actor, (obs) => {
+  if (obs.state.matches('decide')) {
+    return {
+      goal: `Decide what to do based on the given location, which may or may not be a location`,
+      context: {
+        location: obs.state.context.location,
+      },
+    };
+  }
+
+  if (obs.state.matches('reportWeather')) {
+    return {
+      goal: `Report the weather for the given location`,
+      context: obs.state.context,
+    };
+  }
+});

package/examples/wiki.ts

@@ -4,7 +4,7 @@ import { openai } from '@ai-sdk/openai';
 import { CoreMessage, generateText, streamText } from 'ai';
 
 const agent = createAgent({
-  name: 'wiki',
+  id: 'wiki',
   model: openai('gpt-4o-mini'),
   events: {
     provideAnswer: z.object({

package/examples/word.ts

@@ -1,11 +1,16 @@
 import { assign, createActor, log, setup } from 'xstate';
 import { fromTerminal } from './helpers/helpers';
-import { createAgent, fromDecision } from '../src';
+import {
+  ContextFromAgent,
+  createAgent,
+  fromDecision,
+  TypesFromAgent,
+} from '../src';
 import { z } from 'zod';
 import { openai } from '@ai-sdk/openai';
 
 const agent = createAgent({
-  name: 'word',
+  id: 'word',
   model: openai('gpt-4o-mini'),
   context: {
     word: z.string().nullable().describe('The word to guess'),
@@ -36,13 +41,10 @@ const context = {
   word: null,
   guessedWord: null,
   lettersGuessed: [],
-} satisfies typeof agent.types.context;
+} satisfies ContextFromAgent<typeof agent>;
 
 const wordGuesserMachine = setup({
-  types: {
-    context: agent.types.context,
-    events: agent.types.events,
-  },
+  types: {} as TypesFromAgent<typeof agent>,
   actors: {
     agent: fromDecision(agent),
     getFromTerminal: fromTerminal,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@statelyai/agent",
-  "version": "2.0.0-next.0",
+  "version": "2.0.0-next.2",
   "description": "Stateful agents that make decisions based on finite-state machine models",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -15,6 +15,7 @@
   "author": "David Khourshid <david@stately.ai>",
   "license": "MIT",
   "devDependencies": {
+    "@ai-sdk/anthropic": "^0.0.54",
     "@ai-sdk/openai": "^0.0.40",
     "@changesets/changelog-github": "^0.5.0",
     "@changesets/cli": "^2.27.9",
@@ -24,6 +25,7 @@
     "@tavily/core": "^0.0.2",
     "@types/node": "^20.17.1",
     "@types/object-hash": "^3.0.6",
+    "@vitest/coverage-v8": "^2.1.4",
     "dotenv": "^16.4.5",
     "json-schema-to-ts": "^3.1.1",
     "ts-node": "^10.9.2",
@@ -38,7 +40,7 @@
   },
   "dependencies": {
     "@xstate/graph": "^2.0.1",
-    "ai": "^3.4.20",
+    "ai": "^3.4.31",
     "ajv": "^8.17.1",
     "object-hash": "^3.0.0",
     "xstate": "^5.18.2",
@@ -52,6 +54,7 @@
     "example": "ts-node examples/helpers/runner.ts",
     "changeset": "changeset",
     "release": "changeset publish",
-    "version": "changeset version"
+    "version": "changeset version",
+    "coverage": "vitest run --coverage"
   }
 }

package/src/agent.test.ts

@@ -1,5 +1,5 @@
 import { test, expect, vi } from 'vitest';
-import { createAgent } from './';
+import { createAgent, TypesFromAgent } from './';
 import { createActor, createMachine } from 'xstate';
 import { LanguageModelV1CallOptions } from 'ai';
 import { z } from 'zod';
@@ -7,7 +7,7 @@ import { dummyResponseValues, MockLanguageModelV1 } from './mockModel';
 
 test('an agent has the expected interface', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: new MockLanguageModelV1(),
   });
@@ -17,12 +17,12 @@ test('an agent has the expected interface', () => {
   expect(agent.addMessage).toBeDefined();
   expect(agent.addObservation).toBeDefined();
   expect(agent.addFeedback).toBeDefined();
-  expect(agent.addPlan).toBeDefined();
+  expect(agent.addDecision).toBeDefined();
 
   expect(agent.getMessages).toBeDefined();
   expect(agent.getObservations).toBeDefined();
   expect(agent.getFeedback).toBeDefined();
-  expect(agent.getPlans).toBeDefined();
+  expect(agent.getDecisions).toBeDefined();
 
   expect(agent.interact).toBeDefined();
 });
@@ -31,7 +31,7 @@ test('agent.addMessage() adds to message history', () => {
   const model = new MockLanguageModelV1();
 
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model,
   });
@@ -65,7 +65,7 @@ test('agent.addMessage() adds to message history', () => {
 
 test('agent.addFeedback() adds to feedback', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -106,7 +106,7 @@ test('agent.addFeedback() adds to feedback', () => {
 
 test('agent.addObservation() adds to observations', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -132,7 +132,7 @@ test('agent.addObservation() adds to observations', () => {
 
 test('agent.addObservation() adds to observations (initial state)', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -154,7 +154,7 @@ test('agent.addObservation() adds to observations (initial state)', () => {
 
 test('agent.addObservation() adds to observations with machine hash', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -194,7 +194,7 @@ test('agent.addObservation() adds to observations with machine hash', () => {
 
 test('agent.addFeedback() adds to feedback (with observation)', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -251,7 +251,7 @@ test('agent.interact() observes machine actors (no 2nd arg)', () => {
   });
 
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -289,7 +289,7 @@ test('agent.interact() observes machine actors (no 2nd arg)', () => {
 test('You can listen for feedback events', () => {
   const fn = vi.fn();
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -307,7 +307,7 @@ test('You can listen for feedback events', () => {
   expect(fn).toHaveBeenCalled();
 });
 
-test('You can listen for plan events', async () => {
+test('You can listen for decision events', async () => {
   const fn = vi.fn();
   const model = new MockLanguageModelV1({
     doGenerate: async (params: LanguageModelV1CallOptions) => {
@@ -332,14 +332,14 @@ test('You can listen for plan events', async () => {
   });
 
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     model,
     events: {
       WIN: z.object({}),
     },
   });
 
-  agent.on('plan', fn);
+  agent.on('decision', fn);
 
   await agent.decide({
     goal: 'Win the game',
@@ -364,7 +364,7 @@ test('You can listen for plan events', async () => {
 
   expect(fn).toHaveBeenCalledWith(
     expect.objectContaining({
-      plan: expect.objectContaining({
+      decision: expect.objectContaining({
         nextEvent: {
           type: 'WIN',
         },
@@ -386,10 +386,152 @@ test('agent.types provides context and event types', () => {
     },
   });
 
-  agent.types satisfies { context: any; events: any };
+  let types = {} as TypesFromAgent<typeof agent>;
 
-  agent.types.context satisfies { score: number };
+  types satisfies { context: any; events: any };
+
+  types.context satisfies { score: number };
 
   // @ts-expect-error
-  agent.types.context satisfies { score: string };
+  types.context satisfies { score: string };
+});
+
+test('It allows unrecognized events', () => {
+  const agent = createAgent({
+    model: {} as any,
+    events: {},
+    context: {},
+  });
+
+  expect(() => {
+    agent.send({
+      // @ts-expect-error
+      type: 'unrecognized',
+    });
+  }).not.toThrow();
+});
+
+test('You can listen for message events', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  agent.onMessage(fn);
+
+  const message = {
+    role: 'user' as const,
+    content: [{ type: 'text' as const, text: 'test message' }],
+  };
+
+  agent.addMessage(message);
+
+  expect(fn).toHaveBeenCalledWith(
+    expect.objectContaining({
+      role: 'user',
+      content: [{ type: 'text', text: 'test message' }],
+      episodeId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
+});
+
+test('agent.getDecisions() returns decisions from context', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+    strategy: async (agent) => {
+      return {
+        episodeId: agent.episodeId,
+        strategy: 'test-strategy',
+        goal: '',
+        goalState: undefined,
+        paths: [
+          {
+            state: undefined,
+            steps: [],
+          },
+        ],
+        nextEvent: undefined,
+        timestamp: Date.now(),
+      };
+    },
+  });
+
+  const decisions = agent.getDecisions();
+
+  expect(decisions).toBeDefined();
+  expect(Array.isArray(decisions)).toBe(true);
+});
+
+test('Event listeners can be unsubscribed', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const subscription = agent.on('message', fn);
+
+  agent.addMessage({
+    role: 'user',
+    content: [{ type: 'text', text: 'first message' }],
+  });
+
+  expect(fn).toHaveBeenCalledTimes(1);
+
+  subscription.unsubscribe();
+
+  agent.addMessage({
+    role: 'user',
+    content: [{ type: 'text', text: 'second message' }],
+  });
+
+  expect(fn).toHaveBeenCalledTimes(1); // Still only called once
+});
+
+test('agent.observe() adds observations from actor snapshots', () => {
+  const machine = createMachine({
+    initial: 'idle',
+    states: {
+      idle: {
+        on: { START: 'running' },
+      },
+      running: {},
+    },
+  });
+
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const actor = createActor(machine);
+  const subscription = agent.observe(actor);
+
+  actor.start();
+  actor.send({ type: 'START' });
+
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      state: expect.objectContaining({ value: 'idle' }),
+      machineHash: expect.any(String),
+    })
+  );
+
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      prevState: expect.objectContaining({ value: 'idle' }),
+      event: { type: 'START' },
+      state: expect.objectContaining({ value: 'running' }),
+      machineHash: expect.any(String),
+    })
+  );
+
+  subscription.unsubscribe();
 });