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

package/src/agent.ts

@@ -1,23 +1,21 @@
 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,
@@ -25,12 +23,16 @@ import {
   AnyAgent,
   AgentInteractInput,
   AgentDecideInput,
+  EventFromAgent,
+  AgentInsightInput,
+  AgentInsight,
+  AgentDecisionInput,
 } from './types';
-import { simpleStrategy } from './strategies/simpleStrategy';
-import { agentDecide } from './decide';
+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,11 +195,7 @@ 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;
 
@@ -194,7 +208,7 @@ export class Agent<
     events,
     context,
     episodeId,
-    strategy = simpleStrategy,
+    policy = toolPolicy,
   }: {
     logic: AgentLogic<any>;
     id?: string;
@@ -203,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);
@@ -213,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(),
@@ -258,6 +286,7 @@ 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(),
@@ -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,14 +485,16 @@ 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,
+              state: inspEvent.snapshot as SnapshotFrom<TActor>,
               goal: decision?.goal,
               decisionId,
             } satisfies AgentObservationInput<any>;
@@ -478,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,
+              state: inspEvent.snapshot as SnapshotFrom<TActor>,
               goal: decision?.goal,
             } satisfies AgentObservationInput<this>;
 
@@ -501,7 +565,7 @@ export class Agent<
   }
 
   public wrap(modelToWrap: LanguageModelV1) {
-    return experimental_wrapLanguageModel({
+    return wrapLanguageModel({
       model: modelToWrap,
       middleware: createAgentMiddleware(this),
     });
@@ -518,6 +582,56 @@ export class Agent<
   public async decide(
     input: AgentDecideInput<this>
   ): Promise<AgentDecision<this> | undefined> {
-    return agentDecide(this, input);
+    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

package/src/decide.ts

@@ -1,77 +1,34 @@
-import { AnyActor, AnyMachineSnapshot, fromPromise } from 'xstate';
+import {
+  AnyActor,
+  AnyMachineSnapshot,
+  fromPromise,
+  PromiseActorLogic,
+} from 'xstate';
 import {
   AnyAgent,
-  AgentDecisionLogic,
-  AgentDecision,
   AgentDecideInput,
   TransitionData,
-  EventFromAgent,
+  AgentDecision,
 } from './types';
 import { getTransitions } from './utils';
-import { CoreMessage, CoreTool, tool } from 'ai';
+import { CoreTool, generateText, LanguageModel, tool } from 'ai';
 import { ZodEventMapping } from './schemas';
 
-export async function agentDecide<TAgent extends AnyAgent>(
-  agent: TAgent,
-  options: AgentDecideInput<TAgent>
-): Promise<AgentDecision<TAgent> | undefined> {
-  const resolvedOptions = options;
-  const {
-    strategy = agent.strategy,
-    goal,
-    allowedEvents,
-    events = agent.events,
-    state,
-    machine,
-    model = agent.model,
-    messages,
-    episodeId = agent.episodeId,
-    maxAttempts = 2,
-    ...otherDecideInput
-  } = resolvedOptions;
-
-  const filteredEventSchemas = allowedEvents
-    ? Object.fromEntries(
-        Object.entries(events).filter(([key]) => {
-          return allowedEvents.includes(key);
-        })
-      )
-    : events;
-
-  let attempts = 0;
-
-  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: minimalState,
-      machine,
-      messages: messages as CoreMessage[], // TODO: fix UIMessage thing
-      ...otherDecideInput,
-    });
+export type AgentDecideLogicInput = {
+  goal: string;
+  model?: LanguageModel;
+  context?: Record<string, any>;
+} & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt'>;
 
-    if (decision?.nextEvent) {
-      agent.addDecision(decision);
-      break;
-    }
-  }
+export type MachineDecisionLogic<TAgent extends AnyAgent> = PromiseActorLogic<
+  AgentDecision<TAgent> | undefined,
+  AgentDecideLogicInput | string
+>;
 
-  return decision;
-}
-
-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) {
@@ -85,9 +42,10 @@ export function fromDecision<T extends AnyAgent>(
       ...inputObject,
     };
 
-    const decision = await agentDecide<typeof agent>(agent, {
+    const decision = await agent.decide({
       machine: (parentRef as AnyActor).logic,
       state: snapshot,
+      allowedEvents: resolvedInput.allowedEvents as any[],
       ...resolvedInput,
       // @ts-ignore
       messages: resolvedInput.messages,
@@ -98,7 +56,7 @@ export function fromDecision<T extends AnyAgent>(
     }
 
     return decision;
-  }) as AgentDecisionLogic<any>;
+  }) as MachineDecisionLogic<any>;
 }
 
 export function getToolMap<TAgent extends AnyAgent>(

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';

package/src/{strategies/chainOfThought.ts → policies/chainOfThoughtPolicy.ts}

@@ -5,8 +5,8 @@ import {
   AgentDecision,
   PromptTemplate,
 } from '../types';
-import { getMessages } from '../text';
-import { simpleStrategy } from './simpleStrategy';
+import { combinePromptAndMessages } from '../text';
+import { toolPolicy } from './toolPolicy';
 import { convertToXml } from '../utils';
 
 const chainOfThoughtPromptTemplate: PromptTemplate<any> = ({
@@ -14,14 +14,12 @@ const chainOfThoughtPromptTemplate: PromptTemplate<any> = ({
   context,
   goal,
 }) => {
-  return `
-${convertToXml({ stateValue, context, goal })}
+  return `${convertToXml({ stateValue, context, goal })}
 
-How would you achieve the goal? Think step-by-step.
-`.trim();
+How would you achieve the goal? Think step-by-step.`;
 };
 
-export async function chainOfThoughtStrategy<T extends AnyAgent>(
+export async function chainOfThoughtPolicy<T extends AnyAgent>(
   agent: T,
   input: AgentDecideInput<any>
 ): Promise<AgentDecision<any> | undefined> {
@@ -31,7 +29,7 @@ export async function chainOfThoughtStrategy<T extends AnyAgent>(
     goal: input.goal,
   });
 
-  const messages = await getMessages(agent, prompt, input);
+  const messages = combinePromptAndMessages(prompt, input.messages);
 
   const model = input.model ? agent.wrap(input.model) : agent.model;
 
@@ -41,7 +39,7 @@ export async function chainOfThoughtStrategy<T extends AnyAgent>(
     messages,
   });
 
-  const decision = await simpleStrategy(agent, {
+  const decision = await toolPolicy(agent, {
     ...input,
     messages: messages.concat(result.response.messages),
   });

package/src/policies/index.ts

@@ -0,0 +1,3 @@
+export * from './chainOfThoughtPolicy';
+export * from './shortestPathPolicy';
+export * from './toolPolicy';

package/src/{strategies/shortestPath.test.ts → policies/shortestPathPolicy.test.ts}

@@ -1,7 +1,7 @@
 import { createAgent, TypesFromAgent } from '..';
 import { assign, createActor, setup } from 'xstate';
 import { z } from 'zod';
-import { experimental_shortestPathStrategy } from './shortestPath';
+import { experimental_shortestPathPolicy } from './shortestPathPolicy';
 import { test, expect } from 'vitest';
 import { dummyResponseValues, MockLanguageModelV1 } from '../mockModel';
 
@@ -84,7 +84,7 @@ test.skip('should find shortest path to goal', async () => {
     }),
     goal: 'Get the counter to exactly 3',
     state: counterActor.getSnapshot(),
-    strategy: experimental_shortestPathStrategy,
+    policy: experimental_shortestPathPolicy,
   });
 
   expect(decision?.nextEvent?.type).toBe('increment');