@statelyai/agent 1.1.6 → 2.0.0-next.1

package/dist/index.mjs

@@ -1,11 +1,11 @@
 // src/agent.ts
 import {
-  createActor,
+  Actor,
   fromTransition
 } from "xstate";
 
-// src/planners/simplePlanner.ts
-import { tool } from "ai";
+// src/planners/simple.ts
+import { generateText as generateText2 } from "ai";
 
 // src/utils.ts
 import hash from "object-hash";
@@ -43,10 +43,10 @@ function getAllMachineTransitions(stateNode) {
 function wrapInXml(tagName, content) {
   return `<${tagName}>${content}</${tagName}>`;
 }
-function randomId() {
+function randomId(prefix) {
   const timestamp = Date.now().toString(36);
   const random = Math.random().toString(36).substring(2, 9);
-  return timestamp + random;
+  return `${prefix || ""}${timestamp}${random}`;
 }
 var machineHashes = /* @__PURE__ */ new WeakMap();
 function getMachineHash(machine) {
@@ -56,6 +56,26 @@ function getMachineHash(machine) {
   machineHashes.set(machine, machineHash);
   return machineHash;
 }
+function isActorRef(actorRefLike) {
+  return "src" in actorRefLike && "system" in actorRefLike && "sessionId" in actorRefLike;
+}
+function getTransitions(state, machine) {
+  if (!machine) {
+    return [];
+  }
+  const resolvedState = machine.resolveState({
+    ...state,
+    // Need this property defined to make TS happy
+    context: state.context
+  });
+  return getAllTransitions(resolvedState);
+}
+function isMachineActor(actor) {
+  return "src" in actor && typeof actor.src === "object" && actor.src !== null && "definition" in actor.src;
+}
+
+// src/planners/simple.ts
+import { getNextSnapshot } from "xstate";
 
 // src/templates/defaultText.ts
 var defaultTextTemplate = (data) => {
@@ -70,6 +90,10 @@ ${data.goal}
 };
 
 // src/text.ts
+import {
+  generateText,
+  streamText
+} from "ai";
 import {
   fromObservable,
   fromPromise,
@@ -88,120 +112,26 @@ async function getMessages(agent, prompt, options) {
   });
   return messages;
 }
-async function agentGenerateText(agent, options) {
-  const resolvedOptions = {
-    ...agent.defaultOptions,
-    ...options,
-    correlationId: options.correlationId ?? randomId()
-  };
-  const template = resolvedOptions.template ?? defaultTextTemplate;
-  const id = randomId();
-  const goal = typeof resolvedOptions.prompt === "string" ? resolvedOptions.prompt : await resolvedOptions.prompt(agent);
-  const promptWithContext = template({
-    goal,
-    context: resolvedOptions.context
-  });
-  const messages = await getMessages(agent, promptWithContext, resolvedOptions);
-  agent.addMessage({
-    id,
-    role: "user",
-    content: promptWithContext,
-    timestamp: Date.now(),
-    correlationId: resolvedOptions.correlationId,
-    parentCorrelationId: resolvedOptions.parentCorrelationId
-  });
-  const result = await agent.adapter.generateText({
-    ...resolvedOptions,
-    prompt: void 0,
-    messages
-  });
-  agent.addMessage({
-    content: result.text,
-    id,
-    role: "assistant",
-    timestamp: Date.now(),
-    responseId: id,
-    result,
-    correlationId: resolvedOptions.correlationId,
-    parentCorrelationId: resolvedOptions.parentCorrelationId
-  });
-  return {
-    ...result,
-    parentCorrelationId: resolvedOptions.parentCorrelationId,
-    correlationId: resolvedOptions.correlationId
-  };
-}
-async function agentStreamText(agent, options) {
-  const resolvedOptions = {
-    ...agent.defaultOptions,
-    ...options,
-    correlationId: options.correlationId ?? randomId()
-  };
-  const template = resolvedOptions.template ?? defaultTextTemplate;
-  const id = randomId();
-  const goal = typeof resolvedOptions.prompt === "string" ? resolvedOptions.prompt : await resolvedOptions.prompt(agent);
-  const promptWithContext = template({
-    goal,
-    context: resolvedOptions.context
-  });
-  const messages = await getMessages(agent, promptWithContext, resolvedOptions);
-  agent.addMessage({
-    role: "user",
-    content: promptWithContext,
-    id,
-    timestamp: Date.now(),
-    correlationId: resolvedOptions.correlationId,
-    parentCorrelationId: resolvedOptions.parentCorrelationId
-  });
-  const result = await agent.adapter.streamText({
-    ...resolvedOptions,
-    prompt: void 0,
-    messages,
-    onFinish: async (res) => {
-      agent.addMessage({
-        role: "assistant",
-        result: {
-          text: res.text,
-          finishReason: res.finishReason,
-          logprobs: void 0,
-          responseMessages: [],
-          toolCalls: [],
-          toolResults: [],
-          usage: res.usage,
-          warnings: res.warnings,
-          rawResponse: res.rawResponse,
-          roundtrips: [],
-          // TODO: how do we get this information?,
-          steps: res.steps,
-          response: res.response,
-          experimental_providerMetadata: res.experimental_providerMetadata
-        },
-        content: res.text,
-        id: randomId(),
-        timestamp: Date.now(),
-        responseId: id,
-        correlationId: resolvedOptions.correlationId,
-        parentCorrelationId: resolvedOptions.parentCorrelationId
-      });
-    }
-  });
-  return {
-    ...result,
-    textStream: result.textStream,
-    fullStream: result.fullStream,
-    parentCorrelationId: resolvedOptions.parentCorrelationId,
-    correlationId: resolvedOptions.correlationId
-  };
-}
-function fromTextStream(agent, defaultOptions) {
+function fromTextStream(agent, options) {
+  const template = options?.template ?? defaultTextTemplate;
   return fromObservable(({ input }) => {
     const observers = /* @__PURE__ */ new Set();
     (async () => {
-      const result = await agentStreamText(agent, {
-        ...defaultOptions,
-        ...input,
+      const model = input.model ? agent.wrap(input.model) : agent.model;
+      const goal = typeof input.prompt === "string" ? input.prompt : await input.prompt(agent);
+      const promptWithContext = template({
+        goal,
         context: input.context
       });
+      const messages = await getMessages(agent, promptWithContext, input);
+      const result = await streamText({
+        ...options,
+        ...input,
+        prompt: void 0,
+        // overwritten by messages
+        model,
+        messages
+      });
       for await (const part of result.fullStream) {
         if (part.type === "text-delta") {
           observers.forEach((observer) => {
@@ -223,32 +153,97 @@ function fromTextStream(agent, defaultOptions) {
     };
   });
 }
-function fromText(agent, defaultOptions) {
+function fromText(agent, options) {
+  const resolvedOptions = {
+    ...agent.defaultOptions,
+    ...options
+  };
+  const template = resolvedOptions.template ?? defaultTextTemplate;
   return fromPromise(async ({ input }) => {
-    return await agentGenerateText(agent, {
-      ...input,
-      ...defaultOptions,
+    const goal = typeof input.prompt === "string" ? input.prompt : await input.prompt(agent);
+    const promptWithContext = template({
+      goal,
       context: input.context
     });
+    const messages = await getMessages(agent, promptWithContext, input);
+    const model = input.model ? agent.wrap(input.model) : agent.model;
+    return await generateText({
+      ...input,
+      ...options,
+      prompt: void 0,
+      messages,
+      model
+    });
   });
 }
 
-// src/planners/simplePlanner.ts
-function getTransitions(state, machine) {
-  if (!machine) {
-    return [];
+// src/decide.ts
+import { fromPromise as fromPromise2 } from "xstate";
+import { tool } from "ai";
+async function agentDecide(agent, options) {
+  const resolvedOptions = {
+    ...agent.defaultOptions,
+    ...options
+  };
+  const {
+    planner = simplePlanner,
+    goal,
+    events = agent.events,
+    state,
+    machine,
+    model = agent.model,
+    messages,
+    ...otherPlanInput
+  } = resolvedOptions;
+  let attempts = 0;
+  const maxAttempts = resolvedOptions.maxAttempts ?? 2;
+  let plan;
+  while (attempts++ < maxAttempts) {
+    plan = await planner(agent, {
+      model,
+      goal,
+      events,
+      state,
+      machine,
+      messages,
+      // TODO: fix UIMessage thing
+      ...otherPlanInput
+    });
+    if (plan?.nextEvent) {
+      agent.addPlan(plan);
+      await resolvedOptions.execute?.(plan.nextEvent);
+    }
   }
-  const resolvedState = machine.resolveState(state);
-  return getAllTransitions(resolvedState);
+  return plan;
 }
-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) {
+function fromDecision(agent, defaultInput) {
+  return fromPromise2(async ({ input, self }) => {
+    const parentRef = self._parent;
+    if (!parentRef) {
+      return;
+    }
+    const snapshot = parentRef.getSnapshot();
+    const inputObject = typeof input === "string" ? { goal: input } : input;
+    const resolvedInput = {
+      ...defaultInput,
+      ...inputObject
+    };
+    const state = {
+      value: snapshot.value,
+      context: resolvedInput.context
+    };
+    const plan = await agentDecide(agent, {
+      machine: parentRef.logic,
+      state,
+      execute: async (event) => {
+        parentRef.send(event);
+      },
+      ...resolvedInput
+    });
+    return plan;
+  });
+}
+function getToolMap(_agent, input) {
   const transitions = input.machine ? getTransitions(input.state, input.machine) : Object.entries(input.events).map(([eventType, { description }]) => ({
     eventType,
     description
@@ -288,17 +283,56 @@ async function simplePlanner(agent, input) {
   if (!Object.keys(toolMap).length) {
     return void 0;
   }
+  return toolMap;
+}
+
+// src/planners/simple.ts
+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 toolMap = getToolMap(agent, input);
+  if (!toolMap) {
+    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,
+  const model = input.model ? agent.wrap(input.model) : agent.model;
+  const {
+    state,
+    machine,
+    previousPlan,
+    events,
+    goal,
+    model: _,
+    ...rest
+  } = input;
+  const machineState = input.machine ? input.machine.resolveState({
+    ...input.state,
+    context: input.state.context
+  }) : void 0;
+  const result = await generateText2({
+    ...rest,
+    system: input.system ?? agent.description,
+    model,
     messages,
-    tools: toolMap
+    tools: toolMap,
+    toolChoice: input.toolChoice ?? "required"
+  });
+  result.response.messages.forEach((m) => {
+    const message = m;
+    agent.addMessage({
+      ...message,
+      id: randomId(),
+      timestamp: Date.now()
+    });
   });
   const singleResult = result.toolResults[0];
   if (!singleResult) {
@@ -306,89 +340,94 @@ async function simplePlanner(agent, input) {
     return void 0;
   }
   return {
+    planner: "simple",
     goal: input.goal,
-    state: input.state,
-    execute: async (state) => {
-      if (JSON.stringify(state) === JSON.stringify(input.state)) {
-        return singleResult.result;
-      }
-      return void 0;
-    },
+    goalState: input.state,
     nextEvent: singleResult.result,
-    sessionId: agent.sessionId,
-    timestamp: Date.now()
+    episodeId: agent.episodeId,
+    timestamp: Date.now(),
+    paths: [
+      {
+        state: void 0,
+        steps: [
+          {
+            event: singleResult.result,
+            state: machine && machineState ? getNextSnapshot(machine, machineState, singleResult.result) : void 0
+          }
+        ]
+      }
+    ]
   };
 }
 
-// src/decision.ts
-import { fromPromise as fromPromise2 } from "xstate";
-async function agentDecide(agent, options) {
-  const resolvedOptions = {
-    ...agent.defaultOptions,
-    ...options
-  };
-  const {
-    planner = simplePlanner,
-    goal,
-    events = agent.events,
-    state,
-    machine,
-    model = agent.model,
-    ...otherPlanInput
-  } = resolvedOptions;
-  const plan = await planner(agent, {
-    model,
-    goal,
-    events,
-    state,
-    machine,
-    ...otherPlanInput
-  });
-  if (plan?.nextEvent) {
-    agent.addPlan(plan);
-    await resolvedOptions.execute?.(plan.nextEvent);
-  }
-  return plan;
-}
-function fromDecision(agent, defaultInput) {
-  return fromPromise2(async ({ input, self }) => {
-    const parentRef = self._parent;
-    if (!parentRef) {
-      return;
+// src/agent.ts
+import {
+  experimental_wrapLanguageModel
+} from "ai";
+
+// src/middleware.ts
+function createAgentMiddleware(agent) {
+  const middleware = {
+    transformParams: async ({ params }) => {
+      return params;
+    },
+    wrapGenerate: async ({ doGenerate, params }) => {
+      const id = randomId();
+      params.prompt.forEach((message) => {
+        agent.addMessage({
+          id,
+          ...message,
+          timestamp: Date.now()
+        });
+      });
+      const result = await doGenerate();
+      return result;
+    },
+    wrapStream: async ({ doStream, params }) => {
+      const id = randomId();
+      params.prompt.forEach((message) => {
+        message.content;
+        agent.addMessage({
+          id,
+          ...message,
+          timestamp: Date.now()
+        });
+      });
+      const { stream, ...rest } = await doStream();
+      let generatedText = "";
+      const transformStream = new TransformStream({
+        transform(chunk, controller) {
+          if (chunk.type === "text-delta") {
+            generatedText += chunk.textDelta;
+          }
+          controller.enqueue(chunk);
+        },
+        flush() {
+          const content = [];
+          if (generatedText) {
+            content.push({
+              type: "text",
+              text: generatedText
+            });
+          }
+          agent.addMessage({
+            id: randomId(),
+            timestamp: Date.now(),
+            role: "assistant",
+            content,
+            responseId: id
+          });
+        }
+      });
+      return {
+        stream: stream.pipeThrough(transformStream),
+        ...rest
+      };
     }
-    const snapshot = parentRef.getSnapshot();
-    const inputObject = typeof input === "string" ? { goal: input } : input;
-    const resolvedInput = {
-      ...defaultInput,
-      ...inputObject
-    };
-    const contextToInclude = resolvedInput.context === true ? (
-      // include entire context
-      parentRef.getSnapshot().context
-    ) : resolvedInput.context;
-    const state = {
-      value: snapshot.value,
-      context: contextToInclude
-    };
-    const plan = await agentDecide(agent, {
-      machine: parentRef.logic,
-      state,
-      execute: async (event) => {
-        parentRef.send(event);
-      },
-      ...resolvedInput
-    });
-    return plan;
-  });
+  };
+  return middleware;
 }
 
-// src/adapters/vercel.ts
-import { generateText, streamText } from "ai";
-var vercelAdapter = {
-  generateText,
-  streamText
-};
-
 // src/agent.ts
 var agentLogic = fromTransition(
   (state, event, { emit }) => {
@@ -429,8 +468,10 @@ var agentLogic = fromTransition(
         });
         break;
       }
-      default:
+      default: {
+        console.warn("Unrecognized event", event);
         break;
+      }
     }
     return state;
   },
@@ -442,111 +483,149 @@ var agentLogic = fromTransition(
   })
 );
 function createAgent({
-  name,
+  id,
   description,
   model,
   events,
   context,
   planner = simplePlanner,
-  stringify = JSON.stringify,
-  getMemory,
-  logic = agentLogic,
-  adapter = vercelAdapter,
-  ...generateTextOptions
+  logic = agentLogic
 }) {
-  const agent = createActor(logic);
-  agent.events = events;
-  agent.model = model;
-  agent.name = name;
-  agent.description = description;
-  agent.adapter = adapter;
-  agent.defaultOptions = { ...generateTextOptions, model };
-  agent.select = (selector) => {
-    return selector(agent.getSnapshot().context);
-  };
-  agent.memory = getMemory ? getMemory(agent) : void 0;
-  agent.onMessage = (callback) => {
-    agent.on("message", (ev) => callback(ev.message));
-  };
-  agent.decide = (opts) => {
-    return agentDecide(agent, opts);
-  };
-  agent.addMessage = (messageInput) => {
+  return new Agent({
+    id,
+    context,
+    events,
+    description,
+    planner,
+    model,
+    logic
+  });
+}
+var Agent = class extends Actor {
+  // todo
+  constructor({
+    logic = agentLogic,
+    id,
+    name,
+    description,
+    model,
+    events,
+    context,
+    planner = simplePlanner
+  }) {
+    super(logic);
+    this.model = model;
+    this.episodeId = id ?? randomId();
+    this.name = name;
+    this.description = description;
+    this.events = events;
+    this.context = context;
+    this.planner = planner;
+    this.types = {};
+    this.start();
+  }
+  /**
+   * Called whenever the agent (LLM assistant) receives or sends a message.
+   */
+  onMessage(fn) {
+    return this.on("message", (ev) => fn(ev.message));
+  }
+  /**
+   * Retrieves messages from the agent's short-term (local) memory.
+   */
+  addMessage(messageInput) {
     const message = {
       ...messageInput,
       id: messageInput.id ?? randomId(),
       timestamp: messageInput.timestamp ?? Date.now(),
-      sessionId: agent.sessionId,
-      correlationId: messageInput.correlationId ?? randomId()
+      episodeId: this.episodeId
     };
-    agent.send({
+    this.send({
       type: "agent.message",
       message
     });
     return message;
-  };
-  agent.getMessages = () => agent.getSnapshot().context.messages;
-  agent.generateText = (opts) => agentGenerateText(agent, opts);
-  agent.streamText = (opts) => agentStreamText(agent, opts);
-  agent.addFeedback = (feedbackInput) => {
+  }
+  getMessages() {
+    return this.getSnapshot().context.messages;
+  }
+  addFeedback(feedbackInput) {
     const feedback = {
       ...feedbackInput,
       attributes: { ...feedbackInput.attributes },
       reward: feedbackInput.reward ?? 0,
       timestamp: feedbackInput.timestamp ?? Date.now(),
-      sessionId: agent.sessionId
+      episodeId: this.episodeId
     };
-    agent.send({
+    this.send({
       type: "agent.feedback",
       feedback
     });
     return feedback;
-  };
-  agent.getFeedback = () => agent.getSnapshot().context.feedback;
-  agent.addObservation = (observationInput) => {
+  }
+  /**
+   * Retrieves feedback from the agent's short-term (local) memory.
+   */
+  getFeedback() {
+    return this.getSnapshot().context.feedback;
+  }
+  addObservation(observationInput) {
     const { prevState, event, state } = observationInput;
     const observation = {
       prevState,
       event,
       state,
       id: observationInput.id ?? randomId(),
-      sessionId: agent.sessionId,
+      episodeId: this.episodeId,
       timestamp: observationInput.timestamp ?? Date.now(),
       machineHash: observationInput.machine ? getMachineHash(observationInput.machine) : void 0
     };
-    agent.send({
+    this.send({
       type: "agent.observe",
       observation
     });
     return observation;
-  };
-  agent.getObservations = () => agent.getSnapshot().context.observations;
-  agent.addPlan = (plan) => {
-    agent.send({
+  }
+  /**
+   * Retrieves observations from the agent's short-term (local) memory.
+   */
+  getObservations() {
+    return this.getSnapshot().context.observations;
+  }
+  addPlan(plan) {
+    this.send({
       type: "agent.plan",
       plan
     });
-  };
-  agent.getPlans = () => agent.getSnapshot().context.plans;
-  agent.interact = (actorRef, getInput) => {
+  }
+  /**
+   * Retrieves strategies from the agent's short-term (local) memory.
+   */
+  getPlans() {
+    return this.getSnapshot().context.plans;
+  }
+  interact(actorRef, getInput) {
+    const actorRefCheck = isActorRef(actorRef) && actorRef.src;
+    const machine = isMachineActor(actorRef) ? actorRef.src : void 0;
     let prevState = void 0;
     let subscribed = true;
+    const agent = this;
     async function handleObservation(observationInput) {
       const observation = agent.addObservation(observationInput);
       const input = getInput?.(observation);
       if (input) {
-        await agentDecide(agent, {
-          machine: actorRef.src,
+        const res = await agentDecide(agent, {
+          machine,
           state: observation.state,
-          execute: async (event) => {
-            actorRef.send(event);
-          },
           ...input
         });
+        if (res?.nextEvent) {
+          actorRef.send(res.nextEvent);
+        }
       }
       prevState = observationInput.state;
     }
-    actorRef.system.inspect({
+    const sub = actorRefCheck ? actorRef.system.inspect({
       next: async (inspEvent) => {
         if (!subscribed || inspEvent.actorRef !== actorRef || inspEvent.type !== "@xstate.snapshot") {
           return;
@@ -559,7 +638,7 @@ function createAgent({
         };
         await handleObservation(observationInput);
       }
-    });
+    }) : void 0;
     if (actorRef._processingStatus === 1) {
       handleObservation({
         prevState: void 0,
@@ -571,15 +650,50 @@ function createAgent({
     }
     return {
       unsubscribe: () => {
+        sub?.unsubscribe();
         subscribed = false;
       }
-      // TODO: make this actually unsubscribe
     };
-  };
-  agent.types = {};
-  agent.start();
-  return agent;
-}
+  }
+  observe(actorRef) {
+    let prevState = actorRef.getSnapshot();
+    const actorRefCheck = isActorRef(actorRef);
+    const sub = actorRefCheck ? actorRef.system.inspect({
+      next: async (inspEvent) => {
+        if (inspEvent.actorRef !== actorRef || inspEvent.type !== "@xstate.snapshot") {
+          return;
+        }
+        const observationInput = {
+          event: inspEvent.event,
+          prevState,
+          state: inspEvent.snapshot,
+          machine: actorRef.src
+        };
+        prevState = observationInput.state;
+        this.addObservation(observationInput);
+      }
+    }) : void 0;
+    return sub ?? { unsubscribe: () => {
+    } };
+  }
+  wrap(modelToWrap) {
+    return experimental_wrapLanguageModel({
+      model: modelToWrap,
+      middleware: createAgentMiddleware(this)
+    });
+  }
+  /**
+   * Resolves with an `AgentPlan` based on the information provided in the `options`, including:
+   *
+   * - The `goal` for the agent to achieve
+   * - The observed current `state`
+   * - The `machine` (e.g. a state machine) that specifies what can happen next
+   * - Additional `context`
+   */
+  decide(opts) {
+    return agentDecide(this, opts);
+  }
+};
 export {
   createAgent,
   fromDecision,