@statelyai/agent 1.0.0-beta.1 → 1.0.0

package/src/text.ts

@@ -5,13 +5,12 @@ import type {
   StreamTextResult,
 } from 'ai';
 import {
-  Agent,
   AgentGenerateTextOptions,
   AgentStreamTextOptions,
+  AnyAgent,
 } from './types';
 import { defaultTextTemplate } from './templates/defaultText';
 import {
-  AnyMachineSnapshot,
   ObservableActorLogic,
   Observer,
   PromiseActorLogic,
@@ -29,15 +28,13 @@ import { randomId } from './utils';
  * @param options
  * @returns
  */
-async function getMessages(
-  agent: Agent<any>,
+export async function getMessages(
+  agent: AnyAgent,
   prompt: string,
-  options: AgentStreamTextOptions
+  options: Omit<AgentGenerateTextOptions, 'prompt'>
 ): Promise<CoreMessage[]> {
   let messages: CoreMessage[] = [];
-  if (options.messages === true) {
-    messages = agent.select((s) => s.messages);
-  } else if (typeof options.messages === 'function') {
+  if (typeof options.messages === 'function') {
     messages = await options.messages(agent);
   } else if (options.messages) {
     messages = options.messages;
@@ -51,7 +48,7 @@ async function getMessages(
   return messages;
 }
 
-export async function agentGenerateText<T extends Agent<any>>(
+export async function agentGenerateText<T extends AnyAgent>(
   agent: T,
   options: AgentGenerateTextOptions
 ) {
@@ -100,7 +97,7 @@ export async function agentGenerateText<T extends Agent<any>>(
 }
 
 export async function agentStreamText(
-  agent: Agent<any>,
+  agent: AnyAgent,
   options: AgentStreamTextOptions
 ): Promise<StreamTextResult<any>> {
   const resolvedOptions = {
@@ -158,21 +155,16 @@ export async function agentStreamText(
   return result;
 }
 
-export function fromTextStream<T extends Agent<any>>(
+export function fromTextStream<T extends AnyAgent>(
   agent: T,
   defaultOptions?: AgentStreamTextOptions
 ): ObservableActorLogic<
   { textDelta: string },
   Omit<AgentStreamTextOptions, 'context'> & {
-    context?: AgentStreamTextOptions['context'] | boolean;
+    context?: AgentStreamTextOptions['context'];
   }
 > {
-  return fromObservable(({ input, self }) => {
-    const context =
-      input.context === true
-        ? (self._parent?.getSnapshot() as AnyMachineSnapshot).context
-        : input.context;
-
+  return fromObservable(({ input }) => {
     const observers = new Set<Observer<{ textDelta: string }>>();
 
     // TODO: check if messages was provided instead
@@ -181,7 +173,7 @@ export function fromTextStream<T extends Agent<any>>(
       const result = await agentStreamText(agent, {
         ...defaultOptions,
         ...input,
-        context,
+        context: input.context,
       });
 
       for await (const part of result.fullStream) {
@@ -208,24 +200,20 @@ export function fromTextStream<T extends Agent<any>>(
   });
 }
 
-export function fromText<T extends Agent<any>>(
+export function fromText<T extends AnyAgent>(
   agent: T,
   defaultOptions?: AgentGenerateTextOptions
 ): PromiseActorLogic<
   GenerateTextResult<Record<string, CoreTool<any, any>>>,
   Omit<AgentGenerateTextOptions, 'context'> & {
-    context?: AgentGenerateTextOptions['context'] | boolean;
+    context?: AgentGenerateTextOptions['context'];
   }
 > {
-  return fromPromise(async ({ input, self }) => {
-    const context =
-      input.context === true
-        ? (self._parent?.getSnapshot() as AnyMachineSnapshot).context
-        : input.context;
+  return fromPromise(async ({ input }) => {
     return await agentGenerateText(agent, {
       ...input,
       ...defaultOptions,
-      context,
+      context: input.context,
     });
   });
 }

package/src/types.ts

@@ -22,7 +22,7 @@ import {
   streamText,
   StreamTextResult,
 } from 'ai';
-import { ZodEventMapping } from './schemas';
+import { ZodContextMapping, ZodEventMapping } from './schemas';
 import { TypeOf } from 'zod';
 
 export type GenerateTextOptions = Parameters<typeof generateText>[0];
@@ -62,10 +62,11 @@ export type AgentPlan<TEvent extends EventObject> = {
   goal: string;
   state: ObservedState;
   content?: string;
-  steps?: Array<{
-    event: TEvent;
-    state?: ObservedState;
-  }>;
+  /**
+   * Executes the plan based on the given `state` and resolves with
+   * a potential next `event` to trigger to achieve the `goal`.
+   */
+  execute: (state: ObservedState) => Promise<TEvent | undefined>;
   nextEvent: TEvent | undefined;
   sessionId: string;
   timestamp: number;
@@ -103,14 +104,14 @@ export type PromptTemplate<TEvents extends EventObject> = (data: {
    */
   observations?: AgentObservation<any>[]; // TODO
   feedback?: AgentFeedback[];
-  messages?: AgentMessageHistory[];
+  messages?: AgentMessage[];
   plans?: AgentPlan<TEvents>[];
 }) => string;
 
-export type AgentPlanner<T extends Agent<any>> = (
-  agent: T['eventTypes'],
-  options: AgentPlanInput<T['eventTypes']>
-) => Promise<AgentPlan<T['eventTypes']> | undefined>;
+export type AgentPlanner<T extends AnyAgent> = (
+  agent: T,
+  input: AgentPlanInput<T['types']['events']>
+) => Promise<AgentPlan<T['types']['events']> | undefined>;
 
 export type AgentDecideOptions = {
   goal: string;
@@ -141,7 +142,7 @@ export interface AgentFeedbackInput {
   timestamp?: number;
 }
 
-export type AgentMessageHistory = CoreMessage & {
+export type AgentMessage = CoreMessage & {
   timestamp: number;
   id: string;
   /**
@@ -153,7 +154,7 @@ export type AgentMessageHistory = CoreMessage & {
   sessionId: string;
 };
 
-export type AgentMessageHistoryInput = CoreMessage & {
+export type AgentMessageInput = CoreMessage & {
   timestamp?: number;
   id?: string;
   /**
@@ -169,6 +170,7 @@ export interface AgentObservation<TActor extends AnyActorRef> {
   prevState: SnapshotFrom<TActor> | undefined;
   event: EventFrom<TActor>;
   state: SnapshotFrom<TActor>;
+  machineHash: string | undefined;
   sessionId: string;
   timestamp: number;
 }
@@ -178,6 +180,7 @@ export interface AgentObservationInput {
   prevState: ObservedState | undefined;
   event: AnyEventObject;
   state: ObservedState;
+  machine?: AnyStateMachine;
   timestamp?: number;
 }
 
@@ -203,7 +206,7 @@ export type AgentEmitted<TEvents extends EventObject> =
     }
   | {
       type: 'message';
-      message: AgentMessageHistory;
+      message: AgentMessage;
     }
   | {
       type: 'plan';
@@ -222,7 +225,7 @@ export type AgentLogic<TEvents extends EventObject> = ActorLogic<
     }
   | {
       type: 'agent.message';
-      message: AgentMessageHistory;
+      message: AgentMessage;
     }
   | {
       type: 'agent.plan';
@@ -240,21 +243,30 @@ export type EventsFromZodEventMapping<TEventSchemas extends ZodEventMapping> =
     } & TypeOf<TEventSchemas[K]>;
   }>;
 
-export type Agent<TEvents extends EventObject> = ActorRefFrom<
+export type ContextFromZodContextMapping<
+  TContextSchema extends ZodContextMapping
+> = {
+  [K in keyof TContextSchema & string]: TypeOf<TContextSchema[K]>;
+};
+
+export type Agent<TContext, TEvents extends EventObject> = ActorRefFrom<
   AgentLogic<TEvents>
 > & {
   /**
-   * The general name of the agent. All agents with the same name are related and
+   * The name of the agent. All agents with the same name are related and
    * able to share experiences (observations, feedback) with each other.
    */
-  name: string;
+  name?: string;
   /**
-   * The unique id of the agent. This is used to partition message history.
+   * The unique identifier for the agent.
    */
   id?: string;
   description?: string;
   events: ZodEventMapping;
-  eventTypes: TEvents;
+  types: {
+    events: TEvents;
+    context: Compute<TContext>;
+  };
   model: LanguageModel;
   defaultOptions: GenerateTextOptions;
   memory: AgentLongTermMemory | undefined;
@@ -271,7 +283,7 @@ export type Agent<TEvents extends EventObject> = ActorRefFrom<
    *
    * - The `goal` for the agent to achieve
    * - The observed current `state`
-   * - The `logic` (e.g. a state machine) that specifies what can happen next
+   * - The `machine` (e.g. a state machine) that specifies what can happen next
    * - Additional `context`
    */
   decide: (
@@ -288,32 +300,100 @@ export type Agent<TEvents extends EventObject> = ActorRefFrom<
     options: AgentStreamTextOptions
   ) => Promise<StreamTextResult<Record<string, CoreTool<any, any>>>>;
 
-  addObservation: (observation: AgentObservationInput) => AgentObservation<any>; // TODO
-  addMessage: (history: AgentMessageHistoryInput) => AgentMessageHistory;
-  addFeedback: (feedbackItem: AgentFeedbackInput) => AgentFeedback;
+  addObservation: (
+    observationInput: AgentObservationInput
+  ) => AgentObservation<any>; // TODO
+  addMessage: (messageInput: AgentMessageInput) => AgentMessage;
+  addFeedback: (feedbackInput: AgentFeedbackInput) => AgentFeedback;
   addPlan: (plan: AgentPlan<TEvents>) => void;
   /**
    * Called whenever the agent (LLM assistant) receives or sends a message.
    */
-  onMessage: (callback: (message: AgentMessageHistory) => void) => void;
+  onMessage: (callback: (message: AgentMessage) => void) => void;
   /**
    * Selects agent data from its context.
+   *
+   * @deprecated Select from `agent.getSnapshot().context` directly or:
+   * - `agent.getMessages()`
+   * - `agent.getObservations()`
+   * - `agent.getFeedback()`
+   * - `agent.getPlans()`
    */
   select: <T>(selector: (context: AgentMemoryContext) => T) => T;
 
   /**
-   * Inspects state machine actor transitions and automatically observes
-   * (prevState, event, state) tuples.
+   * Retrieves messages from the agent's short-term (local) memory.
+   */
+  getMessages: () => AgentMessage[];
+
+  /**
+   * Retrieves observations from the agent's short-term (local) memory.
+   */
+  getObservations: () => AgentObservation<Agent<TContext, TEvents>>[];
+
+  /**
+   * Retrieves feedback from the agent's short-term (local) memory.
+   */
+  getFeedback: () => AgentFeedback[];
+
+  /**
+   * Retrieves strategies from the agent's short-term (local) memory.
+   */
+  getPlans: () => AgentPlan<TEvents>[];
+
+  /**
+   * 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();
+   * ```
+   */
+  interact<TActor extends AnyActorRef>(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();
+   * ```
    */
-  interact: <TActor extends AnyActorRef>(
+  interact<TActor extends AnyActorRef>(
     actorRef: TActor,
-    getInput?: (
+    getInput: (
       observation: AgentObservation<TActor>
     ) => AgentDecisionInput | undefined
-  ) => Subscription;
+  ): Subscription;
 };
 
-export type AnyAgent = Agent<any>;
+export type AnyAgent = Agent<any, any>;
 
 export type FromAgent<T> = T | ((self: AnyAgent) => T | Promise<T>);
 
@@ -321,7 +401,7 @@ export interface CommonTextOptions {
   prompt: FromAgent<string>;
   model?: LanguageModel;
   context?: Record<string, any>;
-  messages?: FromAgent<CoreMessage[]> | true;
+  messages?: FromAgent<CoreMessage[]>;
   template?: PromptTemplate<any>;
 }
 
@@ -356,7 +436,7 @@ export type ObservedStateFrom<TActor extends AnyActorRef> = Pick<
 
 export type AgentMemoryContext = {
   observations: AgentObservation<any>[]; // TODO
-  messages: AgentMessageHistory[];
+  messages: AgentMessage[];
   plans: AgentPlan<any>[];
   feedback: AgentFeedback[];
 };
@@ -393,3 +473,5 @@ export interface AIAdapter {
   generateText: typeof generateText;
   streamText: typeof streamText;
 }
+
+export type Compute<A extends any> = { [K in keyof A]: A[K] } & unknown;

package/src/utils.ts

@@ -1,18 +1,50 @@
-import { AnyMachineSnapshot, AnyStateNode } from 'xstate';
+import { AnyMachineSnapshot, AnyStateMachine, AnyStateNode } from 'xstate';
+import hash from 'object-hash';
 import { TransitionData } from './types';
 
 export function getAllTransitions(state: AnyMachineSnapshot): TransitionData[] {
   const nodes = state._nodes;
   const transitions = (nodes as AnyStateNode[])
     .map((node) => [...(node as AnyStateNode).transitions.values()])
-    .flat(2)
-    .map((transition) => ({
-      ...transition,
-      guard:
-        typeof transition.guard === 'string'
-          ? { type: transition.guard }
-          : (transition.guard as any), // TODO: fix
-    }));
+    .map((nodeTransitions) => {
+      return nodeTransitions.map((nodeEventTransitions) => {
+        return nodeEventTransitions.map((transition) => {
+          return {
+            ...transition,
+            guard:
+              typeof transition.guard === 'string'
+                ? { type: transition.guard }
+                : (transition.guard as any), // TODO: fix
+          };
+        });
+      });
+    })
+    .flat(2);
+
+  return transitions;
+}
+
+export function getAllMachineTransitions(
+  stateNode: AnyStateNode
+): TransitionData[] {
+  const transitions: TransitionData[] = [...stateNode.transitions.values()]
+    .map((nodeTransitions) => {
+      return nodeTransitions.map((transition) => {
+        return {
+          ...transition,
+          guard:
+            typeof transition.guard === 'string'
+              ? { type: transition.guard }
+              : (transition.guard as any), // TODO: fix
+        };
+      });
+    })
+    .flat(2);
+
+  for (const s of Object.values(stateNode.states)) {
+    const stateTransitions = getAllMachineTransitions(s);
+    transitions.push(...stateTransitions);
+  }
 
   return transitions;
 }
@@ -26,3 +58,15 @@ export function randomId() {
   const random = Math.random().toString(36).substring(2, 9);
   return timestamp + random;
 }
+
+const machineHashes: WeakMap<AnyStateMachine, string> = new WeakMap();
+/**
+ * Returns a string hash representing only the transitions in the state machine.
+ */
+export function getMachineHash(machine: AnyStateMachine): string {
+  if (machineHashes.has(machine)) return machineHashes.get(machine)!;
+  const transitions = getAllMachineTransitions(machine.root);
+  const machineHash = hash(transitions);
+  machineHashes.set(machine, machineHash);
+  return machineHash;
+}

package/src/templates/defaultToolCall.ts

@@ -1,10 +0,0 @@
-import { PromptTemplate } from '../types';
-import { defaultTextTemplate } from './defaultText';
-
-export const defaultToolCallTemplate: PromptTemplate<any> = (data) => {
-  return `
-${defaultTextTemplate(data)}
-
-Only make a single tool call to achieve the above goal.
-  `.trim();
-};