@statelyai/agent 1.0.0 → 1.1.2

package/.github/actions/ci-setup/action.yml

@@ -7,18 +7,18 @@ runs:
     - uses: actions/setup-node@v4
       with:
         node-version: 20.x
+        
+    - name: install pnpm
+      run: npm i pnpm@latest -g
+      shell: bash
 
-    - name: Get pnpm store directory
+    - name: Setup npmrc
+      run: echo "//registry.npmjs.org/:_authToken=${{ secrets.NPM_TOKEN }}" > .npmrc
       shell: bash
-      id: pnpm-cache
-      run: |
-        echo "pnpm_cache_dir=$(pnpm store path)" >> $GITHUB_OUTPUT
-    - uses: actions/cache@v4
-      name: Setup pnpm cache
-      with:
-        path: ${{ steps.pnpm-cache.outputs.pnpm_cache_dir }}
-        key: ${{ runner.os }}-pnpm-store-${{ hashFiles('**/pnpm-lock.yaml') }}
-        restore-keys: |
-          ${{ runner.os }}-pnpm-store-
+
+    - name: setup pnpm config
+      run: pnpm config set store-dir $PNPM_CACHE_FOLDER
+      shell: bash
+
     - run: pnpm install
       shell: bash

package/.github/workflows/release.yml

@@ -22,14 +22,26 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v4
-      - uses: ./.github/actions/ci-setup
-
-      - name: Create Release Pull Request or Publish to npm
+      - name: checkout code repository
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: setup node.js
+        uses: actions/setup-node@v3
+        with:
+          node-version: 20
+      - name: install pnpm
+        run: npm i pnpm@latest -g
+      - name: Setup npmrc
+        run: echo "//registry.npmjs.org/:_authToken=${{ secrets.NPM_TOKEN }}" > .npmrc
+      - name: setup pnpm config
+        run: pnpm config set store-dir $PNPM_CACHE_FOLDER
+      - name: install dependencies
+        run: pnpm install
+      - name: create and publish versions
         uses: changesets/action@v1
         with:
           publish: pnpm run release
           version: pnpm run version
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          NPM_TOKEN: ${{ secrets.NPM_TOKEN }}

package/CHANGELOG.md

@@ -1,5 +1,90 @@
 # @statelyai/agent
 
+## 1.1.2
+
+### Patch Changes
+
+- [#43](https://github.com/statelyai/agent/pull/43) [`8e7629c`](https://github.com/statelyai/agent/commit/8e7629c347b29b704ae9576aa1af97e6cd693bc7) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Update dependencies
+
+## 1.1.1
+
+### Patch Changes
+
+- [#41](https://github.com/statelyai/agent/pull/41) [`b2f2b73`](https://github.com/statelyai/agent/commit/b2f2b7307e96d7722968769aae9db2572ede8ce7) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Update dependencies
+
+## 1.1.0
+
+### Minor Changes
+
+- [#39](https://github.com/statelyai/agent/pull/39) [`3cce30f`](https://github.com/statelyai/agent/commit/3cce30fc77d36dbed0abad805248de9f64bf8086) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Added four new methods for easily retrieving agent messages, observations, feedback, and plans:
+
+  - `agent.getMessages()`
+  - `agent.getObservations()`
+  - `agent.getFeedback()`
+  - `agent.getPlans()`
+
+  The `agent.select(…)` method is deprecated in favor of these methods.
+
+- [#40](https://github.com/statelyai/agent/pull/40) [`8b7c374`](https://github.com/statelyai/agent/commit/8b7c37482d5c35b2b3addc2f88e198526f203da7) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Correlation IDs are now provided as part of the result from `agent.generateText(…)` and `agent.streamText(…)`:
+
+  ```ts
+  const result = await agent.generateText({
+    prompt: "Write me a song",
+    correlationId: "my-correlation-id",
+    // ...
+  });
+
+  result.correlationId; // 'my-correlation-id'
+  ```
+
+  These correlation IDs can be passed to feedback:
+
+  ```ts
+  // ...
+
+  agent.addFeedback({
+    reward: -1,
+    correlationId: result.correlationId,
+  });
+  ```
+
+- [#40](https://github.com/statelyai/agent/pull/40) [`8b7c374`](https://github.com/statelyai/agent/commit/8b7c37482d5c35b2b3addc2f88e198526f203da7) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Changes to agent feedback (the `AgentFeedback` interface):
+
+  - `goal` is now optional
+  - `observationId` is now optional
+  - `correlationId` has been added (optional)
+  - `reward` has been added (optional)
+  - `attributes` are now optional
+
+- [#38](https://github.com/statelyai/agent/pull/38) [`21fb17c`](https://github.com/statelyai/agent/commit/21fb17c65fac1cbb4a8b08a04a58480a6930a0a9) Thanks [@davidkpiano](https://github.com/davidkpiano)! - You can now add `context` Zod schema to your agent. For now, this is meant to be passed directly to the state machine, but in the future, the schema can be shared with the LLM agent to better understand the state machine and its context for decision making.
+
+  Breaking: The `context` and `events` types are now in `agent.types` instead of ~~`agent.eventTypes`.
+
+  ```ts
+  const agent = createAgent({
+    // ...
+    context: {
+      score: z.number().describe("The score of the game"),
+      // ...
+    },
+  });
+
+  const machine = setup({
+    types: agent.types,
+  }).createMachine({
+    context: {
+      score: 0,
+    },
+    // ...
+  });
+  ```
+
+### Patch Changes
+
+- [`5f863bb`](https://github.com/statelyai/agent/commit/5f863bb0d89d90f30d0a9aa1f0dd2a35f0eeb45b) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Use nanoid
+
+- [#37](https://github.com/statelyai/agent/pull/37) [`dafa815`](https://github.com/statelyai/agent/commit/dafa8157cc1b5adbfb222c146dbc84ab2eed8894) Thanks [@davidkpiano](https://github.com/davidkpiano)! - Messages are now properly included in `agent.decide(…)`, when specified.
+
 ## 0.1.0
 
 ### Minor Changes

package/dist/index.d.mts

@@ -97,17 +97,24 @@ type AgentDecideOptions = {
     events?: ZodEventMapping;
 } & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt' | 'messages'>;
 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;
 }
 interface AgentFeedbackInput {
-    goal: string;
-    observationId: string;
-    attributes: Record<string, any>;
+    goal?: string;
+    observationId?: string;
+    correlationId?: string;
+    attributes?: Record<string, any>;
     timestamp?: number;
+    reward?: number;
 }
 type AgentMessage = CoreMessage & {
     timestamp: number;
@@ -119,6 +126,8 @@ type AgentMessage = CoreMessage & {
     responseId?: string;
     result?: GenerateTextResult<any>;
     sessionId: string;
+    correlationId: string;
+    parentCorrelationId?: string;
 };
 type AgentMessageInput = CoreMessage & {
     timestamp?: number;
@@ -128,6 +137,8 @@ type AgentMessageInput = CoreMessage & {
      * which message this message is responding to, if any.
      */
     responseId?: string;
+    correlationId?: string;
+    parentCorrelationId?: string;
     result?: GenerateTextResult<any>;
 };
 interface AgentObservation<TActor extends AnyActorRef> {
@@ -223,8 +234,8 @@ type Agent<TContext, TEvents extends EventObject> = ActorRefFrom<AgentLogic<TEve
      * - Additional `context`
      */
     decide: (options: AgentDecideOptions) => Promise<AgentPlan<TEvents> | undefined>;
-    generateText: (options: AgentGenerateTextOptions) => Promise<GenerateTextResult<Record<string, any>>>;
-    streamText: (options: AgentStreamTextOptions) => Promise<StreamTextResult<Record<string, CoreTool<any, any>>>>;
+    generateText: (options: AgentGenerateTextOptions) => Promise<AgentGenerateTextResult>;
+    streamText: (options: AgentStreamTextOptions) => Promise<AgentStreamTextResult>;
     addObservation: (observationInput: AgentObservationInput) => AgentObservation<any>;
     addMessage: (messageInput: AgentMessageInput) => AgentMessage;
     addFeedback: (feedbackInput: AgentFeedbackInput) => AgentFeedback;
@@ -306,16 +317,24 @@ type Agent<TContext, TEvents extends EventObject> = ActorRefFrom<AgentLogic<TEve
     interact<TActor extends AnyActorRef>(actorRef: TActor, getInput: (observation: AgentObservation<TActor>) => AgentDecisionInput | undefined): Subscription;
 };
 type AnyAgent = Agent<any, any>;
-type FromAgent<T> = T | ((self: AnyAgent) => T | Promise<T>);
-interface CommonTextOptions {
+type FromAgent<T> = T | ((agent: AnyAgent) => T | Promise<T>);
+type CommonTextOptions = {
     prompt: FromAgent<string>;
     model?: LanguageModel;
     context?: Record<string, any>;
     messages?: FromAgent<CoreMessage[]>;
     template?: PromptTemplate<any>;
-}
+    correlationId?: string;
+    parentCorrelationId?: string;
+};
+type TextResultMeta = {
+    correlationId: string;
+    parentCorrelationId?: string;
+};
 type AgentGenerateTextOptions = Omit<GenerateTextOptions, 'model' | 'prompt' | 'messages'> & CommonTextOptions;
+type AgentGenerateTextResult = GenerateTextResult<any> & TextResultMeta;
 type AgentStreamTextOptions = Omit<StreamTextOptions, 'model' | 'prompt' | 'messages'> & CommonTextOptions;
+type AgentStreamTextResult = StreamTextResult<any> & TextResultMeta;
 interface ObservedState {
     /**
      * The current state value of the state machine, e.g.
@@ -406,4 +425,4 @@ declare function fromText<T extends AnyAgent>(agent: T, defaultOptions?: AgentGe
 
 declare function fromDecision(agent: AnyAgent, defaultInput?: AgentDecisionInput): AgentDecisionLogic<any>;
 
-export { type AIAdapter, type Agent, type AgentDecideOptions, type AgentDecisionInput, type AgentDecisionLogic, type AgentEmitted, type AgentFeedback, type AgentFeedbackInput, type AgentGenerateTextOptions, type AgentLogic, type AgentLongTermMemory, type AgentMemory, type AgentMemoryContext, type AgentMessage, type AgentMessageInput, type AgentObservation, type AgentObservationInput, type AgentPlan, type AgentPlanInput, type AgentPlanner, type AgentStreamTextOptions, type AnyAgent, type AppendOnlyStorage, type CommonTextOptions, type Compute, type ContextFromZodContextMapping, type EventsFromZodEventMapping, type FromAgent, type GenerateTextOptions, type ObservedState, type ObservedStateFrom, type PromptTemplate, type StreamTextOptions, type TransitionData, createAgent, fromDecision, fromText, fromTextStream };
+export { type AIAdapter, type Agent, type AgentDecideOptions, type AgentDecisionInput, type AgentDecisionLogic, type AgentEmitted, type AgentFeedback, type AgentFeedbackInput, type AgentGenerateTextOptions, type AgentGenerateTextResult, type AgentLogic, type AgentLongTermMemory, type AgentMemory, type AgentMemoryContext, type AgentMessage, type AgentMessageInput, type AgentObservation, type AgentObservationInput, type AgentPlan, type AgentPlanInput, type AgentPlanner, type AgentStreamTextOptions, type AgentStreamTextResult, type AnyAgent, type AppendOnlyStorage, type CommonTextOptions, type Compute, type ContextFromZodContextMapping, type EventsFromZodEventMapping, type FromAgent, type GenerateTextOptions, type ObservedState, type ObservedStateFrom, type PromptTemplate, type StreamTextOptions, type TextResultMeta, type TransitionData, createAgent, fromDecision, fromText, fromTextStream };

package/dist/index.d.ts

@@ -97,17 +97,24 @@ type AgentDecideOptions = {
     events?: ZodEventMapping;
 } & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt' | 'messages'>;
 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;
 }
 interface AgentFeedbackInput {
-    goal: string;
-    observationId: string;
-    attributes: Record<string, any>;
+    goal?: string;
+    observationId?: string;
+    correlationId?: string;
+    attributes?: Record<string, any>;
     timestamp?: number;
+    reward?: number;
 }
 type AgentMessage = CoreMessage & {
     timestamp: number;
@@ -119,6 +126,8 @@ type AgentMessage = CoreMessage & {
     responseId?: string;
     result?: GenerateTextResult<any>;
     sessionId: string;
+    correlationId: string;
+    parentCorrelationId?: string;
 };
 type AgentMessageInput = CoreMessage & {
     timestamp?: number;
@@ -128,6 +137,8 @@ type AgentMessageInput = CoreMessage & {
      * which message this message is responding to, if any.
      */
     responseId?: string;
+    correlationId?: string;
+    parentCorrelationId?: string;
     result?: GenerateTextResult<any>;
 };
 interface AgentObservation<TActor extends AnyActorRef> {
@@ -223,8 +234,8 @@ type Agent<TContext, TEvents extends EventObject> = ActorRefFrom<AgentLogic<TEve
      * - Additional `context`
      */
     decide: (options: AgentDecideOptions) => Promise<AgentPlan<TEvents> | undefined>;
-    generateText: (options: AgentGenerateTextOptions) => Promise<GenerateTextResult<Record<string, any>>>;
-    streamText: (options: AgentStreamTextOptions) => Promise<StreamTextResult<Record<string, CoreTool<any, any>>>>;
+    generateText: (options: AgentGenerateTextOptions) => Promise<AgentGenerateTextResult>;
+    streamText: (options: AgentStreamTextOptions) => Promise<AgentStreamTextResult>;
     addObservation: (observationInput: AgentObservationInput) => AgentObservation<any>;
     addMessage: (messageInput: AgentMessageInput) => AgentMessage;
     addFeedback: (feedbackInput: AgentFeedbackInput) => AgentFeedback;
@@ -306,16 +317,24 @@ type Agent<TContext, TEvents extends EventObject> = ActorRefFrom<AgentLogic<TEve
     interact<TActor extends AnyActorRef>(actorRef: TActor, getInput: (observation: AgentObservation<TActor>) => AgentDecisionInput | undefined): Subscription;
 };
 type AnyAgent = Agent<any, any>;
-type FromAgent<T> = T | ((self: AnyAgent) => T | Promise<T>);
-interface CommonTextOptions {
+type FromAgent<T> = T | ((agent: AnyAgent) => T | Promise<T>);
+type CommonTextOptions = {
     prompt: FromAgent<string>;
     model?: LanguageModel;
     context?: Record<string, any>;
     messages?: FromAgent<CoreMessage[]>;
     template?: PromptTemplate<any>;
-}
+    correlationId?: string;
+    parentCorrelationId?: string;
+};
+type TextResultMeta = {
+    correlationId: string;
+    parentCorrelationId?: string;
+};
 type AgentGenerateTextOptions = Omit<GenerateTextOptions, 'model' | 'prompt' | 'messages'> & CommonTextOptions;
+type AgentGenerateTextResult = GenerateTextResult<any> & TextResultMeta;
 type AgentStreamTextOptions = Omit<StreamTextOptions, 'model' | 'prompt' | 'messages'> & CommonTextOptions;
+type AgentStreamTextResult = StreamTextResult<any> & TextResultMeta;
 interface ObservedState {
     /**
      * The current state value of the state machine, e.g.
@@ -406,4 +425,4 @@ declare function fromText<T extends AnyAgent>(agent: T, defaultOptions?: AgentGe
 
 declare function fromDecision(agent: AnyAgent, defaultInput?: AgentDecisionInput): AgentDecisionLogic<any>;
 
-export { type AIAdapter, type Agent, type AgentDecideOptions, type AgentDecisionInput, type AgentDecisionLogic, type AgentEmitted, type AgentFeedback, type AgentFeedbackInput, type AgentGenerateTextOptions, type AgentLogic, type AgentLongTermMemory, type AgentMemory, type AgentMemoryContext, type AgentMessage, type AgentMessageInput, type AgentObservation, type AgentObservationInput, type AgentPlan, type AgentPlanInput, type AgentPlanner, type AgentStreamTextOptions, type AnyAgent, type AppendOnlyStorage, type CommonTextOptions, type Compute, type ContextFromZodContextMapping, type EventsFromZodEventMapping, type FromAgent, type GenerateTextOptions, type ObservedState, type ObservedStateFrom, type PromptTemplate, type StreamTextOptions, type TransitionData, createAgent, fromDecision, fromText, fromTextStream };
+export { type AIAdapter, type Agent, type AgentDecideOptions, type AgentDecisionInput, type AgentDecisionLogic, type AgentEmitted, type AgentFeedback, type AgentFeedbackInput, type AgentGenerateTextOptions, type AgentGenerateTextResult, type AgentLogic, type AgentLongTermMemory, type AgentMemory, type AgentMemoryContext, type AgentMessage, type AgentMessageInput, type AgentObservation, type AgentObservationInput, type AgentPlan, type AgentPlanInput, type AgentPlanner, type AgentStreamTextOptions, type AgentStreamTextResult, type AnyAgent, type AppendOnlyStorage, type CommonTextOptions, type Compute, type ContextFromZodContextMapping, type EventsFromZodEventMapping, type FromAgent, type GenerateTextOptions, type ObservedState, type ObservedStateFrom, type PromptTemplate, type StreamTextOptions, type TextResultMeta, type TransitionData, createAgent, fromDecision, fromText, fromTextStream };

package/dist/index.js

@@ -123,7 +123,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();
@@ -137,7 +138,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,
@@ -150,14 +153,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();
@@ -171,7 +181,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,
@@ -189,16 +201,24 @@ 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 (0, import_xstate.fromObservable)(({ input }) => {
@@ -309,7 +329,6 @@ async function simplePlanner(agent, input) {
   });
   const singleResult = result.toolResults[0];
   if (!singleResult) {
-    console.log(toolMap);
     console.warn("No tool call results returned");
     return void 0;
   }
@@ -442,12 +461,12 @@ var agentLogic = (0, import_xstate3.fromTransition)(
     }
     return state;
   },
-  {
+  () => ({
     feedback: [],
     messages: [],
     observations: [],
     plans: []
-  }
+  })
 );
 function createAgent({
   name,
@@ -462,7 +481,6 @@ function createAgent({
   adapter = vercelAdapter,
   ...generateTextOptions
 }) {
-  const messageHistoryListeners = [];
   const agent = (0, import_xstate3.createActor)(logic);
   agent.events = events;
   agent.model = model;
@@ -475,7 +493,7 @@ function createAgent({
   };
   agent.memory = getMemory ? getMemory(agent) : void 0;
   agent.onMessage = (callback) => {
-    messageHistoryListeners.push((0, import_xstate3.toObserver)(callback));
+    agent.on("message", (ev) => callback(ev.message));
   };
   agent.decide = (opts) => {
     return agentDecide(agent, opts);
@@ -485,7 +503,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",
@@ -499,6 +518,8 @@ function createAgent({
   agent.addFeedback = (feedbackInput) => {
     const feedback = {
       ...feedbackInput,
+      attributes: { ...feedbackInput.attributes },
+      reward: feedbackInput.reward ?? 0,
       timestamp: feedbackInput.timestamp ?? Date.now(),
       sessionId: agent.sessionId
     };

package/dist/index.mjs

@@ -1,8 +1,7 @@
 // src/agent.ts
 import {
   createActor,
-  fromTransition,
-  toObserver as toObserver2
+  fromTransition
 } from "xstate";
 
 // src/planners/simplePlanner.ts
@@ -92,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();
@@ -106,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,
@@ -119,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();
@@ -140,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,
@@ -158,16 +169,24 @@ 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 }) => {
@@ -278,7 +297,6 @@ async function simplePlanner(agent, input) {
   });
   const singleResult = result.toolResults[0];
   if (!singleResult) {
-    console.log(toolMap);
     console.warn("No tool call results returned");
     return void 0;
   }
@@ -411,12 +429,12 @@ var agentLogic = fromTransition(
     }
     return state;
   },
-  {
+  () => ({
     feedback: [],
     messages: [],
     observations: [],
     plans: []
-  }
+  })
 );
 function createAgent({
   name,
@@ -431,7 +449,6 @@ function createAgent({
   adapter = vercelAdapter,
   ...generateTextOptions
 }) {
-  const messageHistoryListeners = [];
   const agent = createActor(logic);
   agent.events = events;
   agent.model = model;
@@ -444,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);
@@ -454,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",
@@ -468,6 +486,8 @@ function createAgent({
   agent.addFeedback = (feedbackInput) => {
     const feedback = {
       ...feedbackInput,
+      attributes: { ...feedbackInput.attributes },
+      reward: feedbackInput.reward ?? 0,
       timestamp: feedbackInput.timestamp ?? Date.now(),
       sessionId: agent.sessionId
     };

package/package.json

@@ -1,21 +1,10 @@
 {
   "name": "@statelyai/agent",
-  "version": "1.0.0",
+  "version": "1.1.2",
   "description": "Stateful agents that make decisions based on finite-state machine models",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
   "types": "dist/index.d.ts",
-  "scripts": {
-    "build": "tsup src/index.ts --format cjs,esm --dts",
-    "lint": "tsc --noEmit",
-    "test": "vitest",
-    "test:ci": "vitest --run",
-    "example": "ts-node examples/helpers/runner.ts",
-    "prepublishOnly": "tsup src/index.ts --format cjs,esm --dts",
-    "changeset": "changeset",
-    "release": "changeset publish",
-    "version": "changeset version"
-  },
   "keywords": [
     "ai",
     "state machine",
@@ -31,14 +20,14 @@
     "@langchain/community": "^0.0.53",
     "@langchain/core": "^0.1.63",
     "@langchain/openai": "^0.0.28",
-    "@types/node": "^20.14.10",
+    "@types/node": "^20.14.13",
     "@types/object-hash": "^3.0.6",
     "dotenv": "^16.4.5",
     "json-schema-to-ts": "^3.1.0",
     "ts-node": "^10.9.2",
-    "tsup": "^8.1.0",
-    "typescript": "^5.5.3",
-    "vitest": "^1.6.0",
+    "tsup": "^8.2.3",
+    "typescript": "^5.5.4",
+    "vitest": "^2.0.4",
     "wikipedia": "^2.1.2",
     "zod": "^3.23.8"
   },
@@ -46,11 +35,20 @@
     "access": "public"
   },
   "dependencies": {
-    "@ai-sdk/openai": "^0.0.31",
+    "@ai-sdk/openai": "^0.0.40",
     "@xstate/graph": "^2.0.0",
-    "ai": "^3.2.22",
+    "ai": "^3.2.40",
     "object-hash": "^3.0.0",
-    "xstate": "^5.15.0"
+    "xstate": "^5.16.0"
   },
-  "packageManager": "pnpm@8.11.0"
-}
+  "scripts": {
+    "build": "tsup src/index.ts --format cjs,esm --dts",
+    "lint": "tsc --noEmit",
+    "test": "vitest",
+    "test:ci": "vitest --run",
+    "example": "ts-node examples/helpers/runner.ts",
+    "changeset": "changeset",
+    "release": "changeset publish",
+    "version": "changeset version"
+  }
+}

package/readme.md

@@ -1,8 +1,5 @@
 # Stately Agent
 
-> [!WARNING]  
-> Alpha software! Not ready for production yet.
-
 Stately Agent is a flexible framework for building AI agents using state machines. Stately agents go beyond normal LLM-based AI agents by:
 
 - Using state machines to guide the agent's behavior, powered by [XState](https://stately.ai/docs/xstate)
@@ -10,4 +7,4 @@ Stately Agent is a flexible framework for building AI agents using state machine
 - Enabling custom **planning** abilities for agents to achieve specific goals based on state machine logic, observations, and feedback
 - Wrapping the [Vercel AI SDK](https://sdk.vercel.ai/) to easily support multiple model providers, such as OpenAI, Anthropic, Google, Mistral, Groq, Perplexity, and more
 
-_Documentation coming soon_
+**Read the documentation: [stately.ai/docs/agents](https://stately.ai/docs/agents)**

package/src/agent.test.ts

@@ -1,5 +1,10 @@
 import { test, expect, vi } from 'vitest';
-import { createAgent, type AIAdapter } from './';
+import {
+  AgentGenerateTextResult,
+  AgentMessage,
+  createAgent,
+  type AIAdapter,
+} from './';
 import { createActor, createMachine } from 'xstate';
 import { GenerateTextResult } from 'ai';
 import { z } from 'zod';
@@ -297,7 +302,7 @@ test('You can listen for plan events', async () => {
               },
             },
           ],
-        } as any as GenerateTextResult<any>;
+        } as any as AgentGenerateTextResult;
       },
       streamText: {} as any,
     },
@@ -357,3 +362,145 @@ test('agent.types provides context and event types', () => {
   // @ts-expect-error
   agent.types.context satisfies { score: string };
 });
+
+test.each(['generateText', 'streamText'] as const)(
+  'can provide a correlation ID (%s)',
+  async (method) => {
+    const agent = createAgent({
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts: any) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    const promise = new Promise<AgentMessage>((res) => {
+      agent.onMessage((msg) => {
+        if (msg.role === 'assistant') {
+          res(msg);
+        }
+      });
+    });
+
+    await agent[method]({
+      prompt: 'hi',
+      correlationId: 'c-1',
+    });
+
+    const msg = await promise;
+
+    expect(msg.correlationId).toBe('c-1');
+    expect(msg.parentCorrelationId).toBe(undefined);
+  }
+);
+
+test.each(['generateText', 'streamText'] as const)(
+  'correlation IDs are automatically generated if not provided (%s)',
+  async (method) => {
+    const agent = createAgent({
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts: any) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    await agent[method]({
+      prompt: 'hi',
+    });
+
+    const messages = agent.getMessages();
+
+    expect(messages[0]?.correlationId).toEqual(expect.stringMatching(/.+/));
+    expect(messages[0]?.role).toBe('user');
+    expect(messages[1]?.correlationId).toEqual(expect.stringMatching(/.+/));
+    expect(messages[1]?.role).toBe('assistant');
+
+    expect(messages[0]!.correlationId).toEqual(messages[1]!.correlationId);
+  }
+);
+
+test.each(['generateText', 'streamText'] as const)(
+  'can provide a parent correlation ID (%s)',
+  async (method) => {
+    const agent = createAgent({
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts: any) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    await agent[method]({
+      prompt: 'hi',
+      correlationId: 'c-1',
+      parentCorrelationId: 'c-0',
+    });
+
+    const msg = agent.getMessages().find((msg) => msg.role === 'assistant')!;
+
+    expect(msg.correlationId).toBe('c-1');
+    expect(msg.parentCorrelationId).toBe('c-0');
+  }
+);
+
+test.each(['generateText', 'streamText'] as const)(
+  'can add feedback to a correlation (%s)',
+  async (method) => {
+    const agent = createAgent({
+      name: 'test',
+      model: {} as any,
+      events: {},
+      adapter: {
+        [method]: async (opts: any) => {
+          const res = {
+            text: 'response',
+          };
+
+          opts.onFinish?.(res);
+
+          return res as AgentGenerateTextResult;
+        },
+      } as any as AIAdapter,
+    });
+
+    const res = await agent[method]({
+      prompt: 'test',
+    });
+
+    agent.addFeedback({
+      correlationId: res.correlationId,
+      reward: -1,
+    });
+
+    const message = agent.getMessages()[0]!;
+    const feedback = agent.getFeedback()[0]!;
+
+    expect(message.correlationId).toBeDefined();
+    expect(feedback.correlationId).toEqual(message.correlationId);
+  }
+);

package/src/agent.ts

@@ -22,6 +22,7 @@ import {
   AgentMemoryContext,
   AgentObservation,
   ContextFromZodContextMapping,
+  AgentFeedback,
 } from './types';
 import { simplePlanner } from './planners/simplePlanner';
 import { agentGenerateText, agentStreamText } from './text';
@@ -73,12 +74,13 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
     }
     return state;
   },
-  {
-    feedback: [],
-    messages: [],
-    observations: [],
-    plans: [],
-  } as AgentMemoryContext
+  () =>
+    ({
+      feedback: [],
+      messages: [],
+      observations: [],
+      plans: [],
+    } as AgentMemoryContext)
 );
 
 export function createAgent<
@@ -140,8 +142,6 @@ export function createAgent<
   logic?: AgentLogic<TEvents>;
   adapter?: AIAdapter;
 } & GenerateTextOptions): Agent<TContext, TEvents> {
-  const messageHistoryListeners: Observer<AgentMessage>[] = [];
-
   const agent = createActor(logic) as unknown as Agent<TContext, TEvents>;
   agent.events = events;
   agent.model = model;
@@ -155,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) => {
@@ -168,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,
@@ -185,9 +186,11 @@ 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,

package/src/planners/simplePlanner.ts

@@ -1,4 +1,4 @@
-import { CoreTool, tool } from 'ai';
+import { type CoreTool, tool } from 'ai';
 import {
   AgentPlan,
   AgentPlanInput,
@@ -118,7 +118,6 @@ export async function simplePlanner<T extends AnyAgent>(
   const singleResult = result.toolResults[0];
 
   if (!singleResult) {
-    console.log(toolMap);
     // TODO: retries?
     console.warn('No tool call results returned');
     return undefined;

package/src/text.ts

@@ -1,12 +1,9 @@
-import type {
-  CoreMessage,
-  CoreTool,
-  GenerateTextResult,
-  StreamTextResult,
-} from 'ai';
+import type { CoreMessage, CoreTool, GenerateTextResult } from 'ai';
 import {
   AgentGenerateTextOptions,
+  AgentGenerateTextResult,
   AgentStreamTextOptions,
+  AgentStreamTextResult,
   AnyAgent,
 } from './types';
 import { defaultTextTemplate } from './templates/defaultText';
@@ -51,11 +48,13 @@ export async function getMessages(
 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();
@@ -76,6 +75,8 @@ export async function agentGenerateText<T extends AnyAgent>(
     role: 'user',
     content: promptWithContext,
     timestamp: Date.now(),
+    correlationId: resolvedOptions.correlationId,
+    parentCorrelationId: resolvedOptions.parentCorrelationId,
   });
 
   const result = await agent.adapter.generateText({
@@ -91,18 +92,25 @@ export async function agentGenerateText<T extends AnyAgent>(
     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: AnyAgent,
   options: AgentStreamTextOptions
-): Promise<StreamTextResult<any>> {
+): Promise<AgentStreamTextResult> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
+    correlationId: options.correlationId ?? randomId(),
   };
   const template = resolvedOptions.template ?? defaultTextTemplate;
 
@@ -124,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({
@@ -143,16 +153,23 @@ 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 AnyAgent>(

package/src/types.ts

@@ -128,18 +128,25 @@ 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 AgentMessage = CoreMessage & {
@@ -152,6 +159,8 @@ export type AgentMessage = CoreMessage & {
   responseId?: string;
   result?: GenerateTextResult<any>;
   sessionId: string;
+  correlationId: string;
+  parentCorrelationId?: string;
 };
 
 export type AgentMessageInput = CoreMessage & {
@@ -162,6 +171,8 @@ export type AgentMessageInput = CoreMessage & {
    * which message this message is responding to, if any.
    */
   responseId?: string;
+  correlationId?: string;
+  parentCorrelationId?: string;
   result?: GenerateTextResult<any>;
 };
 
@@ -293,12 +304,12 @@ export type Agent<TContext, 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: (
     observationInput: AgentObservationInput
@@ -395,15 +406,22 @@ export type Agent<TContext, TEvents extends EventObject> = ActorRefFrom<
 
 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[]>;
   template?: PromptTemplate<any>;
-}
+  correlationId?: string;
+  parentCorrelationId?: string;
+};
+
+export type TextResultMeta = {
+  correlationId: string;
+  parentCorrelationId?: string;
+};
 
 export type AgentGenerateTextOptions = Omit<
   GenerateTextOptions,
@@ -411,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.

package/.changeset/heavy-moons-bake.md

@@ -1,12 +0,0 @@
----
-'@statelyai/agent': minor
----
-
-Added four new methods for easily retrieving agent messages, observations, feedback, and plans:
-
-- `agent.getMessages()`
-- `agent.getObservations()`
-- `agent.getFeedback()`
-- `agent.getPlans()`
-
-The `agent.select(…)` method is deprecated in favor of these methods.

package/.changeset/shaggy-buttons-itch.md

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

package/.changeset/silly-berries-drop.md

@@ -1,5 +0,0 @@
----
-'@statelyai/agent': patch
----
-
-Messages are now properly included in `agent.decide(…)`, when specified.

package/.changeset/wild-bobcats-care.md

@@ -1,26 +0,0 @@
----
-'@statelyai/agent': minor
----
-
-You can now add `context` Zod schema to your agent. For now, this is meant to be passed directly to the state machine, but in the future, the schema can be shared with the LLM agent to better understand the state machine and its context for decision making.
-
-Breaking: The `context` and `events` types are now in `agent.types` instead of ~~`agent.eventTypes`.
-
-```ts
-const agent = createAgent({
-  // ...
-  context: {
-    score: z.number().describe('The score of the game'),
-    // ...
-  },
-});
-
-const machine = setup({
-  types: agent.types,
-}).createMachine({
-  context: {
-    score: 0,
-  },
-  // ...
-});
-```