Merge pull request #40 from statelyai/davidkpiano/correlation-id

Correlation ID

.changeset/khaki-emus-jog.md

@@ -0,0 +1,26 @@
+---
+'@statelyai/agent': minor
+---
+
+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,
+});
+```

.changeset/silent-deers-visit.md

@@ -0,0 +1,11 @@
+---
+'@statelyai/agent': minor
+---
+
+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

package.json

@@ -38,7 +38,7 @@
     "ts-node": "^10.9.2",
     "tsup": "^8.2.0",
     "typescript": "^5.5.3",
-    "vitest": "^1.6.0",
+    "vitest": "^2.0.3",
     "wikipedia": "^2.1.2",
     "zod": "^3.23.8"
   },

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) => {
+          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) => {
+          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) => {
+          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) => {
+          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);
+  }
+);

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,

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;

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({
@@ -149,11 +159,17 @@ export async function agentStreamText(
         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>(