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

package/src/agent.test.ts

@@ -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(),
   });
@@ -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,
   });
@@ -332,7 +332,7 @@ test('You can listen for plan events', async () => {
   });
 
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     model,
     events: {
       WIN: z.object({}),
@@ -393,3 +393,143 @@ test('agent.types provides context and event types', () => {
   // @ts-expect-error
   agent.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.getPlans() returns plans from context', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+    planner: async (agent) => {
+      return {
+        episodeId: agent.episodeId,
+        planner: 'test-planner',
+        goal: '',
+        goalState: undefined,
+        paths: [
+          {
+            state: undefined,
+            steps: [],
+          },
+        ],
+        nextEvent: undefined,
+        timestamp: Date.now(),
+      };
+    },
+  });
+
+  const plans = agent.getPlans();
+
+  expect(plans).toBeDefined();
+  expect(Array.isArray(plans)).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();
+});

package/src/agent.ts

@@ -27,10 +27,11 @@ import {
   Compute,
   AgentDecisionInput,
   AgentDecideOptions,
+  AnyAgent,
 } from './types';
-import { simplePlanner } from './planners/simplePlanner';
+import { simplePlanner } from './planners/simple';
 import { agentDecide } from './decide';
-import { getMachineHash, isActorRef, randomId } from './utils';
+import { getMachineHash, isActorRef, isMachineActor, randomId } from './utils';
 import {
   experimental_wrapLanguageModel,
   LanguageModel,
@@ -77,8 +78,11 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
         });
         break;
       }
-      default:
+      default: {
+        // unrecognized
+        console.warn('Unrecognized event', event);
         break;
+      }
     }
     return state;
   },
@@ -98,22 +102,18 @@ export function createAgent<
   TContext = ContextFromZodContextMapping<TContextSchema>
 >({
   id,
-  name,
-  description,
+  description: description,
   model,
   events,
   context,
   planner = simplePlanner as AgentPlanner<Agent<TContextSchema, TEventSchemas>>,
-  stringify = JSON.stringify,
-  getMemory,
   logic = agentLogic as AgentLogic<TEvents>,
-  ...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.
+   * used to retrieve memory for previous episodes of this agent.
    *
    * @example
    * ```ts
@@ -124,10 +124,6 @@ export function createAgent<
    * ```
    */
   id?: string;
-  /**
-   * The name of the agent
-   */
-  name?: string;
   /**
    * A description of the role of the agent
    */
@@ -150,210 +146,17 @@ export function createAgent<
    * Agent logic
    */
   logic?: AgentLogic<TEvents>;
+  model: LanguageModel;
 } & GenerateTextOptions): Agent<TContextSchema, TEventSchemas> {
   return new Agent({
     id,
     context,
     events,
-    name,
     description,
     planner,
     model,
     logic,
   }) as any;
-  // const agent = createActor(logic) as unknown as Agent<TContext, TEvents>;
-  // agent.events = events;
-  // agent.model = model;
-  // agent.name = name;
-  // agent.description = description;
-  // agent.defaultOptions = { ...generateTextOptions, model };
-  // agent.memory = getMemory ? getMemory(agent) : undefined;
-
-  // agent.onMessage = (callback) => {
-  //   agent.on('message', (ev) => callback(ev.message));
-  // };
-
-  // agent.decide = (opts) => {
-  //   return agentDecide(agent, opts);
-  // };
-
-  // agent.addMessage = (messageInput) => {
-  //   const message = {
-  //     ...messageInput,
-  //     id: messageInput.id ?? randomId(),
-  //     timestamp: messageInput.timestamp ?? Date.now(),
-  //     sessionId: agent.sessionId,
-  //   } satisfies AgentMessage;
-  //   agent.send({
-  //     type: 'agent.message',
-  //     message,
-  //   });
-
-  //   return message;
-  // };
-  // agent.getMessages = () => agent.getSnapshot().context.messages;
-
-  // agent.addFeedback = (feedbackInput) => {
-  //   const feedback = {
-  //     ...feedbackInput,
-  //     attributes: { ...feedbackInput.attributes },
-  //     reward: feedbackInput.reward ?? 0,
-  //     timestamp: feedbackInput.timestamp ?? Date.now(),
-  //     sessionId: agent.sessionId,
-  //   } satisfies AgentFeedback;
-  //   agent.send({
-  //     type: 'agent.feedback',
-  //     feedback,
-  //   });
-  //   return feedback;
-  // };
-  // agent.getFeedback = () => agent.getSnapshot().context.feedback;
-
-  // agent.addObservation = (observationInput) => {
-  //   const { prevState, event, state } = observationInput;
-  //   const observedState = { context: state.context, value: state.value };
-  //   const observedPrevState = prevState
-  //     ? {
-  //         context: prevState.context,
-  //         value: prevState.value,
-  //       }
-  //     : undefined;
-  //   const observation = {
-  //     prevState: observedPrevState,
-  //     event,
-  //     state: observedState,
-  //     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',
-  //     observation,
-  //   });
-
-  //   return observation;
-  // };
-  // agent.getObservations = () => agent.getSnapshot().context.observations;
-
-  // agent.addPlan = (plan) => {
-  //   agent.send({
-  //     type: 'agent.plan',
-  //     plan,
-  //   });
-  // };
-  // agent.getPlans = () => agent.getSnapshot().context.plans;
-
-  // agent.interact = ((actorRef, getInput) => {
-  //   let prevState: ObservedState | undefined = undefined;
-  //   let subscribed = true;
-
-  //   async function handleObservation(observationInput: AgentObservationInput) {
-  //     const observation = agent.addObservation(observationInput);
-
-  //     const input = getInput?.(observation);
-
-  //     if (input) {
-  //       await agentDecide(agent, {
-  //         machine: actorRef.src as AnyStateMachine,
-  //         state: observation.state,
-  //         execute: async (event) => {
-  //           actorRef.send(event);
-  //         },
-  //         ...input,
-  //       });
-  //     }
-
-  //     prevState = observationInput.state;
-  //   }
-
-  //   // Inspect system, but only observe specified actor
-  //   const sub = actorRef.system.inspect({
-  //     next: async (inspEvent) => {
-  //       if (
-  //         !subscribed ||
-  //         inspEvent.actorRef !== actorRef ||
-  //         inspEvent.type !== '@xstate.snapshot'
-  //       ) {
-  //         return;
-  //       }
-
-  //       const observationInput = {
-  //         event: inspEvent.event,
-  //         prevState,
-  //         state: inspEvent.snapshot as any,
-  //         machine: (actorRef as any).src,
-  //       } satisfies AgentObservationInput;
-
-  //       await handleObservation(observationInput);
-  //     },
-  //   });
-
-  //   // If actor already started, interact with current state
-  //   if ((actorRef as any)._processingStatus === 1) {
-  //     handleObservation({
-  //       prevState: undefined,
-  //       event: { type: '' }, // TODO: unknown events?
-  //       state: actorRef.getSnapshot(),
-  //       machine: (actorRef as any).src,
-  //     });
-  //   }
-
-  //   return {
-  //     unsubscribe: () => {
-  //       sub.unsubscribe();
-  //       subscribed = false;
-  //     },
-  //   };
-  // }) as typeof agent.interact;
-
-  // agent.observe = (actorRef) => {
-  //   let prevState: ObservedState = actorRef.getSnapshot();
-
-  //   const sub = actorRef.system.inspect({
-  //     next: async (inspEvent) => {
-  //       if (
-  //         inspEvent.actorRef !== actorRef ||
-  //         inspEvent.type !== '@xstate.snapshot'
-  //       ) {
-  //         return;
-  //       }
-
-  //       const observationInput = {
-  //         event: inspEvent.event,
-  //         prevState,
-  //         state: inspEvent.snapshot as any,
-  //         machine: (actorRef as any).src,
-  //       } satisfies AgentObservationInput;
-
-  //       prevState = observationInput.state;
-
-  //       agent.addObservation(observationInput);
-  //     },
-  //   });
-
-  //   return sub;
-  // };
-
-  // agent.types = {} as any;
-
-  // agent.wrap = (modelToWrap) =>
-  //   experimental_wrapLanguageModel({
-  //     model: modelToWrap,
-  //     middleware: createAgentMiddleware(agent),
-  //   });
-
-  // agent.model = experimental_wrapLanguageModel({
-  //   model,
-  //   middleware: createAgentMiddleware(agent),
-  // });
-
-  // agent.start();
-
-  // return agent;
 }
 
 export class Agent<
@@ -381,7 +184,7 @@ export class Agent<
   };
   public model: LanguageModel;
   public memory: AgentLongTermMemory | undefined;
-  public defaultOptions: any; // todo
+  public defaultOptions: AgentDecideOptions<AnyAgent> | undefined; // todo
 
   constructor({
     logic = agentLogic as AgentLogic<TEvents>,
@@ -566,7 +369,8 @@ export class Agent<
       observation: AgentObservation<TActor>
     ) => AgentDecisionInput | undefined
   ): Subscription {
-    const actorRefCheck = isActorRef(actorRef);
+    const actorRefCheck = isActorRef(actorRef) && actorRef.src;
+    const machine = isMachineActor(actorRef) ? actorRef.src : undefined;
 
     let prevState: ObservedState | undefined = undefined;
     let subscribed = true;
@@ -579,16 +383,15 @@ export class Agent<
       const input = getInput?.(observation);
 
       if (input) {
-        await agentDecide(agent, {
-          machine: actorRefCheck
-            ? (actorRef.src as AnyStateMachine)
-            : undefined,
+        const res = await agentDecide(agent, {
+          machine,
           state: observation.state,
-          execute: async (event) => {
-            actorRef.send(event);
-          },
           ...input,
         });
+
+        if (res?.nextEvent) {
+          actorRef.send(res.nextEvent);
+        }
       }
 
       prevState = observationInput.state;
@@ -682,7 +485,7 @@ export class Agent<
    * - The `machine` (e.g. a state machine) that specifies what can happen next
    * - Additional `context`
    */
-  public decide(opts: AgentDecideOptions) {
+  public decide(opts: AgentDecideOptions<this>) {
     return agentDecide(this, opts);
   }
 }

package/src/decide.test.ts

@@ -28,7 +28,7 @@ test('fromDecision() makes a decision', async () => {
     doGenerate,
   });
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     model,
     events: {
       doFirst: z.object({}),
@@ -73,7 +73,7 @@ test('interacts with an actor', async () => {
     doGenerate,
   });
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     model,
     events: {
       doFirst: z.object({}),
@@ -116,7 +116,7 @@ test('interacts with an actor (late interaction)', async () => {
     doGenerate,
   });
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     model,
     events: {
       doFirst: z.object({}),
@@ -153,3 +153,124 @@ 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 () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
+
+  const agent = createAgent({
+    id: 'test',
+    model,
+    events: {
+      MOVE: z.object({}),
+    },
+  });
+
+  const plan = await agent.decide({
+    goal: 'Make the best move',
+    state: {
+      value: 'playing',
+      context: {
+        board: [0, 0, 0],
+      },
+    },
+    machine: createMachine({
+      initial: 'playing',
+      states: {
+        playing: {
+          on: {
+            MOVE: 'next',
+          },
+        },
+        next: {},
+      },
+    }),
+  });
+
+  expect(plan).toBeDefined();
+  expect(plan!.nextEvent).toEqual(
+    expect.objectContaining({
+      type: 'MOVE',
+    })
+  );
+});
+
+test.each([
+  [undefined, true],
+  [undefined, false],
+  [3, true],
+  [3, false],
+])(
+  'agent.decide() retries if a decision is not made (%i attempts, succeed: %s)',
+  async (maxAttempts, succeed) => {
+    let attempts = 0;
+    const doGenerateWithRetry = async (params: LanguageModelV1CallOptions) => {
+      const keys =
+        params.mode.type === 'regular'
+          ? params.mode.tools?.map((t) => t.name)
+          : [];
+
+      console.log('try', attempts, 'max', maxAttempts);
+
+      const toolCalls =
+        succeed && attempts++ === (maxAttempts ?? 2) - 1
+          ? [
+              {
+                toolCallType: 'function',
+                toolCallId: 'call-1',
+                toolName: keys![0],
+                args: `{ "type": "${keys?.[0]}" }`,
+              },
+            ]
+          : [];
+
+      return {
+        ...dummyResponseValues,
+        finishReason: 'tool-calls',
+        toolCalls,
+      } as any;
+    };
+    const model = new MockLanguageModelV1({
+      doGenerate: doGenerateWithRetry,
+    });
+
+    const agent = createAgent({
+      id: 'test',
+      model,
+      events: {
+        MOVE: z.object({}),
+      },
+    });
+
+    const plan = await agent.decide({
+      goal: 'Make the best move',
+      state: {
+        value: 'playing',
+      },
+      machine: createMachine({
+        initial: 'playing',
+        states: {
+          playing: {
+            on: {
+              MOVE: 'win',
+            },
+          },
+          win: {},
+        },
+      }),
+      maxAttempts,
+    });
+
+    if (!succeed) {
+      expect(plan).toBeUndefined();
+    } else {
+      expect(plan).toBeDefined();
+      expect(plan!.nextEvent).toEqual(
+        expect.objectContaining({
+          type: 'MOVE',
+        })
+      );
+    }
+  }
+);