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

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';
@@ -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();
 });
@@ -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) => {
@@ -339,7 +339,7 @@ test('You can listen for plan events', async () => {
     },
   });
 
-  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,12 +386,14 @@ 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', () => {
@@ -436,15 +438,15 @@ test('You can listen for message events', () => {
   );
 });
 
-test('agent.getPlans() returns plans from context', () => {
+test('agent.getDecisions() returns decisions from context', () => {
   const agent = createAgent({
     id: 'test',
     events: {},
     model: {} as any,
-    planner: async (agent) => {
+    strategy: async (agent) => {
       return {
         episodeId: agent.episodeId,
-        planner: 'test-planner',
+        strategy: 'test-strategy',
         goal: '',
         goalState: undefined,
         paths: [
@@ -459,10 +461,10 @@ test('agent.getPlans() returns plans from context', () => {
     },
   });
 
-  const plans = agent.getPlans();
+  const decisions = agent.getDecisions();
 
-  expect(plans).toBeDefined();
-  expect(Array.isArray(plans)).toBe(true);
+  expect(decisions).toBeDefined();
+  expect(Array.isArray(decisions)).toBe(true);
 });
 
 test('Event listeners can be unsubscribed', () => {

package/src/agent.ts

@@ -2,7 +2,6 @@ import {
   Actor,
   ActorRefLike,
   AnyEventObject,
-  AnyStateMachine,
   EventObject,
   fromTransition,
   Subscription,
@@ -11,7 +10,7 @@ import { ZodContextMapping, ZodEventMapping } from './schemas';
 import {
   AgentLogic,
   AgentMessage,
-  AgentPlanner,
+  AgentStrategy,
   EventsFromZodEventMapping,
   GenerateTextOptions,
   AgentLongTermMemory,
@@ -23,13 +22,15 @@ import {
   AgentFeedback,
   AgentMessageInput,
   AgentFeedbackInput,
-  AgentPlan,
+  AgentDecision,
   Compute,
   AgentDecisionInput,
   AgentDecideOptions,
   AnyAgent,
+  EventsFromAgent,
+  AgentInteractInput,
 } from './types';
-import { simplePlanner } from './planners/simple';
+import { simpleStrategy } from './strategies/simple';
 import { agentDecide } from './decide';
 import { getMachineHash, isActorRef, isMachineActor, randomId } from './utils';
 import {
@@ -69,12 +70,11 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
         });
         break;
       }
-      case 'agent.plan': {
-        state.plans.push(event.plan);
+      case 'agent.decision': {
+        state.decisions.push(event.decision);
         emit({
-          type: 'plan',
-          // @ts-ignore TODO: fix types in XState
-          plan: event.plan,
+          type: 'decision',
+          decision: event.decision,
         });
         break;
       }
@@ -91,7 +91,7 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
       feedback: [],
       messages: [],
       observations: [],
-      plans: [],
+      decisions: [],
     } as AgentMemoryContext)
 );
 
@@ -106,7 +106,8 @@ export function createAgent<
   model,
   events,
   context,
-  planner = simplePlanner as AgentPlanner<Agent<TContextSchema, TEventSchemas>>,
+  episodeId,
+  strategy = simpleStrategy,
   logic = agentLogic as AgentLogic<TEvents>,
 }: {
   /**
@@ -134,7 +135,7 @@ export function createAgent<
    */
   events: TEventSchemas;
   context?: TContextSchema;
-  planner?: AgentPlanner<Agent<TContextSchema, TEventSchemas>>;
+  strategy?: AgentStrategy<Agent<TContextSchema, TEventSchemas>>;
   stringify?: typeof JSON.stringify;
   /**
    * A function that retrieves the agent's long term memory
@@ -147,15 +148,17 @@ export function createAgent<
    */
   logic?: AgentLogic<TEvents>;
   model: LanguageModel;
-} & GenerateTextOptions): Agent<TContextSchema, TEventSchemas> {
+  episodeId?: string;
+}): Agent<TContextSchema, TEventSchemas> {
   return new Agent({
     id,
     context,
     events,
     description,
-    planner,
+    strategy: strategy,
     model,
     logic,
+    episodeId,
   }) as any;
 }
 
@@ -177,11 +180,11 @@ export class Agent<
   public description?: string;
   public events: TEventSchemas;
   public context?: TContextSchema;
-  public planner?: AgentPlanner<Agent<TContextSchema, TEventSchemas>>;
-  public types: {
-    events: TEvents;
-    context: Compute<TContext>;
-  };
+  public strategy: AgentStrategy<Agent<TContextSchema, TEventSchemas>>;
+  // public types: {
+  //   events: TEvents;
+  //   context: Compute<TContext>;
+  // };
   public model: LanguageModel;
   public memory: AgentLongTermMemory | undefined;
   public defaultOptions: AgentDecideOptions<AnyAgent> | undefined; // todo
@@ -194,7 +197,8 @@ export class Agent<
     model,
     events,
     context,
-    planner = simplePlanner,
+    episodeId,
+    strategy = simpleStrategy,
   }: {
     logic: AgentLogic<TEvents>;
     id?: string;
@@ -203,17 +207,18 @@ export class Agent<
     model: GenerateTextOptions['model'];
     events: TEventSchemas;
     context?: TContextSchema;
-    planner?: AgentPlanner<Agent<TContextSchema, TEventSchemas>>;
+    strategy?: AgentStrategy<Agent<TContextSchema, TEventSchemas>>;
+    episodeId?: string;
   }) {
     super(logic);
     this.model = model;
-    this.episodeId = id ?? randomId();
+    this.episodeId = episodeId ?? randomId('episode-');
     this.name = name;
     this.description = description;
     this.events = events;
     this.context = context;
-    this.planner = planner;
-    this.types = {} as any;
+    this.strategy = strategy;
+    this.id = id ?? randomId();
 
     this.start();
   }
@@ -225,6 +230,13 @@ export class Agent<
     return this.on('message', (ev) => fn(ev.message));
   }
 
+  /**
+   * Called whenever the agent (LLM assistant) receives some feedback.
+   */
+  public onFeedback(fn: (feedback: AgentFeedback) => void) {
+    return this.on('feedback', (ev) => fn(ev.feedback));
+  }
+
   /**
    * Retrieves messages from the agent's short-term (local) memory.
    */
@@ -251,7 +263,6 @@ export class Agent<
     const feedback = {
       ...feedbackInput,
       attributes: { ...feedbackInput.attributes },
-      reward: feedbackInput.reward ?? 0,
       timestamp: feedbackInput.timestamp ?? Date.now(),
       episodeId: this.episodeId,
     } satisfies AgentFeedback;
@@ -300,17 +311,17 @@ export class Agent<
     return this.getSnapshot().context.observations;
   }
 
-  public addPlan(plan: AgentPlan<TEvents>) {
+  public addDecision(decision: AgentDecision<TEvents>) {
     this.send({
-      type: 'agent.plan',
-      plan,
+      type: 'agent.decision',
+      decision,
     });
   }
   /**
    * Retrieves strategies from the agent's short-term (local) memory.
    */
-  public getPlans() {
-    return this.getSnapshot().context.plans;
+  public getDecisions() {
+    return this.getSnapshot().context.decisions;
   }
 
   /**
@@ -361,13 +372,13 @@ export class Agent<
     actorRef: TActor,
     getInput: (
       observation: AgentObservation<TActor>
-    ) => AgentDecisionInput | undefined
+    ) => AgentInteractInput<this> | void
   ): Subscription;
   public interact<TActor extends ActorRefLike>(
     actorRef: TActor,
     getInput?: (
       observation: AgentObservation<TActor>
-    ) => AgentDecisionInput | undefined
+    ) => AgentInteractInput<this> | void
   ): Subscription {
     const actorRefCheck = isActorRef(actorRef) && actorRef.src;
     const machine = isMachineActor(actorRef) ? actorRef.src : undefined;
@@ -425,7 +436,7 @@ export class Agent<
     if ((actorRef as any)._processingStatus === 1) {
       handleObservation({
         prevState: undefined,
-        event: { type: '' }, // TODO: unknown events?
+        event: undefined,
         state: actorRef.getSnapshot(),
         machine: (actorRef as any).src,
       });
@@ -478,14 +489,16 @@ export class Agent<
   }
 
   /**
-   * Resolves with an `AgentPlan` based on the information provided in the `options`, including:
+   * Resolves with an `AgentDecision` based on the information provided in the `options`, including:
    *
    * - The `goal` for the agent to achieve
    * - The observed current `state`
    * - The `machine` (e.g. a state machine) that specifies what can happen next
    * - Additional `context`
    */
-  public decide(opts: AgentDecideOptions<this>) {
+  public async decide(
+    opts: AgentDecideOptions<this>
+  ): Promise<AgentDecision<EventsFromAgent<this>> | undefined> {
     return agentDecide(this, opts);
   }
 }

package/src/decide.test.ts

@@ -154,7 +154,7 @@ test('interacts with an actor (late interaction)', async () => {
   expect(actor.getSnapshot().value).toBe('third');
 });
 
-test('agent.decide() makes a decision based on goal and state (simple planner)', async () => {
+test('agent.decide() makes a decision based on goal and state (simple strategy)', async () => {
   const model = new MockLanguageModelV1({
     doGenerate,
   });
@@ -167,7 +167,7 @@ test('agent.decide() makes a decision based on goal and state (simple planner)',
     },
   });
 
-  const plan = await agent.decide({
+  const decision = await agent.decide({
     goal: 'Make the best move',
     state: {
       value: 'playing',
@@ -188,8 +188,8 @@ test('agent.decide() makes a decision based on goal and state (simple planner)',
     }),
   });
 
-  expect(plan).toBeDefined();
-  expect(plan!.nextEvent).toEqual(
+  expect(decision).toBeDefined();
+  expect(decision!.nextEvent).toEqual(
     expect.objectContaining({
       type: 'MOVE',
     })
@@ -243,7 +243,7 @@ test.each([
       },
     });
 
-    const plan = await agent.decide({
+    const decision = await agent.decide({
       goal: 'Make the best move',
       state: {
         value: 'playing',
@@ -263,10 +263,10 @@ test.each([
     });
 
     if (!succeed) {
-      expect(plan).toBeUndefined();
+      expect(decision).toBeUndefined();
     } else {
-      expect(plan).toBeDefined();
-      expect(plan!.nextEvent).toEqual(
+      expect(decision).toBeDefined();
+      expect(decision!.nextEvent).toEqual(
         expect.objectContaining({
           type: 'MOVE',
         })
@@ -274,3 +274,51 @@ test.each([
     }
   }
 );
+
+test.each([['MOVE'], ['FORFEIT']] as const)(
+  'agent.decide() respects allowedEvents constraint (event: %s)',
+  async (allowedEventType) => {
+    const model = new MockLanguageModelV1({
+      doGenerate: async (params: LanguageModelV1CallOptions) => {
+        const keys =
+          params.mode.type === 'regular'
+            ? params.mode.tools?.map((t) => t.name)
+            : [];
+
+        return {
+          ...dummyResponseValues,
+          finishReason: 'tool-calls',
+          toolCalls: [
+            {
+              toolCallType: 'function',
+              toolCallId: 'call-1',
+              toolName: keys![0],
+              args: `{ "type": "${keys?.[0]}" }`,
+            },
+          ],
+        } as any;
+      },
+    });
+
+    const agent = createAgent({
+      id: 'test',
+      model,
+      events: {
+        MOVE: z.object({}),
+        SKIP: z.object({}),
+        FORFEIT: z.object({}),
+      },
+    });
+
+    const decision = await agent.decide({
+      goal: 'Make the best move',
+      state: {
+        value: 'playing',
+        context: {},
+      },
+      allowedEvents: [allowedEventType],
+    });
+
+    expect(decision?.nextEvent?.type).toEqual(allowedEventType);
+  }
+);

package/src/decide.ts

@@ -3,65 +3,72 @@ import {
   AnyAgent,
   AgentDecideOptions,
   AgentDecisionLogic,
-  AgentDecisionInput,
-  AgentPlanner,
-  AgentPlan,
-  EventsFromZodEventMapping,
-  AgentPlanInput,
+  AgentDecision,
+  AgentDecideInput,
   TransitionData,
+  EventsFromAgent,
 } from './types';
-import { simplePlanner } from './planners/simple';
 import { getTransitions } from './utils';
 import { CoreMessage, CoreTool, tool } from 'ai';
 
 export async function agentDecide<T extends AnyAgent>(
   agent: T,
   options: AgentDecideOptions<T>
-): Promise<AgentPlan<EventsFromZodEventMapping<T['events']>> | undefined> {
+): Promise<AgentDecision<EventsFromAgent<T>> | undefined> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
   };
   const {
-    planner = simplePlanner as AgentPlanner<any>,
+    strategy = agent.strategy,
     goal,
+    allowedEvents,
     events = agent.events,
     state,
     machine,
     model = agent.model,
     messages,
-    ...otherPlanInput
+    ...otherDecideInput
   } = resolvedOptions;
 
+  const filteredEventSchemas = allowedEvents
+    ? Object.fromEntries(
+        Object.entries(events).filter(([key]) => {
+          return allowedEvents.includes(key);
+        })
+      )
+    : events;
+
   let attempts = 0;
 
   const maxAttempts = resolvedOptions.maxAttempts ?? 2;
 
-  let plan;
+  let decision: AgentDecision<any> | undefined;
 
   while (attempts++ < maxAttempts) {
-    plan = await planner(agent, {
+    decision = await strategy(agent, {
       model,
       goal,
-      events,
+      events: filteredEventSchemas,
       state,
       machine,
       messages: messages as CoreMessage[], // TODO: fix UIMessage thing
-      ...otherPlanInput,
+      ...otherDecideInput,
     });
 
-    if (plan?.nextEvent) {
-      agent.addPlan(plan);
-      await resolvedOptions.execute?.(plan.nextEvent);
+    if (decision?.nextEvent) {
+      agent.addDecision(decision);
+      await resolvedOptions.execute?.(decision.nextEvent);
+      break;
     }
   }
 
-  return plan;
+  return decision;
 }
 
-export function fromDecision(
-  agent: AnyAgent,
-  defaultInput?: AgentDecisionInput
+export function fromDecision<T extends AnyAgent>(
+  agent: T,
+  defaultInput?: AgentDecideInput<EventsFromAgent<T>>
 ): AgentDecisionLogic<any> {
   return fromPromise(async ({ input, self }) => {
     const parentRef = self._parent;
@@ -80,22 +87,25 @@ export function fromDecision(
       context: resolvedInput.context,
     };
 
-    const plan = await agentDecide(agent, {
+    const decision = await agentDecide(agent, {
       machine: (parentRef as AnyActor).logic,
-      state,
+      state: snapshot,
+      context: resolvedInput.context,
       execute: async (event) => {
         parentRef.send(event);
       },
       ...resolvedInput,
+      // @ts-ignore
+      messages: resolvedInput.messages,
     });
 
-    return plan;
+    return decision;
   }) as AgentDecisionLogic<any>;
 }
 
 export function getToolMap<T extends AnyAgent>(
   _agent: T,
-  input: AgentPlanInput<any>
+  input: AgentDecideInput<any>
 ): Record<string, CoreTool<any, any>> | undefined {
   // Get all of the possible next transitions
   const transitions: TransitionData[] = input.machine

package/src/strategies/chainOfThought.ts

@@ -0,0 +1,48 @@
+import { generateText } from 'ai';
+import {
+  AnyAgent,
+  AgentDecideInput,
+  AgentDecision,
+  PromptTemplate,
+} from '../types';
+import { getMessages } from '../text';
+import { simpleStrategy } from './simple';
+import { convertToXml } from '../utils';
+
+const chainOfThoughtPromptTemplate: PromptTemplate<any> = ({
+  context,
+  goal,
+}) => {
+  return `
+${convertToXml({ context, goal })}
+
+How would you achieve the goal? Think step-by-step.
+`.trim();
+};
+
+export async function chainOfThoughtStrategy<T extends AnyAgent>(
+  agent: T,
+  input: AgentDecideInput<any>
+): Promise<AgentDecision<any> | undefined> {
+  const prompt = chainOfThoughtPromptTemplate({
+    context: input.state.context,
+    goal: input.goal,
+  });
+
+  const messages = await getMessages(agent, prompt, input);
+
+  const model = input.model ? agent.wrap(input.model) : agent.model;
+
+  const result = await generateText({
+    model,
+    system: input.system ?? agent.description,
+    messages,
+  });
+
+  const decision = await simpleStrategy(agent, {
+    ...input,
+    messages: messages.concat(result.response.messages),
+  });
+
+  return decision;
+}

package/src/{planners → strategies}/shortestPath.test.ts

@@ -1,7 +1,7 @@
-import { createAgent } from '../';
+import { createAgent, TypesFromAgent } from '..';
 import { assign, createActor, setup } from 'xstate';
 import { z } from 'zod';
-import { experimental_createShortestPathPlanner } from './shortestPath';
+import { experimental_shortestPathStrategy } from './shortestPath';
 import { test, expect } from 'vitest';
 import { dummyResponseValues, MockLanguageModelV1 } from '../mockModel';
 
@@ -35,10 +35,7 @@ test.skip('should find shortest path to goal', async () => {
   });
 
   const counterMachine = setup({
-    types: {
-      context: agent.types.context,
-      events: agent.types.events,
-    },
+    types: {} as TypesFromAgent<typeof agent>,
   }).createMachine({
     initial: 'counting',
     context: { count: 0 },
@@ -87,7 +84,7 @@ test.skip('should find shortest path to goal', async () => {
     }),
     goal: 'Get the counter to exactly 3',
     state: counterActor.getSnapshot(),
-    planner: experimental_createShortestPathPlanner(),
+    strategy: experimental_shortestPathStrategy,
   });
 
   expect(decision?.nextEvent?.type).toBe('increment');