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

package/src/agent.test.ts

@@ -1,5 +1,5 @@
 import { test, expect, vi } from 'vitest';
-import { createAgent, TypesFromAgent } from './';
+import { AgentDecision, createAgent, TypesFromAgent } from './';
 import { createActor, createMachine } from 'xstate';
 import { LanguageModelV1CallOptions } from 'ai';
 import { z } from 'zod';
@@ -66,38 +66,43 @@ test('agent.addMessage() adds to message history', () => {
 test('agent.addFeedback() adds to feedback', () => {
   const agent = createAgent({
     id: 'test',
-    events: {},
+    events: {
+      play: z.object({
+        position: z.number(),
+      }),
+    },
     model: {} as any,
   });
 
-  const feedback = agent.addFeedback({
-    attributes: {
-      score: -1,
-    },
+  const decision: AgentDecision<typeof agent> = {
     goal: 'Win the game',
-    observationId: 'obs-1',
+    episodeId: agent.episodeId,
+    goalState: { value: 'won' },
+    id: 'decision-1',
+    nextEvent: { type: 'play', position: 3 },
+    paths: [],
+    strategy: 'simple',
+    timestamp: Date.now(),
+  };
+
+  const obs = agent.addObservation({
+    decisionId: decision.id,
+    prevState: { value: 'playing' },
+    event: { type: 'play', position: 3 },
+    state: { value: 'lost' },
+  });
+
+  const feedback = agent.addFeedback({
+    score: 0,
+    observationId: obs.id,
   });
 
   expect(feedback.episodeId).toEqual(agent.episodeId);
 
   expect(agent.getFeedback()).toContainEqual(
     expect.objectContaining({
-      attributes: {
-        score: -1,
-      },
-      goal: 'Win the game',
-      observationId: 'obs-1',
-      episodeId: expect.any(String),
-      timestamp: expect.any(Number),
-    })
-  );
-  expect(agent.getFeedback()).toContainEqual(
-    expect.objectContaining({
-      attributes: {
-        score: -1,
-      },
-      goal: 'Win the game',
-      observationId: 'obs-1',
+      score: 0,
+      observationId: obs.id,
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
@@ -115,6 +120,7 @@ test('agent.addObservation() adds to observations', () => {
     prevState: { value: 'playing', context: {} },
     event: { type: 'play', position: 3 },
     state: { value: 'lost', context: {} },
+    goal: 'Win the game',
   });
 
   expect(observation.episodeId).toEqual(agent.episodeId);
@@ -139,6 +145,7 @@ test('agent.addObservation() adds to observations (initial state)', () => {
 
   const observation = agent.addObservation({
     state: { value: 'lost' },
+    goal: 'Win the game',
   });
 
   expect(observation.episodeId).toEqual(agent.episodeId);
@@ -175,7 +182,7 @@ test('agent.addObservation() adds to observations with machine hash', () => {
     prevState: { value: 'playing', context: {} },
     event: { type: 'play', position: 3 },
     state: { value: 'lost', context: {} },
-    machine,
+    goal: 'Win the game',
   });
 
   expect(observation.episodeId).toEqual(agent.episodeId);
@@ -185,7 +192,6 @@ test('agent.addObservation() adds to observations with machine hash', () => {
       prevState: { value: 'playing', context: {} },
       event: { type: 'play', position: 3 },
       state: { value: 'lost', context: {} },
-      machineHash: expect.any(String),
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
@@ -203,13 +209,11 @@ test('agent.addFeedback() adds to feedback (with observation)', () => {
     state: {
       value: 'playing',
     },
+    goal: 'Win the game',
   });
 
   const feedback = agent.addFeedback({
-    attributes: {
-      score: -1,
-    },
-    goal: 'Win the game',
+    score: 0,
     observationId: observation.id,
   });
 
@@ -217,10 +221,7 @@ test('agent.addFeedback() adds to feedback (with observation)', () => {
 
   expect(agent.getFeedback()).toContainEqual(
     expect.objectContaining({
-      attributes: {
-        score: -1,
-      },
-      goal: 'Win the game',
+      score: 0,
       observationId: observation.id,
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
@@ -228,10 +229,7 @@ test('agent.addFeedback() adds to feedback (with observation)', () => {
   );
   expect(agent.getFeedback()).toContainEqual(
     expect.objectContaining({
-      attributes: {
-        score: -1,
-      },
-      goal: 'Win the game',
+      score: 0,
       observationId: observation.id,
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
@@ -297,10 +295,7 @@ test('You can listen for feedback events', () => {
   agent.on('feedback', fn);
 
   agent.addFeedback({
-    attributes: {
-      score: -1,
-    },
-    goal: 'Win the game',
+    score: -1,
     observationId: 'obs-1',
   });
 
@@ -445,6 +440,7 @@ test('agent.getDecisions() returns decisions from context', () => {
     model: {} as any,
     strategy: async (agent) => {
       return {
+        id: Date.now().toString(),
         episodeId: agent.episodeId,
         strategy: 'test-strategy',
         goal: '',
@@ -520,7 +516,6 @@ test('agent.observe() adds observations from actor snapshots', () => {
   expect(agent.getObservations()).toContainEqual(
     expect.objectContaining({
       state: expect.objectContaining({ value: 'idle' }),
-      machineHash: expect.any(String),
     })
   );
 
@@ -529,9 +524,95 @@ test('agent.observe() adds observations from actor snapshots', () => {
       prevState: expect.objectContaining({ value: 'idle' }),
       event: { type: 'START' },
       state: expect.objectContaining({ value: 'running' }),
-      machineHash: expect.any(String),
     })
   );
 
   subscription.unsubscribe();
 });
+
+test('agent.addObservation() accepts custom episodeId', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const customEpisodeId = 'custom-episode-123';
+  const observation = agent.addObservation({
+    state: { value: 'playing' },
+    goal: 'Win the game',
+    episodeId: customEpisodeId,
+  });
+
+  expect(observation.episodeId).toEqual(customEpisodeId);
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      episodeId: customEpisodeId,
+    })
+  );
+});
+
+test('agent.addFeedback() accepts custom episodeId', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const customEpisodeId = 'custom-episode-123';
+  const feedback = agent.addFeedback({
+    score: 1,
+    observationId: 'obs-1',
+    episodeId: customEpisodeId,
+  });
+
+  expect(feedback.episodeId).toEqual(customEpisodeId);
+  expect(agent.getFeedback()).toContainEqual(
+    expect.objectContaining({
+      episodeId: customEpisodeId,
+    })
+  );
+});
+
+test('agent.addObservation() accepts decisionId', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const decisionId = 'decision-123';
+  const observation = agent.addObservation({
+    state: { value: 'playing' },
+    goal: 'Win the game',
+    decisionId,
+  });
+
+  expect(observation.decisionId).toEqual(decisionId);
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      decisionId,
+    })
+  );
+});
+
+test('agent.addFeedback() accepts decisionId', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const decisionId = 'decision-123';
+  const feedback = agent.addFeedback({
+    score: 1,
+    decisionId,
+  });
+
+  expect(feedback.decisionId).toEqual(decisionId);
+  expect(agent.getFeedback()).toContainEqual(
+    expect.objectContaining({
+      decisionId,
+    })
+  );
+});

package/src/agent.ts

@@ -1,7 +1,6 @@
 import {
   Actor,
   ActorRefLike,
-  AnyEventObject,
   EventObject,
   fromTransition,
   Subscription,
@@ -23,16 +22,13 @@ import {
   AgentMessageInput,
   AgentFeedbackInput,
   AgentDecision,
-  Compute,
-  AgentDecisionInput,
-  AgentDecideOptions,
   AnyAgent,
-  EventsFromAgent,
   AgentInteractInput,
+  AgentDecideInput,
 } from './types';
-import { simpleStrategy } from './strategies/simple';
+import { simpleStrategy } from './strategies/simpleStrategy';
 import { agentDecide } from './decide';
-import { getMachineHash, isActorRef, isMachineActor, randomId } from './utils';
+import { isActorRef, isMachineActor, randomId } from './utils';
 import {
   experimental_wrapLanguageModel,
   LanguageModel,
@@ -40,7 +36,7 @@ import {
 } from 'ai';
 import { createAgentMiddleware } from './middleware';
 
-export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
+export const agentLogic: AgentLogic<any> = fromTransition(
   (state, event, { emit }) => {
     switch (event.type) {
       case 'agent.feedback': {
@@ -92,14 +88,15 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
       messages: [],
       observations: [],
       decisions: [],
-    } as AgentMemoryContext)
+    } as AgentMemoryContext<any>)
 );
 
 export function createAgent<
   const TContextSchema extends ZodContextMapping,
   const TEventSchemas extends ZodEventMapping,
   TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
-  TContext = ContextFromZodContextMapping<TContextSchema>
+  TContext = ContextFromZodContextMapping<TContextSchema>,
+  TAgent extends AnyAgent = Agent<TContextSchema, TEventSchemas>
 >({
   id,
   description: description,
@@ -108,7 +105,7 @@ export function createAgent<
   context,
   episodeId,
   strategy = simpleStrategy,
-  logic = agentLogic as AgentLogic<TEvents>,
+  logic = agentLogic as AgentLogic<any>,
 }: {
   /**
    * The unique identifier for the agent.
@@ -142,11 +139,11 @@ export function createAgent<
    */
   getMemory?: (
     agent: Agent<TContextSchema, TEventSchemas>
-  ) => AgentLongTermMemory;
+  ) => AgentLongTermMemory<TAgent>;
   /**
    * Agent logic
    */
-  logic?: AgentLogic<TEvents>;
+  logic?: AgentLogic<TAgent>;
   model: LanguageModel;
   episodeId?: string;
 }): Agent<TContextSchema, TEventSchemas> {
@@ -167,7 +164,7 @@ export class Agent<
   const TEventSchemas extends ZodEventMapping,
   TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
   TContext = ContextFromZodContextMapping<TContextSchema>
-> extends Actor<AgentLogic<TEvents>> {
+> extends Actor<AgentLogic<any>> {
   /**
    * The name of the agent. All agents with the same name are related and
    * able to share experiences (observations, feedback) with each other.
@@ -186,11 +183,10 @@ export class Agent<
   //   context: Compute<TContext>;
   // };
   public model: LanguageModel;
-  public memory: AgentLongTermMemory | undefined;
-  public defaultOptions: AgentDecideOptions<AnyAgent> | undefined; // todo
+  public memory: AgentLongTermMemory<this> | undefined;
 
   constructor({
-    logic = agentLogic as AgentLogic<TEvents>,
+    logic = agentLogic as AgentLogic<any>,
     id,
     name,
     description,
@@ -200,7 +196,7 @@ export class Agent<
     episodeId,
     strategy = simpleStrategy,
   }: {
-    logic: AgentLogic<TEvents>;
+    logic: AgentLogic<any>;
     id?: string;
     name?: string;
     description?: string;
@@ -262,9 +258,10 @@ export class Agent<
   public addFeedback(feedbackInput: AgentFeedbackInput) {
     const feedback = {
       ...feedbackInput,
+      comment: feedbackInput.comment ?? undefined,
       attributes: { ...feedbackInput.attributes },
       timestamp: feedbackInput.timestamp ?? Date.now(),
-      episodeId: this.episodeId,
+      episodeId: feedbackInput.episodeId ?? this.episodeId,
     } satisfies AgentFeedback;
     this.send({
       type: 'agent.feedback',
@@ -281,7 +278,7 @@ export class Agent<
   }
 
   public addObservation(
-    observationInput: AgentObservationInput
+    observationInput: AgentObservationInput<this>
   ): AgentObservation<any> {
     const { prevState, event, state } = observationInput;
     const observation = {
@@ -289,11 +286,12 @@ export class Agent<
       event,
       state,
       id: observationInput.id ?? randomId(),
-      episodeId: this.episodeId,
+      episodeId: observationInput.episodeId ?? this.episodeId,
       timestamp: observationInput.timestamp ?? Date.now(),
-      machineHash: observationInput.machine
-        ? getMachineHash(observationInput.machine)
-        : undefined,
+      decisionId: observationInput.decisionId,
+      // machineHash: observationInput.machine
+      //   ? getMachineHash(observationInput.machine)
+      //   : undefined,
     } satisfies AgentObservation<any>;
 
     this.send({
@@ -311,7 +309,7 @@ export class Agent<
     return this.getSnapshot().context.observations;
   }
 
-  public addDecision(decision: AgentDecision<TEvents>) {
+  public addDecision(decision: AgentDecision<this>) {
     this.send({
       type: 'agent.decision',
       decision,
@@ -383,25 +381,29 @@ export class Agent<
     const actorRefCheck = isActorRef(actorRef) && actorRef.src;
     const machine = isMachineActor(actorRef) ? actorRef.src : undefined;
 
-    let prevState: ObservedState | undefined = undefined;
+    let prevState: ObservedState<this> | undefined = undefined;
     let subscribed = true;
 
     const agent = this;
 
-    async function handleObservation(observationInput: AgentObservationInput) {
+    async function handleObservation(
+      observationInput: AgentObservationInput<any>
+    ) {
       const observation = agent.addObservation(observationInput);
 
-      const input = getInput?.(observation);
+      const interactInput = getInput?.(observation);
 
-      if (input) {
-        const res = await agentDecide(agent, {
+      if (interactInput) {
+        const decision = await agentDecide(agent, {
           machine,
           state: observation.state,
-          ...input,
+          ...interactInput,
         });
 
-        if (res?.nextEvent) {
-          actorRef.send(res.nextEvent);
+        if (decision?.nextEvent) {
+          // @ts-ignore
+          decision.nextEvent['_decision'] = decision.id;
+          actorRef.send(decision.nextEvent);
         }
       }
 
@@ -420,12 +422,21 @@ export class Agent<
               return;
             }
 
+            const decisionId = inspEvent.event['_decision'] as
+              | string
+              | undefined;
+
+            const decision = decisionId
+              ? agent.getDecisions().find((d) => d.id === decisionId)
+              : undefined;
+
             const observationInput = {
               event: inspEvent.event,
               prevState,
               state: inspEvent.snapshot as any,
-              machine: (actorRef as any).src,
-            } satisfies AgentObservationInput;
+              goal: decision?.goal,
+              decisionId,
+            } satisfies AgentObservationInput<any>;
 
             await handleObservation(observationInput);
           },
@@ -435,10 +446,10 @@ export class Agent<
     // If actor already started, interact with current state
     if ((actorRef as any)._processingStatus === 1) {
       handleObservation({
+        decisionId: undefined,
         prevState: undefined,
         event: undefined,
         state: actorRef.getSnapshot(),
-        machine: (actorRef as any).src,
       });
     }
 
@@ -451,7 +462,7 @@ export class Agent<
   }
 
   public observe<TActor extends ActorRefLike>(actorRef: TActor): Subscription {
-    let prevState: ObservedState = actorRef.getSnapshot();
+    let prevState: ObservedState<this> = actorRef.getSnapshot();
     const actorRefCheck = isActorRef(actorRef);
 
     const sub = actorRefCheck
@@ -464,12 +475,20 @@ export class Agent<
               return;
             }
 
+            const decisionId = inspEvent.event['_decision'] as
+              | string
+              | undefined;
+            const decision = decisionId
+              ? this.getDecisions().find((d) => d.id === decisionId)
+              : undefined;
+
             const observationInput = {
+              decisionId,
               event: inspEvent.event,
               prevState,
               state: inspEvent.snapshot as any,
-              machine: (actorRef as any).src,
-            } satisfies AgentObservationInput;
+              goal: decision?.goal,
+            } satisfies AgentObservationInput<this>;
 
             prevState = observationInput.state;
 
@@ -497,8 +516,8 @@ export class Agent<
    * - Additional `context`
    */
   public async decide(
-    opts: AgentDecideOptions<this>
-  ): Promise<AgentDecision<EventsFromAgent<this>> | undefined> {
-    return agentDecide(this, opts);
+    input: AgentDecideInput<this>
+  ): Promise<AgentDecision<this> | undefined> {
+    return agentDecide(this, input);
   }
 }

package/src/decide.test.ts

@@ -322,3 +322,25 @@ test.each([['MOVE'], ['FORFEIT']] as const)(
     expect(decision?.nextEvent?.type).toEqual(allowedEventType);
   }
 );
+
+test('agent.decide() accepts custom episodeId', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
+  const agent = createAgent({
+    id: 'test',
+    events: {
+      WIN: z.object({}),
+    },
+    model,
+  });
+
+  const customEpisodeId = 'custom-episode-123';
+  const decision = await agent.decide({
+    goal: 'Win the game',
+    state: { value: 'playing' },
+    episodeId: customEpisodeId,
+  });
+
+  expect(decision?.episodeId).toEqual(customEpisodeId);
+});

package/src/decide.ts

@@ -1,24 +1,21 @@
 import { AnyActor, AnyMachineSnapshot, fromPromise } from 'xstate';
 import {
   AnyAgent,
-  AgentDecideOptions,
   AgentDecisionLogic,
   AgentDecision,
   AgentDecideInput,
   TransitionData,
-  EventsFromAgent,
+  EventFromAgent,
 } from './types';
 import { getTransitions } from './utils';
 import { CoreMessage, CoreTool, tool } from 'ai';
+import { ZodEventMapping } from './schemas';
 
-export async function agentDecide<T extends AnyAgent>(
-  agent: T,
-  options: AgentDecideOptions<T>
-): Promise<AgentDecision<EventsFromAgent<T>> | undefined> {
-  const resolvedOptions = {
-    ...agent.defaultOptions,
-    ...options,
-  };
+export async function agentDecide<TAgent extends AnyAgent>(
+  agent: TAgent,
+  options: AgentDecideInput<TAgent>
+): Promise<AgentDecision<TAgent> | undefined> {
+  const resolvedOptions = options;
   const {
     strategy = agent.strategy,
     goal,
@@ -28,6 +25,8 @@ export async function agentDecide<T extends AnyAgent>(
     machine,
     model = agent.model,
     messages,
+    episodeId = agent.episodeId,
+    maxAttempts = 2,
     ...otherDecideInput
   } = resolvedOptions;
 
@@ -41,16 +40,20 @@ export async function agentDecide<T extends AnyAgent>(
 
   let attempts = 0;
 
-  const maxAttempts = resolvedOptions.maxAttempts ?? 2;
-
   let decision: AgentDecision<any> | undefined;
 
+  const minimalState = {
+    value: state.value,
+    context: state.context,
+  };
+
   while (attempts++ < maxAttempts) {
     decision = await strategy(agent, {
+      episodeId,
       model,
       goal,
       events: filteredEventSchemas,
-      state,
+      state: minimalState,
       machine,
       messages: messages as CoreMessage[], // TODO: fix UIMessage thing
       ...otherDecideInput,
@@ -58,7 +61,6 @@ export async function agentDecide<T extends AnyAgent>(
 
     if (decision?.nextEvent) {
       agent.addDecision(decision);
-      await resolvedOptions.execute?.(decision.nextEvent);
       break;
     }
   }
@@ -68,7 +70,7 @@ export async function agentDecide<T extends AnyAgent>(
 
 export function fromDecision<T extends AnyAgent>(
   agent: T,
-  defaultInput?: AgentDecideInput<EventsFromAgent<T>>
+  defaultInput?: AgentDecideInput<EventFromAgent<T>>
 ): AgentDecisionLogic<any> {
   return fromPromise(async ({ input, self }) => {
     const parentRef = self._parent;
@@ -82,43 +84,39 @@ export function fromDecision<T extends AnyAgent>(
       ...defaultInput,
       ...inputObject,
     };
-    const state = {
-      value: snapshot.value,
-      context: resolvedInput.context,
-    };
 
-    const decision = await agentDecide(agent, {
+    const decision = await agentDecide<typeof agent>(agent, {
       machine: (parentRef as AnyActor).logic,
       state: snapshot,
-      context: resolvedInput.context,
-      execute: async (event) => {
-        parentRef.send(event);
-      },
       ...resolvedInput,
       // @ts-ignore
       messages: resolvedInput.messages,
     });
 
+    if (decision?.nextEvent) {
+      parentRef.send(decision.nextEvent);
+    }
+
     return decision;
   }) as AgentDecisionLogic<any>;
 }
 
-export function getToolMap<T extends AnyAgent>(
-  _agent: T,
+export function getToolMap<TAgent extends AnyAgent>(
+  agent: TAgent,
   input: AgentDecideInput<any>
 ): Record<string, CoreTool<any, any>> | undefined {
+  const events = input.events ?? (agent.events as ZodEventMapping);
   // Get all of the possible next transitions
   const transitions: TransitionData[] = input.machine
     ? getTransitions(input.state, input.machine)
-    : Object.entries(input.events).map(([eventType, { description }]) => ({
+    : Object.entries(events).map(([eventType, { description }]) => ({
         eventType,
         description,
       }));
 
   // Only keep the transitions that match the event types that are in the event mapping
   // TODO: allow for custom filters
-  const filter = (eventType: string) =>
-    Object.keys(input.events).includes(eventType);
+  const filter = (eventType: string) => Object.keys(events).includes(eventType);
 
   // Mapping of each event type (e.g. "mouse.click")
   // to a valid function name (e.g. "mouse_click")