@statelyai/agent 1.1.6 → 2.0.0-next.1

package/src/agent.ts

@@ -1,34 +1,43 @@
 import {
+  Actor,
+  ActorRefLike,
   AnyEventObject,
   AnyStateMachine,
-  createActor,
   EventObject,
   fromTransition,
-  Observer,
-  toObserver,
+  Subscription,
 } from 'xstate';
 import { ZodContextMapping, ZodEventMapping } from './schemas';
 import {
-  Agent,
   AgentLogic,
   AgentMessage,
   AgentPlanner,
   EventsFromZodEventMapping,
   GenerateTextOptions,
   AgentLongTermMemory,
-  AIAdapter,
   ObservedState,
   AgentObservationInput,
   AgentMemoryContext,
   AgentObservation,
   ContextFromZodContextMapping,
   AgentFeedback,
+  AgentMessageInput,
+  AgentFeedbackInput,
+  AgentPlan,
+  Compute,
+  AgentDecisionInput,
+  AgentDecideOptions,
+  AnyAgent,
 } from './types';
-import { simplePlanner } from './planners/simplePlanner';
-import { agentGenerateText, agentStreamText } from './text';
-import { agentDecide } from './decision';
-import { vercelAdapter } from './adapters/vercel';
-import { getMachineHash, randomId } from './utils';
+import { simplePlanner } from './planners/simple';
+import { agentDecide } from './decide';
+import { getMachineHash, isActorRef, isMachineActor, randomId } from './utils';
+import {
+  experimental_wrapLanguageModel,
+  LanguageModel,
+  LanguageModelV1,
+} from 'ai';
+import { createAgentMiddleware } from './middleware';
 
 export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
   (state, event, { emit }) => {
@@ -69,8 +78,11 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
         });
         break;
       }
-      default:
+      default: {
+        // unrecognized
+        console.warn('Unrecognized event', event);
         break;
+      }
     }
     return state;
   },
@@ -89,23 +101,19 @@ export function createAgent<
   TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
   TContext = ContextFromZodContextMapping<TContextSchema>
 >({
-  name,
-  description,
+  id,
+  description: description,
   model,
   events,
   context,
-  planner = simplePlanner as AgentPlanner<Agent<TContext, TEvents>>,
-  stringify = JSON.stringify,
-  getMemory,
+  planner = simplePlanner as AgentPlanner<Agent<TContextSchema, TEventSchemas>>,
   logic = agentLogic as AgentLogic<TEvents>,
-  adapter = vercelAdapter,
-  ...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
@@ -116,10 +124,6 @@ export function createAgent<
    * ```
    */
   id?: string;
-  /**
-   * The name of the agent
-   */
-  name?: string;
   /**
    * A description of the role of the agent
    */
@@ -130,150 +134,292 @@ export function createAgent<
    */
   events: TEventSchemas;
   context?: TContextSchema;
-  planner?: AgentPlanner<Agent<TContext, TEvents>>;
+  planner?: AgentPlanner<Agent<TContextSchema, TEventSchemas>>;
   stringify?: typeof JSON.stringify;
   /**
    * A function that retrieves the agent's long term memory
    */
-  getMemory?: (agent: Agent<TContext, TEvents>) => AgentLongTermMemory;
+  getMemory?: (
+    agent: Agent<TContextSchema, TEventSchemas>
+  ) => AgentLongTermMemory;
   /**
    * Agent logic
    */
   logic?: AgentLogic<TEvents>;
-  adapter?: AIAdapter;
-} & GenerateTextOptions): Agent<TContext, TEvents> {
-  const agent = createActor(logic) as unknown as Agent<TContext, TEvents>;
-  agent.events = events;
-  agent.model = model;
-  agent.name = name;
-  agent.description = description;
-  agent.adapter = adapter;
-  agent.defaultOptions = { ...generateTextOptions, model };
-  agent.select = (selector) => {
-    return selector(agent.getSnapshot().context);
-  };
-  agent.memory = getMemory ? getMemory(agent) : undefined;
+  model: LanguageModel;
+} & GenerateTextOptions): Agent<TContextSchema, TEventSchemas> {
+  return new Agent({
+    id,
+    context,
+    events,
+    description,
+    planner,
+    model,
+    logic,
+  }) as any;
+}
 
-  agent.onMessage = (callback) => {
-    agent.on('message', (ev) => callback(ev.message));
+export class Agent<
+  const TContextSchema extends ZodContextMapping,
+  const TEventSchemas extends ZodEventMapping,
+  TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
+  TContext = ContextFromZodContextMapping<TContextSchema>
+> extends Actor<AgentLogic<TEvents>> {
+  /**
+   * The name of the agent. All agents with the same name are related and
+   * able to share experiences (observations, feedback) with each other.
+   */
+  public name?: string;
+  /**
+   * The unique identifier for the agent.
+   */
+  public episodeId: string;
+  public description?: string;
+  public events: TEventSchemas;
+  public context?: TContextSchema;
+  public planner?: AgentPlanner<Agent<TContextSchema, TEventSchemas>>;
+  public types: {
+    events: TEvents;
+    context: Compute<TContext>;
   };
+  public model: LanguageModel;
+  public memory: AgentLongTermMemory | undefined;
+  public defaultOptions: AgentDecideOptions<AnyAgent> | undefined; // todo
 
-  agent.decide = (opts) => {
-    return agentDecide(agent, opts);
-  };
+  constructor({
+    logic = agentLogic as AgentLogic<TEvents>,
+    id,
+    name,
+    description,
+    model,
+    events,
+    context,
+    planner = simplePlanner,
+  }: {
+    logic: AgentLogic<TEvents>;
+    id?: string;
+    name?: string;
+    description?: string;
+    model: GenerateTextOptions['model'];
+    events: TEventSchemas;
+    context?: TContextSchema;
+    planner?: AgentPlanner<Agent<TContextSchema, TEventSchemas>>;
+  }) {
+    super(logic);
+    this.model = model;
+    this.episodeId = id ?? randomId();
+    this.name = name;
+    this.description = description;
+    this.events = events;
+    this.context = context;
+    this.planner = planner;
+    this.types = {} as any;
+
+    this.start();
+  }
+
+  /**
+   * Called whenever the agent (LLM assistant) receives or sends a message.
+   */
+  public onMessage(fn: (message: AgentMessage) => void) {
+    return this.on('message', (ev) => fn(ev.message));
+  }
 
-  agent.addMessage = (messageInput) => {
+  /**
+   * Retrieves messages from the agent's short-term (local) memory.
+   */
+  public addMessage(messageInput: AgentMessageInput) {
     const message = {
       ...messageInput,
       id: messageInput.id ?? randomId(),
       timestamp: messageInput.timestamp ?? Date.now(),
-      sessionId: agent.sessionId,
-      correlationId: messageInput.correlationId ?? randomId(),
+      episodeId: this.episodeId,
     } satisfies AgentMessage;
-    agent.send({
+    this.send({
       type: 'agent.message',
       message,
     });
 
     return message;
-  };
-  agent.getMessages = () => agent.getSnapshot().context.messages;
+  }
 
-  agent.generateText = (opts) => agentGenerateText(agent, opts);
+  public getMessages() {
+    return this.getSnapshot().context.messages;
+  }
 
-  agent.streamText = (opts) => agentStreamText(agent, opts);
-
-  agent.addFeedback = (feedbackInput) => {
+  public addFeedback(feedbackInput: AgentFeedbackInput) {
     const feedback = {
       ...feedbackInput,
       attributes: { ...feedbackInput.attributes },
       reward: feedbackInput.reward ?? 0,
       timestamp: feedbackInput.timestamp ?? Date.now(),
-      sessionId: agent.sessionId,
+      episodeId: this.episodeId,
     } satisfies AgentFeedback;
-    agent.send({
+    this.send({
       type: 'agent.feedback',
       feedback,
     });
     return feedback;
-  };
-  agent.getFeedback = () => agent.getSnapshot().context.feedback;
+  }
+
+  /**
+   * Retrieves feedback from the agent's short-term (local) memory.
+   */
+  public getFeedback() {
+    return this.getSnapshot().context.feedback;
+  }
 
-  agent.addObservation = (observationInput) => {
+  public addObservation(
+    observationInput: AgentObservationInput
+  ): AgentObservation<any> {
     const { prevState, event, state } = observationInput;
     const observation = {
       prevState,
       event,
       state,
       id: observationInput.id ?? randomId(),
-      sessionId: agent.sessionId,
+      episodeId: this.episodeId,
       timestamp: observationInput.timestamp ?? Date.now(),
       machineHash: observationInput.machine
         ? getMachineHash(observationInput.machine)
         : undefined,
     } satisfies AgentObservation<any>;
 
-    agent.send({
+    this.send({
       type: 'agent.observe',
       observation,
     });
 
     return observation;
-  };
-  agent.getObservations = () => agent.getSnapshot().context.observations;
+  }
+
+  /**
+   * Retrieves observations from the agent's short-term (local) memory.
+   */
+  public getObservations() {
+    return this.getSnapshot().context.observations;
+  }
 
-  agent.addPlan = (plan) => {
-    agent.send({
+  public addPlan(plan: AgentPlan<TEvents>) {
+    this.send({
       type: 'agent.plan',
       plan,
     });
-  };
-  agent.getPlans = () => agent.getSnapshot().context.plans;
+  }
+  /**
+   * Retrieves strategies from the agent's short-term (local) memory.
+   */
+  public getPlans() {
+    return this.getSnapshot().context.plans;
+  }
+
+  /**
+   * Interacts with this state machine actor by inspecting state transitions and storing them as observations.
+   *
+   * Observations contain the `prevState`, `event`, and current `state` of this
+   * actor, as well as other properties that are useful when recalled.
+   * These observations are stored in the `agent`'s short-term (local) memory
+   * and can be retrieved via `agent.getObservations()`.
+   *
+   * @example
+   * ```ts
+   * // Only observes the actor's state transitions
+   * agent.interact(actor);
+   *
+   * actor.start();
+   * ```
+   */
+  public interact<TActor extends ActorRefLike>(actorRef: TActor): Subscription;
+  /**
+   * Interacts with this state machine actor by:
+   * 1. Inspecting state transitions and storing them as observations
+   * 2. Deciding what to do next (which event to send the actor) based on
+   * the agent input returned from `getInput(observation)`, if `getInput(…)` is provided as the 2nd argument.
+   *
+   * Observations contain the `prevState`, `event`, and current `state` of this
+   * actor, as well as other properties that are useful when recalled.
+   * These observations are stored in the `agent`'s short-term (local) memory
+   * and can be retrieved via `agent.getObservations()`.
+   *
+   * @example
+   * ```ts
+   * // Observes the actor's state transitions and
+   * // makes a decision if on the "summarize" state
+   * agent.interact(actor, observed => {
+   *   if (observed.state.matches('summarize')) {
+   *     return {
+   *       context: observed.state.context,
+   *       goal: 'Summarize the message'
+   *     }
+   *   }
+   * });
+   *
+   * actor.start();
+   * ```
+   */
+  public interact<TActor extends ActorRefLike>(
+    actorRef: TActor,
+    getInput: (
+      observation: AgentObservation<TActor>
+    ) => AgentDecisionInput | undefined
+  ): Subscription;
+  public interact<TActor extends ActorRefLike>(
+    actorRef: TActor,
+    getInput?: (
+      observation: AgentObservation<TActor>
+    ) => AgentDecisionInput | undefined
+  ): Subscription {
+    const actorRefCheck = isActorRef(actorRef) && actorRef.src;
+    const machine = isMachineActor(actorRef) ? actorRef.src : undefined;
 
-  agent.interact = ((actorRef, getInput) => {
     let prevState: ObservedState | undefined = undefined;
     let subscribed = true;
 
+    const agent = this;
+
     async function handleObservation(observationInput: AgentObservationInput) {
       const observation = agent.addObservation(observationInput);
 
       const input = getInput?.(observation);
 
       if (input) {
-        await agentDecide(agent, {
-          machine: actorRef.src as AnyStateMachine,
+        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;
     }
 
     // Inspect system, but only observe specified actor
-    actorRef.system.inspect({
-      next: async (inspEvent) => {
-        if (
-          !subscribed ||
-          inspEvent.actorRef !== actorRef ||
-          inspEvent.type !== '@xstate.snapshot'
-        ) {
-          return;
-        }
+    const sub = actorRefCheck
+      ? 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;
+            const observationInput = {
+              event: inspEvent.event,
+              prevState,
+              state: inspEvent.snapshot as any,
+              machine: (actorRef as any).src,
+            } satisfies AgentObservationInput;
 
-        await handleObservation(observationInput);
-      },
-    });
+            await handleObservation(observationInput);
+          },
+        })
+      : undefined;
 
     // If actor already started, interact with current state
     if ((actorRef as any)._processingStatus === 1) {
@@ -287,14 +433,59 @@ export function createAgent<
 
     return {
       unsubscribe: () => {
+        sub?.unsubscribe();
         subscribed = false;
-      }, // TODO: make this actually unsubscribe
+      },
     };
-  }) as typeof agent.interact;
+  }
+
+  public observe<TActor extends ActorRefLike>(actorRef: TActor): Subscription {
+    let prevState: ObservedState = actorRef.getSnapshot();
+    const actorRefCheck = isActorRef(actorRef);
 
-  agent.types = {} as any;
+    const sub = actorRefCheck
+      ? actorRef.system.inspect({
+          next: async (inspEvent) => {
+            if (
+              inspEvent.actorRef !== actorRef ||
+              inspEvent.type !== '@xstate.snapshot'
+            ) {
+              return;
+            }
 
-  agent.start();
+            const observationInput = {
+              event: inspEvent.event,
+              prevState,
+              state: inspEvent.snapshot as any,
+              machine: (actorRef as any).src,
+            } satisfies AgentObservationInput;
 
-  return agent;
+            prevState = observationInput.state;
+
+            this.addObservation(observationInput);
+          },
+        })
+      : undefined;
+
+    return sub ?? { unsubscribe: () => {} };
+  }
+
+  public wrap(modelToWrap: LanguageModelV1) {
+    return experimental_wrapLanguageModel({
+      model: modelToWrap,
+      middleware: createAgentMiddleware(this),
+    });
+  }
+
+  /**
+   * Resolves with an `AgentPlan` 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>) {
+    return agentDecide(this, opts);
+  }
 }