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

package/src/agent.ts

@@ -1,36 +1,38 @@
 import {
   Actor,
   ActorRefLike,
-  EventObject,
   fromTransition,
+  SnapshotFrom,
   Subscription,
 } from 'xstate';
 import { ZodContextMapping, ZodEventMapping } from './schemas';
 import {
   AgentLogic,
   AgentMessage,
-  AgentStrategy,
-  EventsFromZodEventMapping,
+  AgentPolicy,
   GenerateTextOptions,
   AgentLongTermMemory,
   ObservedState,
   AgentObservationInput,
   AgentMemoryContext,
   AgentObservation,
-  ContextFromZodContextMapping,
   AgentFeedback,
   AgentMessageInput,
   AgentFeedbackInput,
   AgentDecision,
-  AgentDecideOptions,
   AnyAgent,
   AgentInteractInput,
+  AgentDecideInput,
+  EventFromAgent,
+  AgentInsightInput,
+  AgentInsight,
+  AgentDecisionInput,
 } from './types';
-import { simpleStrategy } from './strategies/simple';
-import { agentDecide } from './decide';
-import { getMachineHash, isActorRef, isMachineActor, randomId } from './utils';
+import { toolPolicy } from './policies/toolPolicy';
+import { isActorRef, isMachineActor, randomId } from './utils';
 import {
-  experimental_wrapLanguageModel,
+  CoreMessage,
+  wrapLanguageModel,
   LanguageModel,
   LanguageModelV1,
 } from 'ai';
@@ -43,7 +45,6 @@ export const agentLogic: AgentLogic<any> = fromTransition(
         state.feedback.push(event.feedback);
         emit({
           type: 'feedback',
-          // @ts-ignore TODO: fix types in XState
           feedback: event.feedback,
         });
         break;
@@ -52,7 +53,6 @@ export const agentLogic: AgentLogic<any> = fromTransition(
         state.observations.push(event.observation);
         emit({
           type: 'observation',
-          // @ts-ignore TODO: fix types in XState
           observation: event.observation,
         });
         break;
@@ -61,7 +61,6 @@ export const agentLogic: AgentLogic<any> = fromTransition(
         state.messages.push(event.message);
         emit({
           type: 'message',
-          // @ts-ignore TODO: fix types in XState
           message: event.message,
         });
         break;
@@ -74,8 +73,15 @@ export const agentLogic: AgentLogic<any> = fromTransition(
         });
         break;
       }
+      case 'agent.insight': {
+        state.insights.push(event.insight);
+        emit({
+          type: 'insight',
+          insight: event.insight,
+        });
+        break;
+      }
       default: {
-        // unrecognized
         console.warn('Unrecognized event', event);
         break;
       }
@@ -88,29 +94,28 @@ export const agentLogic: AgentLogic<any> = fromTransition(
       messages: [],
       observations: [],
       decisions: [],
+      insights: [],
     } as AgentMemoryContext<any>)
 );
 
 export function createAgent<
   const TContextSchema extends ZodContextMapping,
   const TEventSchemas extends ZodEventMapping,
-  TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
-  TContext = ContextFromZodContextMapping<TContextSchema>,
   TAgent extends AnyAgent = Agent<TContextSchema, TEventSchemas>
 >({
   id,
-  description: description,
+  description,
   model,
   events,
   context,
   episodeId,
-  strategy = simpleStrategy,
-  logic = agentLogic as AgentLogic<any>,
+  policy = toolPolicy,
+  logic = agentLogic,
 }: {
   /**
    * The unique identifier for the agent.
    *
-   * This should be the same across all sessions of a specific agent, as it can be
+   * This should be the same across all episodes of a specific agent, as it can be
    * used to retrieve memory for previous episodes of this agent.
    *
    * @example
@@ -123,17 +128,23 @@ export function createAgent<
    */
   id?: string;
   /**
-   * A description of the role of the agent
+   * A description of the role of the agent.
    */
   description?: string;
   /**
-   * Events that the agent can cause (send) in an environment
-   * that the agent knows about.
+   * Event schemas for events that the agent can trigger in an environment.
    */
   events: TEventSchemas;
+  /**
+   * The state context schema for the states that the agent can observe.
+   */
   context?: TContextSchema;
-  strategy?: AgentStrategy<Agent<TContextSchema, TEventSchemas>>;
-  stringify?: typeof JSON.stringify;
+  /**
+   * The default policy to use for `agent.decide(…)`.
+   *
+   * A policy is a strategy that the agent uses to decide which event to trigger next.
+   */
+  policy?: AgentPolicy<Agent<TContextSchema, TEventSchemas>>;
   /**
    * A function that retrieves the agent's long term memory
    */
@@ -141,10 +152,19 @@ export function createAgent<
     agent: Agent<TContextSchema, TEventSchemas>
   ) => AgentLongTermMemory<TAgent>;
   /**
-   * Agent logic
+   * Custom agent logic, which receives events for handling feedback,
+   * observations, messages, decisions, and insights.
    */
   logic?: AgentLogic<TAgent>;
+  /**
+   * The default language model for the agent to use in `agent.decide(…)`.
+   */
   model: LanguageModel;
+  /**
+   * The unique episode ID that this agent will run on.
+   *
+   * An episode is an instance of an agent interacting with an environment.
+   */
   episodeId?: string;
 }): Agent<TContextSchema, TEventSchemas> {
   return new Agent({
@@ -152,7 +172,7 @@ export function createAgent<
     context,
     events,
     description,
-    strategy: strategy,
+    policy: policy,
     model,
     logic,
     episodeId,
@@ -161,9 +181,7 @@ export function createAgent<
 
 export class Agent<
   const TContextSchema extends ZodContextMapping,
-  const TEventSchemas extends ZodEventMapping,
-  TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
-  TContext = ContextFromZodContextMapping<TContextSchema>
+  const TEventSchemas extends ZodEventMapping
 > extends Actor<AgentLogic<any>> {
   /**
    * The name of the agent. All agents with the same name are related and
@@ -177,14 +195,9 @@ export class Agent<
   public description?: string;
   public events: TEventSchemas;
   public context?: TContextSchema;
-  public strategy: AgentStrategy<Agent<TContextSchema, TEventSchemas>>;
-  // public types: {
-  //   events: TEvents;
-  //   context: Compute<TContext>;
-  // };
+  public policy: AgentPolicy<Agent<TContextSchema, TEventSchemas>>;
   public model: LanguageModel;
   public memory: AgentLongTermMemory<this> | undefined;
-  public defaultOptions: AgentDecideOptions<AnyAgent> | undefined; // todo
 
   constructor({
     logic = agentLogic as AgentLogic<any>,
@@ -195,7 +208,7 @@ export class Agent<
     events,
     context,
     episodeId,
-    strategy = simpleStrategy,
+    policy = toolPolicy,
   }: {
     logic: AgentLogic<any>;
     id?: string;
@@ -204,7 +217,7 @@ export class Agent<
     model: GenerateTextOptions['model'];
     events: TEventSchemas;
     context?: TContextSchema;
-    strategy?: AgentStrategy<Agent<TContextSchema, TEventSchemas>>;
+    policy?: AgentPolicy<Agent<TContextSchema, TEventSchemas>>;
     episodeId?: string;
   }) {
     super(logic);
@@ -214,30 +227,44 @@ export class Agent<
     this.description = description;
     this.events = events;
     this.context = context;
-    this.strategy = strategy;
+    this.policy = policy;
     this.id = id ?? randomId();
 
     this.start();
   }
 
   /**
-   * Called whenever the agent (LLM assistant) receives or sends a message.
+   * Called whenever the agent detects that a message was sent from the human, assistant, or system.
    */
   public onMessage(fn: (message: AgentMessage) => void) {
     return this.on('message', (ev) => fn(ev.message));
   }
 
   /**
-   * Called whenever the agent (LLM assistant) receives some feedback.
+   * Called whenever the agent 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.
+   * Called whenever the agent receives an observation.
+   */
+  public onObservation(fn: (observation: AgentObservation<this>) => void) {
+    return this.on('observation', (ev) => fn(ev.observation));
+  }
+
+  /**
+   * Called whenever the agent makes a decision.
+   */
+  public onDecision(fn: (decision: AgentDecision<this>) => void) {
+    return this.on('decision', (ev) => fn(ev.decision));
+  }
+
+  /**
+   * Adds a message to the agent's short-term (local) memory.
    */
-  public addMessage(messageInput: AgentMessageInput) {
+  public addMessage(messageInput: AgentMessageInput): AgentMessage {
     const message = {
       ...messageInput,
       id: messageInput.id ?? randomId(),
@@ -259,10 +286,11 @@ export class Agent<
   public addFeedback(feedbackInput: AgentFeedbackInput) {
     const feedback = {
       ...feedbackInput,
+      id: feedbackInput.id ?? randomId(),
       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',
@@ -287,11 +315,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({
@@ -309,10 +338,40 @@ export class Agent<
     return this.getSnapshot().context.observations;
   }
 
-  public addDecision(decision: AgentDecision<this>) {
+  public addInsight(insightInput: AgentInsightInput): AgentInsight {
+    const insight = {
+      ...insightInput,
+      episodeId: insightInput.episodeId ?? this.episodeId,
+      id: insightInput.id ?? randomId(),
+      timestamp: insightInput.timestamp ?? Date.now(),
+    } satisfies AgentInsight;
+
+    this.send({
+      type: 'agent.insight',
+      insight,
+    });
+
+    return insight;
+  }
+
+  public getInsights() {
+    return this.getSnapshot().context.insights;
+  }
+
+  public addDecision(input: AgentDecisionInput<this>) {
     this.send({
       type: 'agent.decision',
-      decision,
+      decision: {
+        id: input.id ?? randomId(),
+        episodeId: input.episodeId ?? this.episodeId,
+        timestamp: input.timestamp ?? Date.now(),
+        decisionId: input.decisionId ?? null,
+        policy: input.policy ?? null,
+        goalState: input.goalState ?? null,
+        nextEvent: input.nextEvent ?? null,
+        paths: input.paths ?? [],
+        ...input,
+      },
     });
   }
   /**
@@ -386,15 +445,15 @@ export class Agent<
 
     const agent = this;
 
-    async function handleObservation(
+    const handleObservation = async (
       observationInput: AgentObservationInput<any>
-    ) {
+    ) => {
       const observation = agent.addObservation(observationInput);
 
       const interactInput = getInput?.(observation);
 
       if (interactInput) {
-        const decision = await agentDecide(agent, {
+        const decision = await this.decide({
           machine,
           state: observation.state,
           ...interactInput,
@@ -408,7 +467,7 @@ export class Agent<
       }
 
       prevState = observationInput.state;
-    }
+    };
 
     // Inspect system, but only observe specified actor
     const sub = actorRefCheck
@@ -426,16 +485,18 @@ export class Agent<
               | string
               | undefined;
 
+            const decisions = agent.getDecisions();
+
             const decision = decisionId
-              ? agent.getDecisions().find((d) => d.id === decisionId)
+              ? decisions.find((d) => d.id === decisionId)
               : undefined;
 
             const observationInput = {
               event: inspEvent.event,
               prevState,
-              state: inspEvent.snapshot as any,
-              machine: (actorRef as any).src,
+              state: inspEvent.snapshot as SnapshotFrom<TActor>,
               goal: decision?.goal,
+              decisionId,
             } satisfies AgentObservationInput<any>;
 
             await handleObservation(observationInput);
@@ -446,11 +507,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,
-        goal: undefined,
       });
     }
 
@@ -479,15 +539,18 @@ export class Agent<
             const decisionId = inspEvent.event['_decision'] as
               | string
               | undefined;
+
+            const decisions = this.getDecisions();
+
             const decision = decisionId
-              ? this.getDecisions().find((d) => d.id === decisionId)
+              ? decisions.find((d) => d.id === decisionId)
               : undefined;
 
             const observationInput = {
+              decisionId,
               event: inspEvent.event,
               prevState,
-              state: inspEvent.snapshot as any,
-              machine: (actorRef as any).src,
+              state: inspEvent.snapshot as SnapshotFrom<TActor>,
               goal: decision?.goal,
             } satisfies AgentObservationInput<this>;
 
@@ -502,7 +565,7 @@ export class Agent<
   }
 
   public wrap(modelToWrap: LanguageModelV1) {
-    return experimental_wrapLanguageModel({
+    return wrapLanguageModel({
       model: modelToWrap,
       middleware: createAgentMiddleware(this),
     });
@@ -517,8 +580,58 @@ export class Agent<
    * - Additional `context`
    */
   public async decide(
-    opts: AgentDecideOptions<this>
+    input: AgentDecideInput<this>
   ): Promise<AgentDecision<this> | undefined> {
-    return agentDecide(this, opts);
+    const resolvedOptions = input;
+    const {
+      policy = this.policy,
+      goal,
+      allowedEvents,
+      events = this.events,
+      state,
+      machine,
+      model = this.model,
+      messages,
+      episodeId = this.episodeId,
+      maxAttempts = 2,
+      ...otherDecideInput
+    } = resolvedOptions;
+
+    const filteredEventSchemas = allowedEvents
+      ? Object.fromEntries(
+          Object.entries(events).filter(([key]) => {
+            return allowedEvents.includes(key as EventFromAgent<this>['type']);
+          })
+        )
+      : events;
+
+    let attempts = 0;
+
+    let decision: AgentDecision<this> | undefined;
+
+    const minimalState = {
+      value: state.value,
+      context: state.context,
+    };
+
+    while (attempts++ < maxAttempts) {
+      decision = await policy(this, {
+        episodeId,
+        model,
+        goal,
+        events: filteredEventSchemas,
+        state: minimalState,
+        machine,
+        messages: messages as CoreMessage[], // TODO: fix UIMessage thing
+        ...otherDecideInput,
+      });
+
+      if (decision?.nextEvent) {
+        this.addDecision(decision);
+        break;
+      }
+    }
+
+    return decision;
   }
 }

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 strategy)', async () => {
+test('agent.decide() makes a decision based on goal and state (tool policy)', async () => {
   const model = new MockLanguageModelV1({
     doGenerate,
   });
@@ -211,7 +211,7 @@ test.each([
           ? params.mode.tools?.map((t) => t.name)
           : [];
 
-      console.log('try', attempts, 'max', maxAttempts);
+      // console.log('try', attempts, 'max', maxAttempts);
 
       const toolCalls =
         succeed && attempts++ === (maxAttempts ?? 2) - 1
@@ -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,80 +1,34 @@
-import { AnyActor, AnyMachineSnapshot, fromPromise } from 'xstate';
+import {
+  AnyActor,
+  AnyMachineSnapshot,
+  fromPromise,
+  PromiseActorLogic,
+} from 'xstate';
 import {
   AnyAgent,
-  AgentDecideOptions,
-  AgentDecisionLogic,
-  AgentDecision,
   AgentDecideInput,
   TransitionData,
-  EventFromAgent,
+  AgentDecision,
 } from './types';
 import { getTransitions } from './utils';
-import { CoreMessage, CoreTool, tool } from 'ai';
-
-export async function agentDecide<TAgent extends AnyAgent>(
-  agent: TAgent,
-  options: AgentDecideOptions<TAgent>
-): Promise<AgentDecision<TAgent> | undefined> {
-  const resolvedOptions = {
-    ...agent.defaultOptions,
-    ...options,
-  };
-  const {
-    strategy = agent.strategy,
-    goal,
-    allowedEvents,
-    events = agent.events,
-    state,
-    machine,
-    model = agent.model,
-    messages,
-    ...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 decision: AgentDecision<any> | undefined;
-
-  const minimalState = {
-    value: state.value,
-    context: state.context,
-  };
-
-  while (attempts++ < maxAttempts) {
-    decision = await strategy(agent, {
-      model,
-      goal,
-      events: filteredEventSchemas,
-      state: minimalState,
-      machine,
-      messages: messages as CoreMessage[], // TODO: fix UIMessage thing
-      ...otherDecideInput,
-    });
+import { CoreTool, generateText, LanguageModel, tool } from 'ai';
+import { ZodEventMapping } from './schemas';
 
-    if (decision?.nextEvent) {
-      agent.addDecision(decision);
-      await resolvedOptions.execute?.(decision.nextEvent);
-      break;
-    }
-  }
+export type AgentDecideLogicInput = {
+  goal: string;
+  model?: LanguageModel;
+  context?: Record<string, any>;
+} & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt'>;
 
-  return decision;
-}
+export type MachineDecisionLogic<TAgent extends AnyAgent> = PromiseActorLogic<
+  AgentDecision<TAgent> | undefined,
+  AgentDecideLogicInput | string
+>;
 
-export function fromDecision<T extends AnyAgent>(
-  agent: T,
-  defaultInput?: AgentDecideInput<EventFromAgent<T>>
-): AgentDecisionLogic<any> {
+export function fromDecision<TAgent extends AnyAgent>(
+  agent: TAgent,
+  defaultInput?: AgentDecideInput<TAgent>
+): MachineDecisionLogic<any> {
   return fromPromise(async ({ input, self }) => {
     const parentRef = self._parent;
     if (!parentRef) {
@@ -88,37 +42,39 @@ export function fromDecision<T extends AnyAgent>(
       ...inputObject,
     };
 
-    const decision = await agentDecide(agent, {
+    const decision = await agent.decide({
       machine: (parentRef as AnyActor).logic,
       state: snapshot,
-      execute: async (event) => {
-        parentRef.send(event);
-      },
+      allowedEvents: resolvedInput.allowedEvents as any[],
       ...resolvedInput,
       // @ts-ignore
       messages: resolvedInput.messages,
     });
 
+    if (decision?.nextEvent) {
+      parentRef.send(decision.nextEvent);
+    }
+
     return decision;
-  }) as AgentDecisionLogic<any>;
+  }) as MachineDecisionLogic<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")

package/src/index.ts

@@ -2,3 +2,4 @@ export { createAgent } from './agent';
 export { fromText, fromTextStream } from './text';
 export { fromDecision } from './decide';
 export * from './types';
+export * from './policies';