@statelyai/agent 2.0.0-next.3 → 2.0.0-next.5

package/src/agent.test.ts

@@ -1,5 +1,10 @@
 import { test, expect, vi } from 'vitest';
-import { createAgent, TypesFromAgent } from './';
+import {
+  AgentDecision,
+  AgentFeedbackInput,
+  createAgent,
+  TypesFromAgent,
+} from './';
 import { createActor, createMachine } from 'xstate';
 import { LanguageModelV1CallOptions } from 'ai';
 import { z } from 'zod';
@@ -66,28 +71,44 @@ test('agent.addMessage() adds to message history', () => {
 test('agent.addFeedback() adds to feedback', () => {
   const agent = createAgent({
     id: 'test',
-    events: {},
+    events: {
+      play: z.object({
+        position: z.number(),
+      }),
+    },
     model: {} as any,
   });
 
+  const decision: AgentDecision<typeof agent> = {
+    goal: 'Win the game',
+    decisionId: null,
+    episodeId: agent.episodeId,
+    goalState: { value: 'won' },
+    id: 'decision-1',
+    nextEvent: { type: 'play', position: 3 },
+    paths: [],
+    policy: 'simple',
+    timestamp: Date.now(),
+  };
+
   const obs = agent.addObservation({
+    decisionId: decision.id,
     prevState: { value: 'playing' },
-    state: { value: 'lost' },
     event: { type: 'play', position: 3 },
-    goal: 'Win the game',
+    state: { value: 'lost' },
   });
 
   const feedback = agent.addFeedback({
-    score: 0,
-    observationId: obs.id,
+    reward: 0,
+    decisionId: decision.id,
   });
 
   expect(feedback.episodeId).toEqual(agent.episodeId);
 
   expect(agent.getFeedback()).toContainEqual(
     expect.objectContaining({
-      score: 0,
-      observationId: obs.id,
+      reward: 0,
+      decisionId: decision.id,
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
@@ -144,7 +165,7 @@ test('agent.addObservation() adds to observations (initial state)', () => {
   );
 });
 
-test('agent.addObservation() adds to observations with machine hash', () => {
+test.skip('agent.addObservation() adds to observations with machine hash', () => {
   const agent = createAgent({
     id: 'test',
     events: {},
@@ -167,7 +188,6 @@ test('agent.addObservation() adds to observations with machine hash', () => {
     prevState: { value: 'playing', context: {} },
     event: { type: 'play', position: 3 },
     state: { value: 'lost', context: {} },
-    machine,
     goal: 'Win the game',
   });
 
@@ -178,14 +198,13 @@ test('agent.addObservation() adds to observations with machine hash', () => {
       prevState: { value: 'playing', context: {} },
       event: { type: 'play', position: 3 },
       state: { value: 'lost', context: {} },
-      machineHash: expect.any(String),
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
   );
 });
 
-test('agent.addFeedback() adds to feedback (with observation)', () => {
+test('agent.addInsight() adds to insights (with observation)', () => {
   const agent = createAgent({
     id: 'test',
     events: {},
@@ -199,24 +218,26 @@ test('agent.addFeedback() adds to feedback (with observation)', () => {
     goal: 'Win the game',
   });
 
-  const feedback = agent.addFeedback({
-    score: 0,
+  const insight = agent.addInsight({
     observationId: observation.id,
+    attributes: {
+      advantage: 50,
+    },
   });
 
-  expect(feedback.episodeId).toEqual(agent.episodeId);
+  expect(insight.episodeId).toEqual(agent.episodeId);
 
-  expect(agent.getFeedback()).toContainEqual(
+  expect(agent.getInsights()).toContainEqual(
     expect.objectContaining({
-      score: 0,
+      attributes: { advantage: 50 },
       observationId: observation.id,
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
   );
-  expect(agent.getFeedback()).toContainEqual(
+  expect(agent.getInsights()).toContainEqual(
     expect.objectContaining({
-      score: 0,
+      attributes: { advantage: 50 },
       observationId: observation.id,
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
@@ -279,14 +300,25 @@ test('You can listen for feedback events', () => {
     model: {} as any,
   });
 
-  agent.on('feedback', fn);
+  agent.onFeedback(fn);
 
   agent.addFeedback({
-    score: -1,
-    observationId: 'obs-1',
+    reward: 1,
+    decisionId: 'dec-1',
+    comment: 'Good move',
+    attributes: { confidence: 'high' },
   });
 
-  expect(fn).toHaveBeenCalled();
+  expect(fn).toHaveBeenCalledWith(
+    expect.objectContaining({
+      reward: 1,
+      decisionId: 'dec-1',
+      comment: 'Good move',
+      attributes: { confidence: 'high' },
+      episodeId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
 });
 
 test('You can listen for decision events', async () => {
@@ -425,20 +457,21 @@ test('agent.getDecisions() returns decisions from context', () => {
     id: 'test',
     events: {},
     model: {} as any,
-    strategy: async (agent) => {
+    policy: async (agent) => {
       return {
         id: Date.now().toString(),
+        decisionId: null,
         episodeId: agent.episodeId,
-        strategy: 'test-strategy',
+        policy: 'test-policy',
         goal: '',
-        goalState: undefined,
+        goalState: null,
         paths: [
           {
-            state: undefined,
+            state: null,
             steps: [],
           },
         ],
-        nextEvent: undefined,
+        nextEvent: null,
         timestamp: Date.now(),
       };
     },
@@ -503,7 +536,6 @@ test('agent.observe() adds observations from actor snapshots', () => {
   expect(agent.getObservations()).toContainEqual(
     expect.objectContaining({
       state: expect.objectContaining({ value: 'idle' }),
-      machineHash: expect.any(String),
     })
   );
 
@@ -512,9 +544,334 @@ test('agent.observe() adds observations from actor snapshots', () => {
       prevState: expect.objectContaining({ value: 'idle' }),
       event: { type: 'START' },
       state: expect.objectContaining({ value: 'running' }),
-      machineHash: expect.any(String),
     })
   );
 
   subscription.unsubscribe();
 });
+
+test('agent.addObservation() accepts custom episodeId', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const customEpisodeId = 'custom-episode-123';
+  const observation = agent.addObservation({
+    state: { value: 'playing' },
+    goal: 'Win the game',
+    episodeId: customEpisodeId,
+  });
+
+  expect(observation.episodeId).toEqual(customEpisodeId);
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      episodeId: customEpisodeId,
+    })
+  );
+});
+
+test('agent.addFeedback() accepts custom episodeId', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const customEpisodeId = 'custom-episode-123';
+  const feedback = agent.addFeedback({
+    reward: 1,
+    decisionId: 'dec-1',
+    episodeId: customEpisodeId,
+  });
+
+  expect(feedback.episodeId).toEqual(customEpisodeId);
+  expect(agent.getFeedback()).toContainEqual(
+    expect.objectContaining({
+      episodeId: customEpisodeId,
+    })
+  );
+});
+
+test('agent.addObservation() accepts decisionId', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const decisionId = 'decision-123';
+  const observation = agent.addObservation({
+    state: { value: 'playing' },
+    goal: 'Win the game',
+    decisionId,
+  });
+
+  expect(observation.decisionId).toEqual(decisionId);
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      decisionId,
+    })
+  );
+});
+
+test('agent.addFeedback() accepts decisionId', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const decisionId = 'decision-123';
+  const feedback = agent.addFeedback({
+    reward: 1,
+    decisionId,
+  });
+
+  expect(feedback.decisionId).toEqual(decisionId);
+  expect(agent.getFeedback()).toContainEqual(
+    expect.objectContaining({
+      decisionId,
+    })
+  );
+});
+
+test('You can listen for observation events', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  agent.onObservation(fn);
+
+  agent.addObservation({
+    state: { value: 'playing' },
+    goal: 'Win the game',
+  });
+
+  expect(fn).toHaveBeenCalledWith(
+    expect.objectContaining({
+      state: { value: 'playing' },
+      episodeId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
+});
+
+test('You can listen for decision events', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {
+      MOVE: z.object({}),
+    },
+    model: {} as any,
+  });
+
+  agent.onDecision(fn);
+
+  const decision = {
+    id: 'decision-1',
+    decisionId: null,
+    episodeId: agent.episodeId,
+    policy: 'test-policy',
+    goal: 'Win the game',
+    goalState: { value: 'won' },
+    paths: [],
+    nextEvent: { type: 'MOVE' },
+    timestamp: Date.now(),
+  } satisfies AgentDecision<typeof agent>;
+
+  agent.addDecision(decision);
+
+  expect(fn).toHaveBeenCalledWith(
+    expect.objectContaining({
+      id: 'decision-1',
+      episodeId: agent.episodeId,
+      policy: 'test-policy',
+      goal: 'Win the game',
+      nextEvent: { type: 'MOVE' },
+    })
+  );
+});
+
+test('Event listeners can be unsubscribed (onObservation)', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const subscription = agent.onObservation(fn);
+
+  agent.addObservation({
+    state: { value: 'playing' },
+    goal: 'Win the game',
+  });
+
+  expect(fn).toHaveBeenCalledTimes(1);
+
+  subscription.unsubscribe();
+
+  agent.addObservation({
+    state: { value: 'playing' },
+    goal: 'Win the game again',
+  });
+
+  expect(fn).toHaveBeenCalledTimes(1); // Still only called once
+});
+
+test('Event listeners can be unsubscribed (onDecision)', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {
+      MOVE: z.object({}),
+    },
+    model: {} as any,
+  });
+
+  const subscription = agent.onDecision(fn);
+
+  const decision = {
+    id: 'decision-1',
+    decisionId: null,
+    episodeId: agent.episodeId,
+    policy: 'test-policy',
+    goal: 'Win the game',
+    goalState: { value: 'won' },
+    paths: [],
+    nextEvent: { type: 'MOVE' },
+    timestamp: Date.now(),
+  } satisfies AgentDecision<typeof agent>;
+
+  agent.addDecision(decision);
+
+  expect(fn).toHaveBeenCalledTimes(1);
+
+  subscription.unsubscribe();
+
+  agent.addDecision({
+    ...decision,
+    id: 'decision-2',
+  });
+
+  expect(fn).toHaveBeenCalledTimes(1); // Still only called once
+});
+
+test('Event listeners can be unsubscribed (onFeedback)', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const subscription = agent.onFeedback(fn);
+
+  agent.addFeedback({
+    reward: 1,
+    decisionId: 'dec-1',
+  });
+
+  expect(fn).toHaveBeenCalledTimes(1);
+
+  subscription.unsubscribe();
+
+  agent.addFeedback({
+    reward: 0,
+    decisionId: 'dec-2',
+  });
+
+  expect(fn).toHaveBeenCalledTimes(1); // Still only called once
+});
+
+test('Feedback events include optional fields', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  agent.onFeedback(fn);
+
+  // Test with minimal feedback
+  agent.addFeedback({
+    reward: 1,
+    decisionId: 'dec-1',
+  });
+
+  expect(fn).toHaveBeenCalledWith(
+    expect.objectContaining({
+      reward: 1,
+      decisionId: 'dec-1',
+      comment: undefined,
+      attributes: {},
+      episodeId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
+
+  // Test with all optional fields
+  agent.addFeedback({
+    reward: 0,
+    decisionId: 'dec-2',
+    comment: 'Could be better',
+    attributes: { reason: 'suboptimal' },
+  } satisfies AgentFeedbackInput);
+
+  expect(fn).toHaveBeenLastCalledWith(
+    expect.objectContaining({
+      reward: 0,
+      decisionId: 'dec-2',
+      comment: 'Could be better',
+      attributes: { reason: 'suboptimal' },
+      episodeId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
+});
+
+test('Feedback events maintain episodeId consistency', () => {
+  const fn = vi.fn();
+  const customEpisodeId = 'custom-episode-123';
+
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+    episodeId: customEpisodeId,
+  });
+
+  agent.onFeedback(fn);
+
+  agent.addFeedback({
+    reward: 1,
+    decisionId: 'dec-1',
+  });
+
+  expect(fn).toHaveBeenCalledWith(
+    expect.objectContaining({
+      episodeId: customEpisodeId,
+    })
+  );
+
+  // Test with explicit different episodeId
+  const differentEpisodeId = 'different-episode-456';
+  agent.addFeedback({
+    reward: 0,
+    decisionId: 'dec-2',
+    episodeId: differentEpisodeId,
+  });
+
+  expect(fn).toHaveBeenLastCalledWith(
+    expect.objectContaining({
+      episodeId: differentEpisodeId,
+    })
+  );
+});