@statelyai/agent 1.0.0-beta.1 → 1.1.2

package/dist/index.mjs

@@ -1,21 +1,43 @@
 // src/agent.ts
 import {
   createActor,
-  fromTransition,
-  toObserver as toObserver2
+  fromTransition
 } from "xstate";
 
 // src/planners/simplePlanner.ts
 import { tool } from "ai";
 
 // src/utils.ts
+import hash from "object-hash";
 function getAllTransitions(state) {
   const nodes = state._nodes;
-  const transitions = nodes.map((node) => [...node.transitions.values()]).flat(2).map((transition) => ({
-    ...transition,
-    guard: typeof transition.guard === "string" ? { type: transition.guard } : transition.guard
-    // TODO: fix
-  }));
+  const transitions = nodes.map((node) => [...node.transitions.values()]).map((nodeTransitions) => {
+    return nodeTransitions.map((nodeEventTransitions) => {
+      return nodeEventTransitions.map((transition) => {
+        return {
+          ...transition,
+          guard: typeof transition.guard === "string" ? { type: transition.guard } : transition.guard
+          // TODO: fix
+        };
+      });
+    });
+  }).flat(2);
+  return transitions;
+}
+function getAllMachineTransitions(stateNode) {
+  const transitions = [...stateNode.transitions.values()].map((nodeTransitions) => {
+    return nodeTransitions.map((transition) => {
+      return {
+        ...transition,
+        guard: typeof transition.guard === "string" ? { type: transition.guard } : transition.guard
+        // TODO: fix
+      };
+    });
+  }).flat(2);
+  for (const s of Object.values(stateNode.states)) {
+    const stateTransitions = getAllMachineTransitions(s);
+    transitions.push(...stateTransitions);
+  }
   return transitions;
 }
 function wrapInXml(tagName, content) {
@@ -26,6 +48,14 @@ function randomId() {
   const random = Math.random().toString(36).substring(2, 9);
   return timestamp + random;
 }
+var machineHashes = /* @__PURE__ */ new WeakMap();
+function getMachineHash(machine) {
+  if (machineHashes.has(machine)) return machineHashes.get(machine);
+  const transitions = getAllMachineTransitions(machine.root);
+  const machineHash = hash(transitions);
+  machineHashes.set(machine, machineHash);
+  return machineHash;
+}
 
 // src/templates/defaultText.ts
 var defaultTextTemplate = (data) => {
@@ -39,87 +69,6 @@ ${data.goal}
   `.trim();
 };
 
-// src/planners/simplePlanner.ts
-function getTransitions(state, machine) {
-  if (!machine) {
-    return [];
-  }
-  const resolvedState = machine.resolveState(state);
-  return getAllTransitions(resolvedState);
-}
-var simplePlannerPromptTemplate = (data) => {
-  return `
-${defaultTextTemplate(data)}
-
-Only make a single tool call to achieve the above goal.
-  `.trim();
-};
-async function simplePlanner(agent, input) {
-  const transitions = input.machine ? getTransitions(input.state, input.machine) : Object.entries(input.events).map(([eventType, { description }]) => ({
-    eventType,
-    description
-  }));
-  const filter = (eventType) => Object.keys(input.events).includes(eventType);
-  const functionNameMapping = {};
-  const toolTransitions = transitions.filter((t) => {
-    return filter(t.eventType);
-  }).map((t) => {
-    const name = t.eventType.replace(/\./g, "_");
-    functionNameMapping[name] = t.eventType;
-    return {
-      type: "function",
-      eventType: t.eventType,
-      description: t.description,
-      name
-    };
-  });
-  const toolMap = {};
-  for (const toolTransitionData of toolTransitions) {
-    const toolZodType = input.events?.[toolTransitionData.eventType];
-    if (!toolZodType) {
-      continue;
-    }
-    toolMap[toolTransitionData.name] = tool({
-      description: toolZodType?.description ?? toolTransitionData.description,
-      parameters: toolZodType,
-      execute: async (params) => {
-        const event = {
-          type: toolTransitionData.eventType,
-          ...params
-        };
-        return event;
-      }
-    });
-  }
-  const prompt = simplePlannerPromptTemplate({
-    context: input.state.context,
-    goal: input.goal
-  });
-  const result = await agent.generateText({
-    ...input,
-    prompt,
-    tools: toolMap,
-    toolChoice: "required"
-  });
-  const singleResult = result.toolResults[0];
-  if (!singleResult) {
-    console.warn("No tool call results returned");
-    return void 0;
-  }
-  return {
-    goal: input.goal,
-    state: input.state,
-    steps: [
-      {
-        event: singleResult.result
-      }
-    ],
-    nextEvent: singleResult.result,
-    sessionId: agent.sessionId,
-    timestamp: Date.now()
-  };
-}
-
 // src/text.ts
 import {
   fromObservable,
@@ -128,9 +77,7 @@ import {
 } from "xstate";
 async function getMessages(agent, prompt, options) {
   let messages = [];
-  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;
@@ -144,7 +91,8 @@ async function getMessages(agent, prompt, options) {
 async function agentGenerateText(agent, options) {
   const resolvedOptions = {
     ...agent.defaultOptions,
-    ...options
+    ...options,
+    correlationId: options.correlationId ?? randomId()
   };
   const template = resolvedOptions.template ?? defaultTextTemplate;
   const id = randomId();
@@ -158,7 +106,9 @@ async function agentGenerateText(agent, options) {
     id,
     role: "user",
     content: promptWithContext,
-    timestamp: Date.now()
+    timestamp: Date.now(),
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId
   });
   const result = await agent.adapter.generateText({
     ...resolvedOptions,
@@ -171,14 +121,21 @@ async function agentGenerateText(agent, options) {
     role: "assistant",
     timestamp: Date.now(),
     responseId: id,
-    result
+    result,
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId
   });
-  return result;
+  return {
+    ...result,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
+    correlationId: resolvedOptions.correlationId
+  };
 }
 async function agentStreamText(agent, options) {
   const resolvedOptions = {
     ...agent.defaultOptions,
-    ...options
+    ...options,
+    correlationId: options.correlationId ?? randomId()
   };
   const template = resolvedOptions.template ?? defaultTextTemplate;
   const id = randomId();
@@ -192,7 +149,9 @@ async function agentStreamText(agent, options) {
     role: "user",
     content: promptWithContext,
     id,
-    timestamp: Date.now()
+    timestamp: Date.now(),
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId
   });
   const result = await agent.adapter.streamText({
     ...resolvedOptions,
@@ -210,26 +169,33 @@ async function agentStreamText(agent, options) {
           toolResults: [],
           usage: res.usage,
           warnings: res.warnings,
-          rawResponse: res.rawResponse
+          rawResponse: res.rawResponse,
+          roundtrips: []
+          // TODO: how do we get this information?
         },
         content: res.text,
         id: randomId(),
         timestamp: Date.now(),
-        responseId: id
+        responseId: id,
+        correlationId: resolvedOptions.correlationId,
+        parentCorrelationId: resolvedOptions.parentCorrelationId
       });
     }
   });
-  return result;
+  return {
+    ...result,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
+    correlationId: resolvedOptions.correlationId
+  };
 }
 function fromTextStream(agent, defaultOptions) {
-  return fromObservable(({ input, self }) => {
-    const context = input.context === true ? (self._parent?.getSnapshot()).context : input.context;
+  return fromObservable(({ input }) => {
     const observers = /* @__PURE__ */ new Set();
     (async () => {
       const result = await agentStreamText(agent, {
         ...defaultOptions,
         ...input,
-        context
+        context: input.context
       });
       for await (const part of result.fullStream) {
         if (part.type === "text-delta") {
@@ -253,16 +219,102 @@ function fromTextStream(agent, defaultOptions) {
   });
 }
 function fromText(agent, defaultOptions) {
-  return fromPromise(async ({ input, self }) => {
-    const context = input.context === true ? (self._parent?.getSnapshot()).context : input.context;
+  return fromPromise(async ({ input }) => {
     return await agentGenerateText(agent, {
       ...input,
       ...defaultOptions,
-      context
+      context: input.context
     });
   });
 }
 
+// src/planners/simplePlanner.ts
+function getTransitions(state, machine) {
+  if (!machine) {
+    return [];
+  }
+  const resolvedState = machine.resolveState(state);
+  return getAllTransitions(resolvedState);
+}
+var simplePlannerPromptTemplate = (data) => {
+  return `
+${defaultTextTemplate(data)}
+
+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();
+};
+async function simplePlanner(agent, input) {
+  const transitions = input.machine ? getTransitions(input.state, input.machine) : Object.entries(input.events).map(([eventType, { description }]) => ({
+    eventType,
+    description
+  }));
+  const filter = (eventType) => Object.keys(input.events).includes(eventType);
+  const functionNameMapping = {};
+  const toolTransitions = transitions.filter((t) => {
+    return filter(t.eventType);
+  }).map((t) => {
+    const name = t.eventType.replace(/\./g, "_");
+    functionNameMapping[name] = t.eventType;
+    return {
+      type: "function",
+      eventType: t.eventType,
+      description: t.description,
+      name
+    };
+  });
+  const toolMap = {};
+  for (const toolTransitionData of toolTransitions) {
+    const toolZodType = input.events?.[toolTransitionData.eventType];
+    if (!toolZodType) {
+      continue;
+    }
+    toolMap[toolTransitionData.name] = tool({
+      description: toolZodType?.description ?? toolTransitionData.description,
+      parameters: toolZodType,
+      execute: async (params) => {
+        const event = {
+          type: toolTransitionData.eventType,
+          ...params
+        };
+        return event;
+      }
+    });
+  }
+  if (!Object.keys(toolMap).length) {
+    return void 0;
+  }
+  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
+  });
+  const singleResult = result.toolResults[0];
+  if (!singleResult) {
+    console.warn("No tool call results returned");
+    return void 0;
+  }
+  return {
+    goal: input.goal,
+    state: input.state,
+    execute: async (state) => {
+      if (JSON.stringify(state) === JSON.stringify(input.state)) {
+        return singleResult.result;
+      }
+      return void 0;
+    },
+    nextEvent: singleResult.result,
+    sessionId: agent.sessionId,
+    timestamp: Date.now()
+  };
+}
+
 // src/decision.ts
 import { fromPromise as fromPromise2 } from "xstate";
 async function agentDecide(agent, options) {
@@ -377,18 +429,19 @@ var agentLogic = fromTransition(
     }
     return state;
   },
-  {
+  () => ({
     feedback: [],
     messages: [],
     observations: [],
     plans: []
-  }
+  })
 );
 function createAgent({
   name,
   description,
   model,
   events,
+  context,
   planner = simplePlanner,
   stringify = JSON.stringify,
   getMemory,
@@ -396,7 +449,6 @@ function createAgent({
   adapter = vercelAdapter,
   ...generateTextOptions
 }) {
-  const messageHistoryListeners = [];
   const agent = createActor(logic);
   agent.events = events;
   agent.model = model;
@@ -409,7 +461,7 @@ function createAgent({
   };
   agent.memory = getMemory ? getMemory(agent) : void 0;
   agent.onMessage = (callback) => {
-    messageHistoryListeners.push(toObserver2(callback));
+    agent.on("message", (ev) => callback(ev.message));
   };
   agent.decide = (opts) => {
     return agentDecide(agent, opts);
@@ -419,7 +471,8 @@ function createAgent({
       ...messageInput,
       id: messageInput.id ?? randomId(),
       timestamp: messageInput.timestamp ?? Date.now(),
-      sessionId: agent.sessionId
+      sessionId: agent.sessionId,
+      correlationId: messageInput.correlationId ?? randomId()
     };
     agent.send({
       type: "agent.message",
@@ -427,11 +480,14 @@ function createAgent({
     });
     return message;
   };
+  agent.getMessages = () => agent.getSnapshot().context.messages;
   agent.generateText = (opts) => agentGenerateText(agent, opts);
   agent.streamText = (opts) => agentStreamText(agent, opts);
   agent.addFeedback = (feedbackInput) => {
     const feedback = {
       ...feedbackInput,
+      attributes: { ...feedbackInput.attributes },
+      reward: feedbackInput.reward ?? 0,
       timestamp: feedbackInput.timestamp ?? Date.now(),
       sessionId: agent.sessionId
     };
@@ -441,12 +497,17 @@ function createAgent({
     });
     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()
+      timestamp: observationInput.timestamp ?? Date.now(),
+      machineHash: observationInput.machine ? getMachineHash(observationInput.machine) : void 0
     };
     agent.send({
       type: "agent.observe",
@@ -454,12 +515,14 @@ function createAgent({
     });
     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 = void 0;
     let subscribed = true;
@@ -486,7 +549,8 @@ function createAgent({
         const observationInput = {
           event: inspEvent.event,
           prevState,
-          state: inspEvent.snapshot
+          state: inspEvent.snapshot,
+          machine: actorRef.src
         };
         await handleObservation(observationInput);
       }
@@ -496,7 +560,8 @@ function createAgent({
         prevState: void 0,
         event: { type: "" },
         // TODO: unknown events?
-        state: actorRef.getSnapshot()
+        state: actorRef.getSnapshot(),
+        machine: actorRef.src
       });
     }
     return {
@@ -506,12 +571,11 @@ function createAgent({
       // TODO: make this actually unsubscribe
     };
   };
+  agent.types = {};
   agent.start();
   return agent;
 }
 export {
-  agentDecide,
-  agentGenerateText,
   createAgent,
   fromDecision,
   fromText,

package/examples/chatbot.ts

@@ -13,20 +13,18 @@ const agent = createAgent({
     }),
     'agent.endConversation': z.object({}).describe('Stop the conversation'),
   },
+  context: {
+    userMessage: z.string(),
+  },
 });
 
 const machine = setup({
-  types: {
-    context: {} as {
-      conversation: string[];
-    },
-    events: agent.eventTypes,
-  },
+  types: agent.types,
   actors: { agent: fromDecision(agent), getFromTerminal },
 }).createMachine({
   initial: 'listening',
   context: {
-    conversation: [],
+    userMessage: '',
   },
   states: {
     listening: {
@@ -35,8 +33,7 @@ const machine = setup({
         input: 'User:',
         onDone: {
           actions: assign({
-            conversation: (x) =>
-              x.context.conversation.concat('User: ' + x.event.output),
+            userMessage: (x) => x.event.output,
           }),
           target: 'responding',
         },
@@ -47,20 +44,15 @@ const machine = setup({
         src: 'agent',
         input: (x) => ({
           context: {
-            conversation: x.context.conversation,
+            userMessage: 'User says: ' + x.context.userMessage,
           },
+          messages: agent.getMessages(),
           goal: 'Respond to the user, unless they want to end the conversation.',
         }),
       },
       on: {
         'agent.respond': {
-          actions: [
-            assign({
-              conversation: (x) =>
-                x.context.conversation.concat('Assistant: ' + x.event.response),
-            }),
-            log((x) => `Agent: ${x.event.response}`),
-          ],
+          actions: [log((x) => `Agent: ${x.event.response}`)],
           target: 'listening',
         },
         'agent.endConversation': 'finished',

package/examples/cot.ts

@@ -17,16 +17,14 @@ const agent = createAgent({
       answer: z.string().describe('The answer to the question'),
     }),
   },
+  context: {
+    question: z.string().nullable(),
+    thought: z.string().nullable(),
+  },
 });
 
 const machine = setup({
-  types: {
-    context: {} as {
-      question: string | null;
-      thought: string | null;
-    },
-    events: agent.eventTypes,
-  },
+  types: agent.types,
   actors: { agent: fromDecision(agent), getFromTerminal },
 }).createMachine({
   initial: 'asking',

package/examples/email.ts

@@ -19,7 +19,7 @@ const agent = createAgent({
 
 const machine = setup({
   types: {
-    events: agent.eventTypes,
+    events: agent.types.events,
     input: {} as {
       email: string;
       instructions: string;
@@ -50,7 +50,7 @@ const machine = setup({
             instructions: x.context.instructions,
             clarifications: x.context.clarifications,
           },
-          messages: agent.select((ctx) => ctx.messages),
+          messages: agent.getMessages(),
           goal: 'Respond to the email given the instructions and the provided clarifications. If not enough information is provided, ask for clarification. Otherwise, if you are absolutely sure that there is no ambiguous or missing information, create and submit a response email.',
         }),
       },

package/examples/example.ts

@@ -18,7 +18,7 @@ const agent = createAgent({
 
 const machine = setup({
   types: {
-    events: agent.eventTypes,
+    events: agent.types.events,
   },
   actors: { agent: fromDecision(agent), summarizer: fromText(agent) },
 }).createMachine({
@@ -32,11 +32,11 @@ const machine = setup({
       invoke: [
         {
           src: 'summarizer',
-          input: {
-            context: true,
+          input: (x) => ({
+            context: x.context,
             prompt:
               'Summarize the patient visit in a single sentence. The summary should be in English.',
-          },
+          }),
           onDone: {
             actions: assign({
               englishSummary: ({ event }) => event.output.text,
@@ -45,11 +45,11 @@ const machine = setup({
         },
         {
           src: 'summarizer',
-          input: {
-            context: true,
+          input: (x) => ({
+            context: x.context,
             prompt:
               'Summarize the patient visit in a single sentence. The summary should be in Spanish.',
-          },
+          }),
           onDone: {
             actions: assign({
               spanishSummary: ({ event }) => event.output.text,

package/examples/goal.ts

@@ -25,7 +25,7 @@ const machine = setup({
       question: string | null;
       goal: string | null;
     },
-    events: agent.eventTypes,
+    events: agent.types.events,
   },
   actors: { decider, getFromTerminal },
 }).createMachine({

package/examples/joke.ts

@@ -67,19 +67,17 @@ const agent = createAgent({
       .describe('Explains why the joke was irrelevant'),
     'agent.markAsRelevant': z.object({}).describe('The joke was relevant'),
   },
+  context: {
+    topic: z.string().describe('The topic for the joke'),
+    jokes: z.array(z.string()).describe('The jokes told so far'),
+    desire: z.string().nullable().describe('The user desire'),
+    lastRating: z.number().nullable().describe('The last joke rating'),
+    loader: z.string().nullable().describe('The loader text'),
+  },
 });
 
 const jokeMachine = setup({
-  types: {
-    context: {} as {
-      topic: string;
-      jokes: string[];
-      desire: string | null;
-      lastRating: number | null;
-      loader: string | null;
-    },
-    events: agent.eventTypes,
-  },
+  types: agent.types,
   actors: {
     agent: fromDecision(agent),
     loader,

package/examples/number.ts

@@ -20,7 +20,7 @@ const machine = setup({
       previousGuesses: number[];
       answer: number | null;
     },
-    events: agent.eventTypes,
+    events: agent.types.events,
   },
   actors: {
     agent: fromDecision(agent),

package/examples/raffle.ts

@@ -27,7 +27,7 @@ const machine = setup({
       lastInput: string | null;
       entries: string[];
     },
-    events: {} as typeof agent.eventTypes | { type: 'draw' },
+    events: {} as typeof agent.types.events | { type: 'draw' },
   },
   actors: { agent: fromDecision(agent), getFromTerminal },
 }).createMachine({