@statelyai/agent 1.0.0-beta.1 → 1.1.2

package/src/agent.ts

@@ -7,11 +7,11 @@ import {
   Observer,
   toObserver,
 } from 'xstate';
-import { ZodEventMapping } from './schemas';
+import { ZodContextMapping, ZodEventMapping } from './schemas';
 import {
   Agent,
   AgentLogic,
-  AgentMessageHistory,
+  AgentMessage,
   AgentPlanner,
   EventsFromZodEventMapping,
   GenerateTextOptions,
@@ -20,12 +20,15 @@ import {
   ObservedState,
   AgentObservationInput,
   AgentMemoryContext,
+  AgentObservation,
+  ContextFromZodContextMapping,
+  AgentFeedback,
 } from './types';
 import { simplePlanner } from './planners/simplePlanner';
 import { agentGenerateText, agentStreamText } from './text';
 import { agentDecide } from './decision';
 import { vercelAdapter } from './adapters/vercel';
-import { randomId } from './utils';
+import { getMachineHash, randomId } from './utils';
 
 export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
   (state, event, { emit }) => {
@@ -71,33 +74,52 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
     }
     return state;
   },
-  {
-    feedback: [],
-    messages: [],
-    observations: [],
-    plans: [],
-  } as AgentMemoryContext
+  () =>
+    ({
+      feedback: [],
+      messages: [],
+      observations: [],
+      plans: [],
+    } as AgentMemoryContext)
 );
 
 export function createAgent<
+  const TContextSchema extends ZodContextMapping,
   const TEventSchemas extends ZodEventMapping,
-  TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>
+  TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
+  TContext = ContextFromZodContextMapping<TContextSchema>
 >({
   name,
   description,
   model,
   events,
-  planner = simplePlanner as AgentPlanner<Agent<TEvents>>,
+  context,
+  planner = simplePlanner as AgentPlanner<Agent<TContext, TEvents>>,
   stringify = JSON.stringify,
   getMemory,
   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.
+   *
+   * @example
+   * ```ts
+   * const agent = createAgent({
+   *  id: 'recipe-assistant',
+   *  // ...
+   * });
+   * ```
+   */
+  id?: string;
   /**
    * The name of the agent
    */
-  name: string;
+  name?: string;
   /**
    * A description of the role of the agent
    */
@@ -107,21 +129,20 @@ export function createAgent<
    * that the agent knows about.
    */
   events: TEventSchemas;
-  planner?: AgentPlanner<Agent<TEvents>>;
+  context?: TContextSchema;
+  planner?: AgentPlanner<Agent<TContext, TEvents>>;
   stringify?: typeof JSON.stringify;
   /**
    * A function that retrieves the agent's long term memory
    */
-  getMemory?: (agent: Agent<any>) => AgentLongTermMemory;
+  getMemory?: (agent: Agent<TContext, TEvents>) => AgentLongTermMemory;
   /**
    * Agent logic
    */
   logic?: AgentLogic<TEvents>;
   adapter?: AIAdapter;
-} & GenerateTextOptions): Agent<TEvents> {
-  const messageHistoryListeners: Observer<AgentMessageHistory>[] = [];
-
-  const agent = createActor(logic) as unknown as Agent<TEvents>;
+} & GenerateTextOptions): Agent<TContext, TEvents> {
+  const agent = createActor(logic) as unknown as Agent<TContext, TEvents>;
   agent.events = events;
   agent.model = model;
   agent.name = name;
@@ -134,7 +155,7 @@ export function createAgent<
   agent.memory = getMemory ? getMemory(agent) : undefined;
 
   agent.onMessage = (callback) => {
-    messageHistoryListeners.push(toObserver(callback));
+    agent.on('message', (ev) => callback(ev.message));
   };
 
   agent.decide = (opts) => {
@@ -147,7 +168,8 @@ export function createAgent<
       id: messageInput.id ?? randomId(),
       timestamp: messageInput.timestamp ?? Date.now(),
       sessionId: agent.sessionId,
-    };
+      correlationId: messageInput.correlationId ?? randomId(),
+    } satisfies AgentMessage;
     agent.send({
       type: 'agent.message',
       message,
@@ -155,6 +177,7 @@ export function createAgent<
 
     return message;
   };
+  agent.getMessages = () => agent.getSnapshot().context.messages;
 
   agent.generateText = (opts) => agentGenerateText(agent, opts);
 
@@ -163,23 +186,32 @@ export function createAgent<
   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 observation = {
-      ...observationInput,
+      prevState,
+      event,
+      state,
       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',
@@ -188,6 +220,7 @@ export function createAgent<
 
     return observation;
   };
+  agent.getObservations = () => agent.getSnapshot().context.observations;
 
   agent.addPlan = (plan) => {
     agent.send({
@@ -195,8 +228,9 @@ export function createAgent<
       plan,
     });
   };
+  agent.getPlans = () => agent.getSnapshot().context.plans;
 
-  agent.interact = (actorRef, getInput) => {
+  agent.interact = ((actorRef, getInput) => {
     let prevState: ObservedState | undefined = undefined;
     let subscribed = true;
 
@@ -234,7 +268,8 @@ export function createAgent<
           event: inspEvent.event,
           prevState,
           state: inspEvent.snapshot as any,
-        };
+          machine: (actorRef as any).src,
+        } satisfies AgentObservationInput;
 
         await handleObservation(observationInput);
       },
@@ -246,6 +281,7 @@ export function createAgent<
         prevState: undefined,
         event: { type: '' }, // TODO: unknown events?
         state: actorRef.getSnapshot(),
+        machine: (actorRef as any).src,
       });
     }
 
@@ -254,7 +290,9 @@ export function createAgent<
         subscribed = false;
       }, // TODO: make this actually unsubscribe
     };
-  };
+  }) as typeof agent.interact;
+
+  agent.types = {} as any;
 
   agent.start();
 

package/src/decision.ts

@@ -1,17 +1,18 @@
 import { AnyMachineSnapshot, fromPromise } from 'xstate';
 import {
-  Agent,
+  AnyAgent,
   AgentDecideOptions,
   AgentDecisionLogic,
   AgentDecisionInput,
   AgentPlanner,
+  AgentPlan,
 } from './types';
 import { simplePlanner } from './planners/simplePlanner';
 
-export async function agentDecide<T extends Agent<any>>(
+export async function agentDecide<T extends AnyAgent>(
   agent: T,
   options: AgentDecideOptions
-) {
+): Promise<AgentPlan<any> | undefined> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
@@ -44,9 +45,9 @@ export async function agentDecide<T extends Agent<any>>(
 }
 
 export function fromDecision(
-  agent: Agent<any>,
+  agent: AnyAgent,
   defaultInput?: AgentDecisionInput
-) {
+): AgentDecisionLogic<any> {
   return fromPromise(async ({ input, self }) => {
     const parentRef = self._parent;
     if (!parentRef) {

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 { type 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,18 +93,26 @@ 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({
+    toolChoice: 'required',
     ...input,
     prompt,
+    messages,
     tools: toolMap,
-    toolChoice: 'required',
   });
 
   const singleResult = result.toolResults[0];
@@ -117,11 +126,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

@@ -1,17 +1,13 @@
-import type {
-  CoreMessage,
-  CoreTool,
-  GenerateTextResult,
-  StreamTextResult,
-} from 'ai';
+import type { CoreMessage, CoreTool, GenerateTextResult } from 'ai';
 import {
-  Agent,
   AgentGenerateTextOptions,
+  AgentGenerateTextResult,
   AgentStreamTextOptions,
+  AgentStreamTextResult,
+  AnyAgent,
 } from './types';
 import { defaultTextTemplate } from './templates/defaultText';
 import {
-  AnyMachineSnapshot,
   ObservableActorLogic,
   Observer,
   PromiseActorLogic,
@@ -29,15 +25,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,14 +45,16 @@ async function getMessages(
   return messages;
 }
 
-export async function agentGenerateText<T extends Agent<any>>(
+export async function agentGenerateText<T extends AnyAgent>(
   agent: T,
   options: AgentGenerateTextOptions
-) {
+): Promise<AgentGenerateTextResult> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
+    correlationId: options.correlationId ?? randomId(),
   };
+  // Generate a correlation ID if one is not provided
   const template = resolvedOptions.template ?? defaultTextTemplate;
   // TODO: check if messages was provided instead
   const id = randomId();
@@ -79,6 +75,8 @@ export async function agentGenerateText<T extends Agent<any>>(
     role: 'user',
     content: promptWithContext,
     timestamp: Date.now(),
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
   });
 
   const result = await agent.adapter.generateText({
@@ -94,18 +92,25 @@ export async function agentGenerateText<T extends Agent<any>>(
     timestamp: Date.now(),
     responseId: id,
     result,
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
   });
 
-  return result;
+  return {
+    ...result,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
+    correlationId: resolvedOptions.correlationId,
+  };
 }
 
 export async function agentStreamText(
-  agent: Agent<any>,
+  agent: AnyAgent,
   options: AgentStreamTextOptions
-): Promise<StreamTextResult<any>> {
+): Promise<AgentStreamTextResult> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
+    correlationId: options.correlationId ?? randomId(),
   };
   const template = resolvedOptions.template ?? defaultTextTemplate;
 
@@ -127,6 +132,8 @@ export async function agentStreamText(
     content: promptWithContext,
     id,
     timestamp: Date.now(),
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
   });
 
   const result = await agent.adapter.streamText({
@@ -146,33 +153,35 @@ export async function agentStreamText(
           usage: res.usage,
           warnings: res.warnings,
           rawResponse: res.rawResponse,
+          roundtrips: [], // TODO: how do we get this information?
         },
         content: res.text,
         id: randomId(),
         timestamp: Date.now(),
         responseId: id,
+        correlationId: resolvedOptions.correlationId,
+        parentCorrelationId: resolvedOptions.parentCorrelationId,
       });
     },
   });
 
-  return result;
+  return {
+    ...result,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
+    correlationId: resolvedOptions.correlationId,
+  } as unknown as AgentStreamTextResult; // TODO: fix
 }
 
-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 +190,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 +217,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,
     });
   });
 }