@statelyai/agent 1.1.6 → 2.0.0-next.1

package/src/decide.test.ts

@@ -0,0 +1,276 @@
+import { test, expect } from 'vitest';
+import { createAgent, fromDecision } from '.';
+import { createActor, createMachine, waitFor } from 'xstate';
+import { z } from 'zod';
+import { LanguageModelV1CallOptions } from 'ai';
+import { dummyResponseValues, MockLanguageModelV1 } from './mockModel';
+
+const doGenerate = async (params: LanguageModelV1CallOptions) => {
+  const keys =
+    params.mode.type === 'regular' ? params.mode.tools?.map((t) => t.name) : [];
+
+  return {
+    ...dummyResponseValues,
+    finishReason: 'tool-calls',
+    toolCalls: [
+      {
+        toolCallType: 'function',
+        toolCallId: 'call-1',
+        toolName: keys![0],
+        args: `{ "type": "${keys?.[0]}" }`,
+      },
+    ],
+  } as any;
+};
+
+test('fromDecision() makes a decision', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
+  const agent = createAgent({
+    id: 'test',
+    model,
+    events: {
+      doFirst: z.object({}),
+      doSecond: z.object({}),
+    },
+  });
+
+  const machine = createMachine({
+    initial: 'first',
+    states: {
+      first: {
+        invoke: {
+          src: fromDecision(agent),
+        },
+        on: {
+          doFirst: 'second',
+        },
+      },
+      second: {
+        invoke: {
+          src: fromDecision(agent),
+        },
+        on: {
+          doSecond: 'third',
+        },
+      },
+      third: {},
+    },
+  });
+
+  const actor = createActor(machine);
+
+  actor.start();
+
+  await waitFor(actor, (s) => s.matches('third'));
+
+  expect(actor.getSnapshot().value).toBe('third');
+});
+
+test('interacts with an actor', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
+  const agent = createAgent({
+    id: 'test',
+    model,
+    events: {
+      doFirst: z.object({}),
+      doSecond: z.object({}),
+    },
+  });
+
+  const machine = createMachine({
+    initial: 'first',
+    states: {
+      first: {
+        on: {
+          doFirst: 'second',
+        },
+      },
+      second: {
+        on: {
+          doSecond: 'third',
+        },
+      },
+      third: {},
+    },
+  });
+
+  const actor = createActor(machine);
+
+  agent.interact(actor, () => ({
+    goal: 'Some goal',
+  }));
+
+  actor.start();
+
+  await waitFor(actor, (s) => s.matches('third'));
+
+  expect(actor.getSnapshot().value).toBe('third');
+});
+
+test('interacts with an actor (late interaction)', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
+  const agent = createAgent({
+    id: 'test',
+    model,
+    events: {
+      doFirst: z.object({}),
+      doSecond: z.object({}),
+    },
+  });
+
+  const machine = createMachine({
+    initial: 'first',
+    states: {
+      first: {
+        on: {
+          doFirst: 'second',
+        },
+      },
+      second: {
+        on: {
+          doSecond: 'third',
+        },
+      },
+      third: {},
+    },
+  });
+
+  const actor = createActor(machine);
+
+  actor.start();
+
+  agent.interact(actor, () => ({
+    goal: 'Some goal',
+  }));
+
+  await waitFor(actor, (s) => s.matches('third'));
+
+  expect(actor.getSnapshot().value).toBe('third');
+});
+
+test('agent.decide() makes a decision based on goal and state (simple planner)', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
+
+  const agent = createAgent({
+    id: 'test',
+    model,
+    events: {
+      MOVE: z.object({}),
+    },
+  });
+
+  const plan = await agent.decide({
+    goal: 'Make the best move',
+    state: {
+      value: 'playing',
+      context: {
+        board: [0, 0, 0],
+      },
+    },
+    machine: createMachine({
+      initial: 'playing',
+      states: {
+        playing: {
+          on: {
+            MOVE: 'next',
+          },
+        },
+        next: {},
+      },
+    }),
+  });
+
+  expect(plan).toBeDefined();
+  expect(plan!.nextEvent).toEqual(
+    expect.objectContaining({
+      type: 'MOVE',
+    })
+  );
+});
+
+test.each([
+  [undefined, true],
+  [undefined, false],
+  [3, true],
+  [3, false],
+])(
+  'agent.decide() retries if a decision is not made (%i attempts, succeed: %s)',
+  async (maxAttempts, succeed) => {
+    let attempts = 0;
+    const doGenerateWithRetry = async (params: LanguageModelV1CallOptions) => {
+      const keys =
+        params.mode.type === 'regular'
+          ? params.mode.tools?.map((t) => t.name)
+          : [];
+
+      console.log('try', attempts, 'max', maxAttempts);
+
+      const toolCalls =
+        succeed && attempts++ === (maxAttempts ?? 2) - 1
+          ? [
+              {
+                toolCallType: 'function',
+                toolCallId: 'call-1',
+                toolName: keys![0],
+                args: `{ "type": "${keys?.[0]}" }`,
+              },
+            ]
+          : [];
+
+      return {
+        ...dummyResponseValues,
+        finishReason: 'tool-calls',
+        toolCalls,
+      } as any;
+    };
+    const model = new MockLanguageModelV1({
+      doGenerate: doGenerateWithRetry,
+    });
+
+    const agent = createAgent({
+      id: 'test',
+      model,
+      events: {
+        MOVE: z.object({}),
+      },
+    });
+
+    const plan = await agent.decide({
+      goal: 'Make the best move',
+      state: {
+        value: 'playing',
+      },
+      machine: createMachine({
+        initial: 'playing',
+        states: {
+          playing: {
+            on: {
+              MOVE: 'win',
+            },
+          },
+          win: {},
+        },
+      }),
+      maxAttempts,
+    });
+
+    if (!succeed) {
+      expect(plan).toBeUndefined();
+    } else {
+      expect(plan).toBeDefined();
+      expect(plan!.nextEvent).toEqual(
+        expect.objectContaining({
+          type: 'MOVE',
+        })
+      );
+    }
+  }
+);

package/src/decide.ts

@@ -0,0 +1,163 @@
+import { AnyActor, AnyMachineSnapshot, fromPromise } from 'xstate';
+import {
+  AnyAgent,
+  AgentDecideOptions,
+  AgentDecisionLogic,
+  AgentDecisionInput,
+  AgentPlanner,
+  AgentPlan,
+  EventsFromZodEventMapping,
+  AgentPlanInput,
+  TransitionData,
+} from './types';
+import { simplePlanner } from './planners/simple';
+import { getTransitions } from './utils';
+import { CoreMessage, CoreTool, tool } from 'ai';
+
+export async function agentDecide<T extends AnyAgent>(
+  agent: T,
+  options: AgentDecideOptions<T>
+): Promise<AgentPlan<EventsFromZodEventMapping<T['events']>> | undefined> {
+  const resolvedOptions = {
+    ...agent.defaultOptions,
+    ...options,
+  };
+  const {
+    planner = simplePlanner as AgentPlanner<any>,
+    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: messages as CoreMessage[], // TODO: fix UIMessage thing
+      ...otherPlanInput,
+    });
+
+    if (plan?.nextEvent) {
+      agent.addPlan(plan);
+      await resolvedOptions.execute?.(plan.nextEvent);
+    }
+  }
+
+  return plan;
+}
+
+export function fromDecision(
+  agent: AnyAgent,
+  defaultInput?: AgentDecisionInput
+): AgentDecisionLogic<any> {
+  return fromPromise(async ({ input, self }) => {
+    const parentRef = self._parent;
+    if (!parentRef) {
+      return;
+    }
+
+    const snapshot = parentRef.getSnapshot() as AnyMachineSnapshot;
+    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 as AnyActor).logic,
+      state,
+      execute: async (event) => {
+        parentRef.send(event);
+      },
+      ...resolvedInput,
+    });
+
+    return plan;
+  }) as AgentDecisionLogic<any>;
+}
+
+export function getToolMap<T extends AnyAgent>(
+  _agent: T,
+  input: AgentPlanInput<any>
+): Record<string, CoreTool<any, any>> | undefined {
+  // Get all of the possible next transitions
+  const transitions: TransitionData[] = input.machine
+    ? getTransitions(input.state, input.machine)
+    : Object.entries(input.events).map(([eventType, { description }]) => ({
+        eventType,
+        description,
+      }));
+
+  // Only keep the transitions that match the event types that are in the event mapping
+  // TODO: allow for custom filters
+  const filter = (eventType: string) =>
+    Object.keys(input.events).includes(eventType);
+
+  // Mapping of each event type (e.g. "mouse.click")
+  // to a valid function name (e.g. "mouse_click")
+  const functionNameMapping: Record<string, string> = {};
+
+  const toolTransitions = transitions
+    .filter((t) => {
+      return filter(t.eventType);
+    })
+    .map((t) => {
+      const name = t.eventType.replace(/\./g, '_');
+      functionNameMapping[name] = t.eventType;
+
+      return {
+        type: 'function',
+        eventType: t.eventType,
+        description: t.description,
+        name,
+      } as const;
+    });
+
+  // Convert the transition data to a tool map that the
+  // Vercel AI SDK can use
+  const toolMap: Record<string, CoreTool<any, any>> = {};
+  for (const toolTransitionData of toolTransitions) {
+    const toolZodType = input.events?.[toolTransitionData.eventType];
+
+    if (!toolZodType) {
+      continue;
+    }
+
+    toolMap[toolTransitionData.name] = tool({
+      description: toolZodType?.description ?? toolTransitionData.description,
+      parameters: toolZodType,
+      execute: async (params: Record<string, any>) => {
+        const event = {
+          type: toolTransitionData.eventType,
+          ...params,
+        };
+
+        return event;
+      },
+    });
+  }
+
+  if (!Object.keys(toolMap).length) {
+    // No valid transitions for the specified tools
+    return undefined;
+  }
+
+  return toolMap;
+}

package/src/index.ts

@@ -1,4 +1,4 @@
 export { createAgent } from './agent';
 export { fromText, fromTextStream } from './text';
-export { fromDecision } from './decision';
+export { fromDecision } from './decide';
 export * from './types';

package/src/middleware.ts

@@ -0,0 +1,91 @@
+import {
+  Experimental_LanguageModelV1Middleware as LanguageModelV1Middleware,
+  LanguageModelV1StreamPart,
+} from 'ai';
+import {
+  AnyAgent,
+  LanguageModelV1TextPart,
+  LanguageModelV1ToolCallPart,
+} from './types';
+import { randomId } from './utils';
+
+export function createAgentMiddleware(agent: AnyAgent) {
+  const middleware: LanguageModelV1Middleware = {
+    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<
+        LanguageModelV1StreamPart,
+        LanguageModelV1StreamPart
+      >({
+        transform(chunk, controller) {
+          if (chunk.type === 'text-delta') {
+            generatedText += chunk.textDelta;
+          }
+
+          controller.enqueue(chunk);
+        },
+
+        flush() {
+          const content: (
+            | LanguageModelV1TextPart
+            | LanguageModelV1ToolCallPart
+          )[] = [];
+
+          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,
+      };
+    },
+  };
+  return middleware;
+}

package/src/mockModel.ts

@@ -0,0 +1,47 @@
+import { LanguageModelV1 } from 'ai';
+
+export class MockLanguageModelV1 implements LanguageModelV1 {
+  readonly specificationVersion = 'v1';
+
+  readonly provider: LanguageModelV1['provider'];
+  readonly modelId: LanguageModelV1['modelId'];
+
+  doGenerate: LanguageModelV1['doGenerate'];
+  doStream: LanguageModelV1['doStream'];
+
+  readonly defaultObjectGenerationMode: LanguageModelV1['defaultObjectGenerationMode'];
+  readonly supportsStructuredOutputs: LanguageModelV1['supportsStructuredOutputs'];
+  constructor({
+    provider = 'mock-provider',
+    modelId = 'mock-model-id',
+    doGenerate = notImplemented,
+    doStream = notImplemented,
+    defaultObjectGenerationMode = undefined,
+    supportsStructuredOutputs = undefined,
+  }: {
+    provider?: LanguageModelV1['provider'];
+    modelId?: LanguageModelV1['modelId'];
+    doGenerate?: LanguageModelV1['doGenerate'];
+    doStream?: LanguageModelV1['doStream'];
+    defaultObjectGenerationMode?: LanguageModelV1['defaultObjectGenerationMode'];
+    supportsStructuredOutputs?: LanguageModelV1['supportsStructuredOutputs'];
+  } = {}) {
+    this.provider = provider;
+    this.modelId = modelId;
+    this.doGenerate = doGenerate;
+    this.doStream = doStream;
+
+    this.defaultObjectGenerationMode = defaultObjectGenerationMode;
+    this.supportsStructuredOutputs = supportsStructuredOutputs;
+  }
+}
+
+function notImplemented(): never {
+  throw new Error('Not implemented');
+}
+
+export const dummyResponseValues = {
+  rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+  finishReason: 'stop' as const,
+  usage: { promptTokens: 10, completionTokens: 20 },
+};

package/src/planners/shortestPath.test.ts

@@ -0,0 +1,94 @@
+import { createAgent } from '../';
+import { assign, createActor, setup } from 'xstate';
+import { z } from 'zod';
+import { experimental_createShortestPathPlanner } from './shortestPath';
+import { test, expect } from 'vitest';
+import { dummyResponseValues, MockLanguageModelV1 } from '../mockModel';
+
+test.skip('should find shortest path to goal', async () => {
+  const agent = createAgent({
+    id: 'counter',
+    model: new MockLanguageModelV1({
+      doGenerate: async () => {
+        return {
+          ...dummyResponseValues,
+          text: JSON.stringify({
+            type: 'object',
+            properties: {
+              count: {
+                type: 'number',
+                const: 3,
+              },
+            },
+            required: ['count'],
+          }),
+        };
+      },
+    }),
+    events: {
+      increment: z.object({}).describe('Increment the counter by 1'),
+      decrement: z.object({}).describe('Decrement the counter by 1'),
+    },
+    context: {
+      count: z.number().int().describe('Current count value'),
+    },
+  });
+
+  const counterMachine = setup({
+    types: {
+      context: agent.types.context,
+      events: agent.types.events,
+    },
+  }).createMachine({
+    initial: 'counting',
+    context: { count: 0 },
+    states: {
+      counting: {
+        always: {
+          guard: ({ context }) => context.count === 3,
+          target: 'success',
+        },
+        on: {
+          increment: {
+            actions: assign({ count: ({ context }) => context.count + 1 }),
+          },
+          decrement: {
+            actions: assign({ count: ({ context }) => context.count - 1 }),
+          },
+        },
+      },
+      success: {
+        type: 'final',
+      },
+    },
+  });
+
+  const counterActor = createActor(counterMachine).start();
+
+  const decision = await agent.decide({
+    machine: counterMachine,
+    model: new MockLanguageModelV1({
+      defaultObjectGenerationMode: 'tool',
+      doGenerate: async () => {
+        return {
+          ...dummyResponseValues,
+          text: JSON.stringify({
+            type: 'object',
+            properties: {
+              count: {
+                type: 'number',
+                const: 3,
+              },
+            },
+            required: ['count'],
+          }),
+        };
+      },
+    }),
+    goal: 'Get the counter to exactly 3',
+    state: counterActor.getSnapshot(),
+    planner: experimental_createShortestPathPlanner(),
+  });
+
+  expect(decision?.nextEvent?.type).toBe('increment');
+});