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

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,14 +22,17 @@ import {
   AgentFeedback,
   AgentMessageInput,
   AgentFeedbackInput,
-  AgentPlan,
+  AgentDecision,
   Compute,
   AgentDecisionInput,
   AgentDecideOptions,
+  AnyAgent,
+  EventsFromAgent,
+  AgentInteractInput,
 } from './types';
-import { simplePlanner } from './planners/simplePlanner';
+import { simpleStrategy } from './strategies/simple';
 import { agentDecide } from './decide';
-import { getMachineHash, isActorRef, randomId } from './utils';
+import { getMachineHash, isActorRef, isMachineActor, randomId } from './utils';
 import {
   experimental_wrapLanguageModel,
   LanguageModel,
@@ -68,17 +70,19 @@ 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;
       }
-      default:
+      default: {
+        // unrecognized
+        console.warn('Unrecognized event', event);
         break;
+      }
     }
     return state;
   },
@@ -87,7 +91,7 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
       feedback: [],
       messages: [],
       observations: [],
-      plans: [],
+      decisions: [],
     } as AgentMemoryContext)
 );
 
@@ -98,22 +102,19 @@ 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,
+  episodeId,
+  strategy = simpleStrategy,
   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 +125,6 @@ export function createAgent<
    * ```
    */
   id?: string;
-  /**
-   * The name of the agent
-   */
-  name?: string;
   /**
    * A description of the role of the agent
    */
@@ -138,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
@@ -150,210 +147,19 @@ export function createAgent<
    * Agent logic
    */
   logic?: AgentLogic<TEvents>;
-} & GenerateTextOptions): Agent<TContextSchema, TEventSchemas> {
+  model: LanguageModel;
+  episodeId?: string;
+}): Agent<TContextSchema, TEventSchemas> {
   return new Agent({
     id,
     context,
     events,
-    name,
     description,
-    planner,
+    strategy: strategy,
     model,
     logic,
+    episodeId,
   }) 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<
@@ -374,14 +180,14 @@ 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: any; // todo
+  public defaultOptions: AgentDecideOptions<AnyAgent> | undefined; // todo
 
   constructor({
     logic = agentLogic as AgentLogic<TEvents>,
@@ -391,7 +197,8 @@ export class Agent<
     model,
     events,
     context,
-    planner = simplePlanner,
+    episodeId,
+    strategy = simpleStrategy,
   }: {
     logic: AgentLogic<TEvents>;
     id?: string;
@@ -400,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();
   }
@@ -422,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.
    */
@@ -448,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;
@@ -497,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;
   }
 
   /**
@@ -558,15 +372,16 @@ 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);
+    const actorRefCheck = isActorRef(actorRef) && actorRef.src;
+    const machine = isMachineActor(actorRef) ? actorRef.src : undefined;
 
     let prevState: ObservedState | undefined = undefined;
     let subscribed = true;
@@ -579,16 +394,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;
@@ -622,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,
       });
@@ -675,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) {
+  public async decide(
+    opts: AgentDecideOptions<this>
+  ): Promise<AgentDecision<EventsFromAgent<this>> | undefined> {
     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,172 @@ 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 strategy)', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
+
+  const agent = createAgent({
+    id: 'test',
+    model,
+    events: {
+      MOVE: z.object({}),
+    },
+  });
+
+  const decision = 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(decision).toBeDefined();
+  expect(decision!.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 decision = 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(decision).toBeUndefined();
+    } else {
+      expect(decision).toBeDefined();
+      expect(decision!.nextEvent).toEqual(
+        expect.objectContaining({
+          type: 'MOVE',
+        })
+      );
+    }
+  }
+);
+
+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);
+  }
+);