@statelyai/agent 1.0.0-beta.1 → 1.0.0

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,6 +1,6 @@
 {
   "name": "@statelyai/agent",
-  "version": "1.0.0-beta.1",
+  "version": "1.0.0",
   "description": "Stateful agents that make decisions based on finite-state machine models",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -27,16 +27,17 @@
   "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.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"
@@ -47,8 +48,9 @@
   "dependencies": {
     "@ai-sdk/openai": "^0.0.31",
     "@xstate/graph": "^2.0.0",
-    "ai": "^3.2.1",
-    "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 { randomId } from './utils';
+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;
@@ -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,
+      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) {

package/src/index.ts

@@ -1,4 +1,4 @@
 export { createAgent } from './agent';
-export { fromText, fromTextStream, agentGenerateText } from './text';
-export { fromDecision, agentDecide } from './decision';
+export { fromText, fromTextStream } from './text';
+export { fromDecision } from './decision';
 export * from './types';

package/src/planners/shortestPathPlanner.ts

@@ -1,7 +1,7 @@
-import { Agent, AgentPlan, AgentPlanInput } from '../types';
+import { Agent, AgentPlan, AgentPlanInput, AnyAgent } from '../types';
 import { getShortestPaths } from '@xstate/graph';
 
-export async function simplePlanner<T extends Agent<any>>(
+export async function simplePlanner<T extends AnyAgent>(
   agent: T,
   input: AgentPlanInput<any>
 ): Promise<AgentPlan<any> | undefined> {

package/src/planners/simplePlanner.ts

@@ -1,15 +1,16 @@
 import { CoreTool, tool } from 'ai';
 import {
-  Agent,
   AgentPlan,
   AgentPlanInput,
   ObservedState,
   PromptTemplate,
   TransitionData,
+  AnyAgent,
 } from '../types';
 import { getAllTransitions } from '../utils';
 import { AnyStateMachine } from 'xstate';
 import { defaultTextTemplate } from '../templates/defaultText';
+import { getMessages } from '../text';
 
 function getTransitions(
   state: ObservedState,
@@ -27,11 +28,11 @@ const simplePlannerPromptTemplate: PromptTemplate<any> = (data) => {
   return `
 ${defaultTextTemplate(data)}
 
-Only make a single tool call to achieve the above goal.
+Make at most one tool call to achieve the above goal. If the goal cannot be achieved with any tool calls, do not make any tool call.
   `.trim();
 };
 
-export async function simplePlanner<T extends Agent<any>>(
+export async function simplePlanner<T extends AnyAgent>(
   agent: T,
   input: AgentPlanInput<any>
 ): Promise<AgentPlan<any> | undefined> {
@@ -81,7 +82,7 @@ export async function simplePlanner<T extends Agent<any>>(
     toolMap[toolTransitionData.name] = tool({
       description: toolZodType?.description ?? toolTransitionData.description,
       parameters: toolZodType,
-      execute: async (params) => {
+      execute: async (params: Record<string, any>) => {
         const event = {
           type: toolTransitionData.eventType,
           ...params,
@@ -92,23 +93,32 @@ export async function simplePlanner<T extends Agent<any>>(
     });
   }
 
+  if (!Object.keys(toolMap).length) {
+    // No valid transitions for the specified tools
+    return undefined;
+  }
+
   // Create a prompt with the given context and goal.
-  // The template is used to ensure that a single tool call is made.
+  // The template is used to ensure that a single tool call at most is made.
   const prompt = simplePlannerPromptTemplate({
     context: input.state.context,
     goal: input.goal,
   });
 
+  const messages = await getMessages(agent, prompt, input);
+
   const result = await agent.generateText({
+    toolChoice: 'required',
     ...input,
     prompt,
+    messages,
     tools: toolMap,
-    toolChoice: 'required',
   });
 
   const singleResult = result.toolResults[0];
 
   if (!singleResult) {
+    console.log(toolMap);
     // TODO: retries?
     console.warn('No tool call results returned');
     return undefined;
@@ -117,11 +127,12 @@ export async function simplePlanner<T extends Agent<any>>(
   return {
     goal: input.goal,
     state: input.state,
-    steps: [
-      {
-        event: singleResult.result,
-      },
-    ],
+    execute: async (state) => {
+      if (JSON.stringify(state) === JSON.stringify(input.state)) {
+        return singleResult.result;
+      }
+      return undefined;
+    },
     nextEvent: singleResult.result,
     sessionId: agent.sessionId,
     timestamp: Date.now(),

package/src/schemas.ts

@@ -1,15 +1,11 @@
-import type { SomeZodObject } from 'zod';
-import { AnyEventObject } from 'xstate';
-import { ObservedState } from './types';
+import { ZodType, type SomeZodObject } from 'zod';
 
 export type ZodEventMapping = {
   // map event types to Zod types
   [eventType: string]: SomeZodObject;
 };
 
-export type ZodActionMapping = {
-  [eventType: string]: {
-    schema: SomeZodObject;
-    action: (state: ObservedState, event: AnyEventObject) => Promise<void>;
-  };
+export type ZodContextMapping = {
+  // map context keys to Zod types
+  [contextKey: string]: ZodType;
 };

package/src/strategies/chain-of-note.ts

@@ -1,7 +1,7 @@
 import { GenerateTextResult, LanguageModel } from 'ai';
 import wiki, { wikiSearchResult, wikiSummary } from 'wikipedia';
 import { assign, fromPromise, setup } from 'xstate';
-import { Agent } from '../types';
+import { AnyAgent } from '../types';
 
 const searchWiki = fromPromise(
   async ({
@@ -44,7 +44,7 @@ export const chainOfNote = setup({
   types: {
     input: {} as {
       model: LanguageModel;
-      agent: Agent<any>;
+      agent: AnyAgent;
       prompt: string;
     },
     context: {} as {
@@ -56,7 +56,7 @@ export const chainOfNote = setup({
           }[]
         | null;
       model: LanguageModel;
-      agent: Agent<any>;
+      agent: AnyAgent;
       prompt: string;
     },
     output: {} as GenerateTextResult<any>,