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

package/src/agent.ts

@@ -1,8 +1,6 @@
 import {
   Actor,
   ActorRefLike,
-  AnyEventObject,
-  AnyStateMachine,
   EventObject,
   fromTransition,
   Subscription,
@@ -11,7 +9,7 @@ import { ZodContextMapping, ZodEventMapping } from './schemas';
 import {
   AgentLogic,
   AgentMessage,
-  AgentPlanner,
+  AgentStrategy,
   EventsFromZodEventMapping,
   GenerateTextOptions,
   AgentLongTermMemory,
@@ -23,13 +21,12 @@ import {
   AgentFeedback,
   AgentMessageInput,
   AgentFeedbackInput,
-  AgentPlan,
-  Compute,
-  AgentDecisionInput,
+  AgentDecision,
   AgentDecideOptions,
   AnyAgent,
+  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 {
@@ -39,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': {
@@ -69,12 +66,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,23 +87,25 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
       feedback: [],
       messages: [],
       observations: [],
-      plans: [],
-    } as AgentMemoryContext)
+      decisions: [],
+    } 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,
   model,
   events,
   context,
-  planner = simplePlanner as AgentPlanner<Agent<TContextSchema, TEventSchemas>>,
-  logic = agentLogic as AgentLogic<TEvents>,
+  episodeId,
+  strategy = simpleStrategy,
+  logic = agentLogic as AgentLogic<any>,
 }: {
   /**
    * The unique identifier for the agent.
@@ -134,28 +132,30 @@ 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
    */
   getMemory?: (
     agent: Agent<TContextSchema, TEventSchemas>
-  ) => AgentLongTermMemory;
+  ) => AgentLongTermMemory<TAgent>;
   /**
    * Agent logic
    */
-  logic?: AgentLogic<TEvents>;
+  logic?: AgentLogic<TAgent>;
   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;
 }
 
@@ -164,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.
@@ -177,43 +177,45 @@ 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 memory: AgentLongTermMemory<this> | undefined;
   public defaultOptions: AgentDecideOptions<AnyAgent> | undefined; // todo
 
   constructor({
-    logic = agentLogic as AgentLogic<TEvents>,
+    logic = agentLogic as AgentLogic<any>,
     id,
     name,
     description,
     model,
     events,
     context,
-    planner = simplePlanner,
+    episodeId,
+    strategy = simpleStrategy,
   }: {
-    logic: AgentLogic<TEvents>;
+    logic: AgentLogic<any>;
     id?: string;
     name?: string;
     description?: string;
     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 +227,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.
    */
@@ -250,8 +259,8 @@ export class Agent<
   public addFeedback(feedbackInput: AgentFeedbackInput) {
     const feedback = {
       ...feedbackInput,
+      comment: feedbackInput.comment ?? undefined,
       attributes: { ...feedbackInput.attributes },
-      reward: feedbackInput.reward ?? 0,
       timestamp: feedbackInput.timestamp ?? Date.now(),
       episodeId: this.episodeId,
     } satisfies AgentFeedback;
@@ -270,7 +279,7 @@ export class Agent<
   }
 
   public addObservation(
-    observationInput: AgentObservationInput
+    observationInput: AgentObservationInput<this>
   ): AgentObservation<any> {
     const { prevState, event, state } = observationInput;
     const observation = {
@@ -300,17 +309,17 @@ export class Agent<
     return this.getSnapshot().context.observations;
   }
 
-  public addPlan(plan: AgentPlan<TEvents>) {
+  public addDecision(decision: AgentDecision<this>) {
     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,36 +370,40 @@ 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;
 
-    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);
         }
       }
 
@@ -409,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,
+            } satisfies AgentObservationInput<any>;
 
             await handleObservation(observationInput);
           },
@@ -425,9 +447,10 @@ 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,
+        goal: undefined,
       });
     }
 
@@ -440,7 +463,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
@@ -453,12 +476,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 = {
               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;
 
@@ -478,14 +509,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<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,77 @@ import {
   AnyAgent,
   AgentDecideOptions,
   AgentDecisionLogic,
-  AgentDecisionInput,
-  AgentPlanner,
-  AgentPlan,
-  EventsFromZodEventMapping,
-  AgentPlanInput,
+  AgentDecision,
+  AgentDecideInput,
   TransitionData,
+  EventFromAgent,
 } 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> {
+export async function agentDecide<TAgent extends AnyAgent>(
+  agent: TAgent,
+  options: AgentDecideOptions<TAgent>
+): Promise<AgentDecision<TAgent> | 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;
+
+  const minimalState = {
+    value: state.value,
+    context: state.context,
+  };
 
   while (attempts++ < maxAttempts) {
-    plan = await planner(agent, {
+    decision = await strategy(agent, {
       model,
       goal,
-      events,
-      state,
+      events: filteredEventSchemas,
+      state: minimalState,
       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<EventFromAgent<T>>
 ): AgentDecisionLogic<any> {
   return fromPromise(async ({ input, self }) => {
     const parentRef = self._parent;
@@ -75,27 +87,25 @@ export function fromDecision(
       ...defaultInput,
       ...inputObject,
     };
-    const state = {
-      value: snapshot.value,
-      context: resolvedInput.context,
-    };
 
-    const plan = await agentDecide(agent, {
+    const decision = await agentDecide(agent, {
       machine: (parentRef as AnyActor).logic,
-      state,
+      state: snapshot,
       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,50 @@
+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> = ({
+  stateValue,
+  context,
+  goal,
+}) => {
+  return `
+${convertToXml({ stateValue, 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({
+    stateValue: input.state.value,
+    context: input.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;
+}