@statelyai/agent 1.0.0-beta.1 → 1.1.2

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;
@@ -127,21 +128,28 @@ export type AgentDecideOptions = {
 >;
 
 export interface AgentFeedback {
-  goal: string;
-  observationId: string;
+  goal?: string;
+  observationId?: string;
+  /**
+   * The message correlation that the feedback is relevant for
+   */
+  correlationId?: string;
   attributes: Record<string, any>;
+  reward: number;
   timestamp: number;
   sessionId: string;
 }
 
 export interface AgentFeedbackInput {
-  goal: string;
-  observationId: string; // Observation ID;
-  attributes: Record<string, any>;
+  goal?: string;
+  observationId?: string;
+  correlationId?: string;
+  attributes?: Record<string, any>;
   timestamp?: number;
+  reward?: number;
 }
 
-export type AgentMessageHistory = CoreMessage & {
+export type AgentMessage = CoreMessage & {
   timestamp: number;
   id: string;
   /**
@@ -151,9 +159,11 @@ export type AgentMessageHistory = CoreMessage & {
   responseId?: string;
   result?: GenerateTextResult<any>;
   sessionId: string;
+  correlationId: string;
+  parentCorrelationId?: string;
 };
 
-export type AgentMessageHistoryInput = CoreMessage & {
+export type AgentMessageInput = CoreMessage & {
   timestamp?: number;
   id?: string;
   /**
@@ -161,6 +171,8 @@ export type AgentMessageHistoryInput = CoreMessage & {
    * which message this message is responding to, if any.
    */
   responseId?: string;
+  correlationId?: string;
+  parentCorrelationId?: string;
   result?: GenerateTextResult<any>;
 };
 
@@ -169,6 +181,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 +191,7 @@ export interface AgentObservationInput {
   prevState: ObservedState | undefined;
   event: AnyEventObject;
   state: ObservedState;
+  machine?: AnyStateMachine;
   timestamp?: number;
 }
 
@@ -203,7 +217,7 @@ export type AgentEmitted<TEvents extends EventObject> =
     }
   | {
       type: 'message';
-      message: AgentMessageHistory;
+      message: AgentMessage;
     }
   | {
       type: 'plan';
@@ -222,7 +236,7 @@ export type AgentLogic<TEvents extends EventObject> = ActorLogic<
     }
   | {
       type: 'agent.message';
-      message: AgentMessageHistory;
+      message: AgentMessage;
     }
   | {
       type: 'agent.plan';
@@ -240,21 +254,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 +294,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: (
@@ -281,49 +304,124 @@ export type Agent<TEvents extends EventObject> = ActorRefFrom<
   // Generate text
   generateText: (
     options: AgentGenerateTextOptions
-  ) => Promise<GenerateTextResult<Record<string, any>>>;
+  ) => Promise<AgentGenerateTextResult>;
 
   // Stream text
   streamText: (
     options: AgentStreamTextOptions
-  ) => Promise<StreamTextResult<Record<string, CoreTool<any, any>>>>;
+  ) => Promise<AgentStreamTextResult>;
 
-  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.
    */
-  interact: <TActor extends AnyActorRef>(
+  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>(
     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>);
+export type FromAgent<T> = T | ((agent: AnyAgent) => T | Promise<T>);
 
-export interface CommonTextOptions {
+export type CommonTextOptions = {
   prompt: FromAgent<string>;
   model?: LanguageModel;
   context?: Record<string, any>;
-  messages?: FromAgent<CoreMessage[]> | true;
+  messages?: FromAgent<CoreMessage[]>;
   template?: PromptTemplate<any>;
-}
+  correlationId?: string;
+  parentCorrelationId?: string;
+};
+
+export type TextResultMeta = {
+  correlationId: string;
+  parentCorrelationId?: string;
+};
 
 export type AgentGenerateTextOptions = Omit<
   GenerateTextOptions,
@@ -331,12 +429,16 @@ export type AgentGenerateTextOptions = Omit<
 > &
   CommonTextOptions;
 
+export type AgentGenerateTextResult = GenerateTextResult<any> & TextResultMeta;
+
 export type AgentStreamTextOptions = Omit<
   StreamTextOptions,
   'model' | 'prompt' | 'messages'
 > &
   CommonTextOptions;
 
+export type AgentStreamTextResult = StreamTextResult<any> & TextResultMeta;
+
 export interface ObservedState {
   /**
    * The current state value of the state machine, e.g.
@@ -356,7 +458,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 +495,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/.changeset/shaggy-buttons-itch.md

@@ -1,5 +0,0 @@
----
-"@statelyai/agent": patch
----
-
-Use nanoid

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();
-};