@statelyai/agent 1.0.0-beta.0 → 1.0.0

package/src/index.ts

@@ -1,4 +1,4 @@
 export { createAgent } from './agent';
-export { fromText, fromTextStream, agentGenerateText } from './text';
-export { fromDecision, agentDecide } from './decision';
+export { fromText, fromTextStream } from './text';
+export { fromDecision } from './decision';
 export * from './types';

package/src/planners/shortestPathPlanner.ts

@@ -1,7 +1,7 @@
-import { Agent, AgentPlan, AgentPlanInput } from '../types';
+import { Agent, AgentPlan, AgentPlanInput, AnyAgent } from '../types';
 import { getShortestPaths } from '@xstate/graph';
 
-export async function simplePlanner<T extends Agent<any>>(
+export async function simplePlanner<T extends AnyAgent>(
   agent: T,
   input: AgentPlanInput<any>
 ): Promise<AgentPlan<any> | undefined> {

package/src/planners/simplePlanner.ts

@@ -1,15 +1,16 @@
 import { CoreTool, tool } from 'ai';
 import {
-  Agent,
   AgentPlan,
   AgentPlanInput,
   ObservedState,
   PromptTemplate,
   TransitionData,
+  AnyAgent,
 } from '../types';
 import { getAllTransitions } from '../utils';
 import { AnyStateMachine } from 'xstate';
 import { defaultTextTemplate } from '../templates/defaultText';
+import { getMessages } from '../text';
 
 function getTransitions(
   state: ObservedState,
@@ -27,11 +28,11 @@ const simplePlannerPromptTemplate: PromptTemplate<any> = (data) => {
   return `
 ${defaultTextTemplate(data)}
 
-Only make a single tool call to achieve the above goal.
+Make at most one tool call to achieve the above goal. If the goal cannot be achieved with any tool calls, do not make any tool call.
   `.trim();
 };
 
-export async function simplePlanner<T extends Agent<any>>(
+export async function simplePlanner<T extends AnyAgent>(
   agent: T,
   input: AgentPlanInput<any>
 ): Promise<AgentPlan<any> | undefined> {
@@ -81,7 +82,7 @@ export async function simplePlanner<T extends Agent<any>>(
     toolMap[toolTransitionData.name] = tool({
       description: toolZodType?.description ?? toolTransitionData.description,
       parameters: toolZodType,
-      execute: async (params) => {
+      execute: async (params: Record<string, any>) => {
         const event = {
           type: toolTransitionData.eventType,
           ...params,
@@ -92,23 +93,32 @@ export async function simplePlanner<T extends Agent<any>>(
     });
   }
 
+  if (!Object.keys(toolMap).length) {
+    // No valid transitions for the specified tools
+    return undefined;
+  }
+
   // Create a prompt with the given context and goal.
-  // The template is used to ensure that a single tool call is made.
+  // The template is used to ensure that a single tool call at most is made.
   const prompt = simplePlannerPromptTemplate({
     context: input.state.context,
     goal: input.goal,
   });
 
+  const messages = await getMessages(agent, prompt, input);
+
   const result = await agent.generateText({
-    prompt,
-    tools: toolMap,
     toolChoice: 'required',
     ...input,
+    prompt,
+    messages,
+    tools: toolMap,
   });
 
   const singleResult = result.toolResults[0];
 
   if (!singleResult) {
+    console.log(toolMap);
     // TODO: retries?
     console.warn('No tool call results returned');
     return undefined;
@@ -117,11 +127,12 @@ export async function simplePlanner<T extends Agent<any>>(
   return {
     goal: input.goal,
     state: input.state,
-    steps: [
-      {
-        event: singleResult.result,
-      },
-    ],
+    execute: async (state) => {
+      if (JSON.stringify(state) === JSON.stringify(input.state)) {
+        return singleResult.result;
+      }
+      return undefined;
+    },
     nextEvent: singleResult.result,
     sessionId: agent.sessionId,
     timestamp: Date.now(),

package/src/schemas.ts

@@ -1,15 +1,11 @@
-import type { SomeZodObject } from 'zod';
-import { AnyEventObject } from 'xstate';
-import { ObservedState } from './types';
+import { ZodType, type SomeZodObject } from 'zod';
 
 export type ZodEventMapping = {
   // map event types to Zod types
   [eventType: string]: SomeZodObject;
 };
 
-export type ZodActionMapping = {
-  [eventType: string]: {
-    schema: SomeZodObject;
-    action: (state: ObservedState, event: AnyEventObject) => Promise<void>;
-  };
+export type ZodContextMapping = {
+  // map context keys to Zod types
+  [contextKey: string]: ZodType;
 };

package/src/strategies/chain-of-note.ts

@@ -1,7 +1,7 @@
 import { GenerateTextResult, LanguageModel } from 'ai';
 import wiki, { wikiSearchResult, wikiSummary } from 'wikipedia';
 import { assign, fromPromise, setup } from 'xstate';
-import { Agent } from '../types';
+import { AnyAgent } from '../types';
 
 const searchWiki = fromPromise(
   async ({
@@ -44,7 +44,7 @@ export const chainOfNote = setup({
   types: {
     input: {} as {
       model: LanguageModel;
-      agent: Agent<any>;
+      agent: AnyAgent;
       prompt: string;
     },
     context: {} as {
@@ -56,7 +56,7 @@ export const chainOfNote = setup({
           }[]
         | null;
       model: LanguageModel;
-      agent: Agent<any>;
+      agent: AnyAgent;
       prompt: string;
     },
     output: {} as GenerateTextResult<any>,

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,
@@ -19,7 +18,7 @@ import {
   fromPromise,
   toObserver,
 } from 'xstate';
-import { nanoid } from 'nanoid';
+import { randomId } from './utils';
 
 /**
  * Gets an array of messages from the given prompt, based on the agent and options.
@@ -29,15 +28,13 @@ import { nanoid } from 'nanoid';
  * @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
 ) {
@@ -61,7 +58,7 @@ export async function agentGenerateText<T extends Agent<any>>(
   };
   const template = resolvedOptions.template ?? defaultTextTemplate;
   // TODO: check if messages was provided instead
-  const id = nanoid();
+  const id = randomId();
   const goal =
     typeof resolvedOptions.prompt === 'string'
       ? resolvedOptions.prompt
@@ -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 = {
@@ -109,7 +106,7 @@ export async function agentStreamText(
   };
   const template = resolvedOptions.template ?? defaultTextTemplate;
 
-  const id = nanoid();
+  const id = randomId();
   const goal =
     typeof resolvedOptions.prompt === 'string'
       ? resolvedOptions.prompt
@@ -148,7 +145,7 @@ export async function agentStreamText(
           rawResponse: res.rawResponse,
         },
         content: res.text,
-        id: nanoid(),
+        id: randomId(),
         timestamp: Date.now(),
         responseId: id,
       });
@@ -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,21 +22,38 @@ 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];
 
 export type StreamTextOptions = Parameters<typeof streamText>[0];
 
-export type AgentPlanInput<TEvent extends EventObject> = {
-  model: LanguageModel;
+export type AgentPlanInput<TEvent extends EventObject> = Omit<
+  GenerateTextOptions,
+  'prompt' | 'messages' | 'tools'
+> & {
+  /**
+   * The currently observed state.
+   */
   state: ObservedState;
+  /**
+   * The goal for the agent to accomplish.
+   * The agent will create a plan based on this goal.
+   */
   goal: string;
+  /**
+   * The events that the agent can trigger. This is a mapping of
+   * event types to Zod event schemas.
+   */
   events: ZodEventMapping;
+  /**
+   * The state machine that represents the environment the agent
+   * is interacting with.
+   */
   machine?: AnyStateMachine;
   /**
-   * The previous plan
+   * The previous plan.
    */
   previousPlan?: AgentPlan<TEvent>;
 };
@@ -45,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;
@@ -86,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;
@@ -124,7 +142,7 @@ export interface AgentFeedbackInput {
   timestamp?: number;
 }
 
-export type AgentMessageHistory = CoreMessage & {
+export type AgentMessage = CoreMessage & {
   timestamp: number;
   id: string;
   /**
@@ -136,7 +154,7 @@ export type AgentMessageHistory = CoreMessage & {
   sessionId: string;
 };
 
-export type AgentMessageHistoryInput = CoreMessage & {
+export type AgentMessageInput = CoreMessage & {
   timestamp?: number;
   id?: string;
   /**
@@ -152,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;
 }
@@ -161,6 +180,7 @@ export interface AgentObservationInput {
   prevState: ObservedState | undefined;
   event: AnyEventObject;
   state: ObservedState;
+  machine?: AnyStateMachine;
   timestamp?: number;
 }
 
@@ -186,7 +206,7 @@ export type AgentEmitted<TEvents extends EventObject> =
     }
   | {
       type: 'message';
-      message: AgentMessageHistory;
+      message: AgentMessage;
     }
   | {
       type: 'plan';
@@ -205,7 +225,7 @@ export type AgentLogic<TEvents extends EventObject> = ActorLogic<
     }
   | {
       type: 'agent.message';
-      message: AgentMessageHistory;
+      message: AgentMessage;
     }
   | {
       type: 'agent.plan';
@@ -223,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;
@@ -254,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: (
@@ -271,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.
    */
-  interact: <TActor extends AnyActorRef>(
+  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>(
     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>);
 
@@ -304,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>;
 }
 
@@ -339,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[];
 };
@@ -376,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;