@statelyai/agent 1.1.6 → 2.0.0-next.1

package/src/agent.test.ts

@@ -1,24 +1,18 @@
 import { test, expect, vi } from 'vitest';
-import {
-  AgentGenerateTextResult,
-  AgentMessage,
-  createAgent,
-  type AIAdapter,
-} from './';
+import { createAgent } from './';
 import { createActor, createMachine } from 'xstate';
-import { GenerateTextResult } from 'ai';
+import { LanguageModelV1CallOptions } from 'ai';
 import { z } from 'zod';
+import { dummyResponseValues, MockLanguageModelV1 } from './mockModel';
 
 test('an agent has the expected interface', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
-    model: {} as any,
+    model: new MockLanguageModelV1(),
   });
 
   expect(agent.decide).toBeDefined();
-  expect(agent.generateText).toBeDefined();
-  expect(agent.streamText).toBeDefined();
 
   expect(agent.addMessage).toBeDefined();
   expect(agent.addObservation).toBeDefined();
@@ -34,46 +28,36 @@ test('an agent has the expected interface', () => {
 });
 
 test('agent.addMessage() adds to message history', () => {
+  const model = new MockLanguageModelV1();
+
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
-    model: {} as any,
+    model,
   });
 
   agent.addMessage({
-    content: 'msg 1',
     role: 'user',
+    content: [{ type: 'text', text: 'msg 1' }],
   });
 
   const messageHistory = agent.addMessage({
-    content: 'response 1',
     role: 'assistant',
+    content: [{ type: 'text', text: 'response 1' }],
   });
 
-  expect(messageHistory.sessionId).toEqual(agent.sessionId);
+  expect(messageHistory.episodeId).toEqual(agent.episodeId);
 
-  expect(agent.select((c) => c.messages)).toContainEqual(
-    expect.objectContaining({
-      content: 'msg 1',
-    })
-  );
   expect(agent.getMessages()).toContainEqual(
     expect.objectContaining({
-      content: 'msg 1',
+      content: [expect.objectContaining({ text: 'msg 1' })],
     })
   );
 
-  expect(agent.select((c) => c.messages)).toContainEqual(
-    expect.objectContaining({
-      content: 'response 1',
-      sessionId: expect.any(String),
-      timestamp: expect.any(Number),
-    })
-  );
   expect(agent.getMessages()).toContainEqual(
     expect.objectContaining({
-      content: 'response 1',
-      sessionId: expect.any(String),
+      content: [expect.objectContaining({ text: 'response 1' })],
+      episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
   );
@@ -81,7 +65,7 @@ test('agent.addMessage() adds to message history', () => {
 
 test('agent.addFeedback() adds to feedback', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -94,16 +78,16 @@ test('agent.addFeedback() adds to feedback', () => {
     observationId: 'obs-1',
   });
 
-  expect(feedback.sessionId).toEqual(agent.sessionId);
+  expect(feedback.episodeId).toEqual(agent.episodeId);
 
-  expect(agent.select((c) => c.feedback)).toContainEqual(
+  expect(agent.getFeedback()).toContainEqual(
     expect.objectContaining({
       attributes: {
         score: -1,
       },
       goal: 'Win the game',
       observationId: 'obs-1',
-      sessionId: expect.any(String),
+      episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
   );
@@ -114,7 +98,7 @@ test('agent.addFeedback() adds to feedback', () => {
       },
       goal: 'Win the game',
       observationId: 'obs-1',
-      sessionId: expect.any(String),
+      episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
   );
@@ -122,7 +106,7 @@ test('agent.addFeedback() adds to feedback', () => {
 
 test('agent.addObservation() adds to observations', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -133,14 +117,36 @@ test('agent.addObservation() adds to observations', () => {
     state: { value: 'lost', context: {} },
   });
 
-  expect(observation.sessionId).toEqual(agent.sessionId);
+  expect(observation.episodeId).toEqual(agent.episodeId);
 
-  expect(agent.select((c) => c.observations)).toContainEqual(
+  expect(agent.getObservations()).toContainEqual(
     expect.objectContaining({
       prevState: { value: 'playing', context: {} },
       event: { type: 'play', position: 3 },
       state: { value: 'lost', context: {} },
-      sessionId: expect.any(String),
+      episodeId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
+});
+
+test('agent.addObservation() adds to observations (initial state)', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const observation = agent.addObservation({
+    state: { value: 'lost' },
+  });
+
+  expect(observation.episodeId).toEqual(agent.episodeId);
+
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      state: { value: 'lost', context: undefined },
+      episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
   );
@@ -148,7 +154,7 @@ test('agent.addObservation() adds to observations', () => {
 
 test('agent.addObservation() adds to observations with machine hash', () => {
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -172,15 +178,62 @@ test('agent.addObservation() adds to observations with machine hash', () => {
     machine,
   });
 
-  expect(observation.sessionId).toEqual(agent.sessionId);
+  expect(observation.episodeId).toEqual(agent.episodeId);
 
-  expect(agent.select((c) => c.observations)).toContainEqual(
+  expect(agent.getObservations()).toContainEqual(
     expect.objectContaining({
       prevState: { value: 'playing', context: {} },
       event: { type: 'play', position: 3 },
       state: { value: 'lost', context: {} },
       machineHash: expect.any(String),
-      sessionId: expect.any(String),
+      episodeId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
+});
+
+test('agent.addFeedback() adds to feedback (with observation)', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
+
+  const observation = agent.addObservation({
+    state: {
+      value: 'playing',
+    },
+  });
+
+  const feedback = agent.addFeedback({
+    attributes: {
+      score: -1,
+    },
+    goal: 'Win the game',
+    observationId: observation.id,
+  });
+
+  expect(feedback.episodeId).toEqual(agent.episodeId);
+
+  expect(agent.getFeedback()).toContainEqual(
+    expect.objectContaining({
+      attributes: {
+        score: -1,
+      },
+      goal: 'Win the game',
+      observationId: observation.id,
+      episodeId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
+  expect(agent.getFeedback()).toContainEqual(
+    expect.objectContaining({
+      attributes: {
+        score: -1,
+      },
+      goal: 'Win the game',
+      observationId: observation.id,
+      episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
   );
@@ -198,7 +251,7 @@ test('agent.interact() observes machine actors (no 2nd arg)', () => {
   });
 
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
     model: {} as any,
   });
@@ -209,7 +262,7 @@ test('agent.interact() observes machine actors (no 2nd arg)', () => {
 
   actor.start();
 
-  expect(agent.select((c) => c.observations)).toContainEqual(
+  expect(agent.getObservations()).toContainEqual(
     expect.objectContaining({
       prevState: undefined,
       state: expect.objectContaining({ value: 'a' }),
@@ -224,7 +277,7 @@ test('agent.interact() observes machine actors (no 2nd arg)', () => {
 
   actor.send({ type: 'NEXT' });
 
-  expect(agent.select((c) => c.observations)).toContainEqual(
+  expect(agent.getObservations()).toContainEqual(
     expect.objectContaining({
       prevState: expect.objectContaining({ value: 'a' }),
       event: { type: 'NEXT' },
@@ -233,35 +286,11 @@ test('agent.interact() observes machine actors (no 2nd arg)', () => {
   );
 });
 
-test('Agents can use a custom adapter', async () => {
-  const adapter = {
-    generateText: async () => {
-      return {
-        text: 'Response',
-      } as any;
-    },
-  } as unknown as AIAdapter;
-
-  const agent = createAgent({
-    name: 'test',
-    events: {},
-    adapter,
-    model: {} as any,
-  });
-
-  const res = await agent.generateText({
-    prompt: 'Question?',
-  });
-
-  expect(res.text).toEqual('Response');
-});
-
 test('You can listen for feedback events', () => {
   const fn = vi.fn();
   const agent = createAgent({
-    name: 'test',
+    id: 'test',
     events: {},
-    adapter: {} as any,
     model: {} as any,
   });
 
@@ -280,32 +309,34 @@ test('You can listen for feedback events', () => {
 
 test('You can listen for plan events', async () => {
   const fn = vi.fn();
+  const model = new MockLanguageModelV1({
+    doGenerate: async (params: LanguageModelV1CallOptions) => {
+      const keys =
+        params.mode.type === 'regular'
+          ? params.mode.tools?.map((tool) => tool.name)
+          : [];
+
+      return {
+        ...dummyResponseValues,
+        finishReason: 'tool-calls',
+        toolCalls: [
+          {
+            toolCallType: 'function',
+            toolCallId: 'call-1',
+            toolName: keys![0],
+            args: `{ "type": "${keys?.[0]}" }`,
+          },
+        ],
+      } as any;
+    },
+  });
+
   const agent = createAgent({
-    name: 'test',
-    model: {} as any,
+    id: 'test',
+    model,
     events: {
       WIN: z.object({}),
     },
-    adapter: {
-      generateText: async (arg) => {
-        const keys = Object.keys(arg.tools!);
-
-        if (keys.length !== 1) {
-          throw new Error('Expected only 1 choice');
-        }
-
-        return {
-          toolResults: [
-            {
-              result: {
-                type: keys[0],
-              },
-            },
-          ],
-        } as any as AgentGenerateTextResult;
-      },
-      streamText: {} as any,
-    },
   });
 
   agent.on('plan', fn);
@@ -363,144 +394,142 @@ test('agent.types provides context and event types', () => {
   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,
-    });
+test('It allows unrecognized events', () => {
+  const agent = createAgent({
+    model: {} as any,
+    events: {},
+    context: {},
+  });
 
-    const promise = new Promise<AgentMessage>((res) => {
-      agent.onMessage((msg) => {
-        if (msg.role === 'assistant') {
-          res(msg);
-        }
-      });
+  expect(() => {
+    agent.send({
+      // @ts-expect-error
+      type: 'unrecognized',
     });
+  }).not.toThrow();
+});
 
-    await agent[method]({
-      prompt: 'hi',
-      correlationId: 'c-1',
-    });
+test('You can listen for message events', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
 
-    const msg = await promise;
+  agent.onMessage(fn);
 
-    expect(msg.correlationId).toBe('c-1');
-    expect(msg.parentCorrelationId).toBe(undefined);
-  }
-);
+  const message = {
+    role: 'user' as const,
+    content: [{ type: 'text' as const, text: 'test message' }],
+  };
 
-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',
-          };
+  agent.addMessage(message);
 
-          opts.onFinish?.(res);
+  expect(fn).toHaveBeenCalledWith(
+    expect.objectContaining({
+      role: 'user',
+      content: [{ type: 'text', text: 'test message' }],
+      episodeId: expect.any(String),
+      timestamp: expect.any(Number),
+    })
+  );
+});
 
-          return res as AgentGenerateTextResult;
-        },
-      } as any as AIAdapter,
-    });
+test('agent.getPlans() returns plans from context', () => {
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+    planner: async (agent) => {
+      return {
+        episodeId: agent.episodeId,
+        planner: 'test-planner',
+        goal: '',
+        goalState: undefined,
+        paths: [
+          {
+            state: undefined,
+            steps: [],
+          },
+        ],
+        nextEvent: undefined,
+        timestamp: Date.now(),
+      };
+    },
+  });
 
-    await agent[method]({
-      prompt: 'hi',
-    });
+  const plans = agent.getPlans();
+
+  expect(plans).toBeDefined();
+  expect(Array.isArray(plans)).toBe(true);
+});
+
+test('Event listeners can be unsubscribed', () => {
+  const fn = vi.fn();
+  const agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
 
-    const messages = agent.getMessages();
+  const subscription = agent.on('message', fn);
 
-    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');
+  agent.addMessage({
+    role: 'user',
+    content: [{ type: 'text', text: 'first message' }],
+  });
 
-    expect(messages[0]!.correlationId).toEqual(messages[1]!.correlationId);
-  }
-);
+  expect(fn).toHaveBeenCalledTimes(1);
 
-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',
-          };
+  subscription.unsubscribe();
 
-          opts.onFinish?.(res);
+  agent.addMessage({
+    role: 'user',
+    content: [{ type: 'text', text: 'second message' }],
+  });
 
-          return res as AgentGenerateTextResult;
-        },
-      } as any as AIAdapter,
-    });
+  expect(fn).toHaveBeenCalledTimes(1); // Still only called once
+});
 
-    await agent[method]({
-      prompt: 'hi',
-      correlationId: 'c-1',
-      parentCorrelationId: 'c-0',
-    });
+test('agent.observe() adds observations from actor snapshots', () => {
+  const machine = createMachine({
+    initial: 'idle',
+    states: {
+      idle: {
+        on: { START: 'running' },
+      },
+      running: {},
+    },
+  });
 
-    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 agent = createAgent({
+    id: 'test',
+    events: {},
+    model: {} as any,
+  });
 
-    const res = await agent[method]({
-      prompt: 'test',
-    });
+  const actor = createActor(machine);
+  const subscription = agent.observe(actor);
 
-    agent.addFeedback({
-      correlationId: res.correlationId,
-      reward: -1,
-    });
+  actor.start();
+  actor.send({ type: 'START' });
 
-    const message = agent.getMessages()[0]!;
-    const feedback = agent.getFeedback()[0]!;
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      state: expect.objectContaining({ value: 'idle' }),
+      machineHash: expect.any(String),
+    })
+  );
+
+  expect(agent.getObservations()).toContainEqual(
+    expect.objectContaining({
+      prevState: expect.objectContaining({ value: 'idle' }),
+      event: { type: 'START' },
+      state: expect.objectContaining({ value: 'running' }),
+      machineHash: expect.any(String),
+    })
+  );
 
-    expect(message.correlationId).toBeDefined();
-    expect(feedback.correlationId).toEqual(message.correlationId);
-  }
-);
+  subscription.unsubscribe();
+});