@statelyai/agent 1.1.5 → 2.0.0-next.0

package/src/agent.ts

@@ -1,34 +1,42 @@
 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,
 } 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 { agentDecide } from './decide';
+import { getMachineHash, isActorRef, randomId } from './utils';
+import {
+  experimental_wrapLanguageModel,
+  LanguageModel,
+  LanguageModelV1,
+} from 'ai';
+import { createAgentMiddleware } from './middleware';
 
 export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
   (state, event, { emit }) => {
@@ -89,16 +97,16 @@ export function createAgent<
   TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
   TContext = ContextFromZodContextMapping<TContextSchema>
 >({
+  id,
   name,
   description,
   model,
   events,
   context,
-  planner = simplePlanner as AgentPlanner<Agent<TContext, TEvents>>,
+  planner = simplePlanner as AgentPlanner<Agent<TContextSchema, TEventSchemas>>,
   stringify = JSON.stringify,
   getMemory,
   logic = agentLogic as AgentLogic<TEvents>,
-  adapter = vercelAdapter,
   ...generateTextOptions
 }: {
   /**
@@ -130,110 +138,441 @@ 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;
+} & GenerateTextOptions): Agent<TContextSchema, TEventSchemas> {
+  return new Agent({
+    id,
+    context,
+    events,
+    name,
+    description,
+    planner,
+    model,
+    logic,
+  }) 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.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.decide = (opts) => {
-    return agentDecide(agent, opts);
+  //       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<
+  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: any; // todo
+
+  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;
 
-  agent.addMessage = (messageInput) => {
+    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));
+  }
+
+  /**
+   * 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);
+  }
 
-  agent.streamText = (opts) => agentStreamText(agent, opts);
+  public getMessages() {
+    return this.getSnapshot().context.messages;
+  }
 
-  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);
 
-  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);
 
@@ -241,7 +580,9 @@ export function createAgent<
 
       if (input) {
         await agentDecide(agent, {
-          machine: actorRef.src as AnyStateMachine,
+          machine: actorRefCheck
+            ? (actorRef.src as AnyStateMachine)
+            : undefined,
           state: observation.state,
           execute: async (event) => {
             actorRef.send(event);
@@ -254,26 +595,28 @@ export function createAgent<
     }
 
     // Inspect system, but only observe specified actor
-    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);
-      },
-    });
+    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;
+
+            await handleObservation(observationInput);
+          },
+        })
+      : undefined;
 
     // If actor already started, interact with current state
     if ((actorRef as any)._processingStatus === 1) {
@@ -287,14 +630,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);
+
+    const sub = actorRefCheck
+      ? 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.types = {} as any;
+            this.addObservation(observationInput);
+          },
+        })
+      : undefined;
+
+    return sub ?? { unsubscribe: () => {} };
+  }
 
-  agent.start();
+  public wrap(modelToWrap: LanguageModelV1) {
+    return experimental_wrapLanguageModel({
+      model: modelToWrap,
+      middleware: createAgentMiddleware(this),
+    });
+  }
 
-  return agent;
+  /**
+   * 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) {
+    return agentDecide(this, opts);
+  }
 }