@statelyai/agent 1.1.6 → 2.0.0-next.0

package/src/{decision.test.ts → decide.test.ts}

@@ -1,61 +1,39 @@
 import { test, expect } from 'vitest';
-import { createAgent, fromDecision, type AIAdapter } from './';
+import { createAgent, fromDecision } from '.';
 import { createActor, createMachine, waitFor } from 'xstate';
 import { z } from 'zod';
-import { GenerateTextResult } from 'ai';
+import { LanguageModelV1CallOptions } from 'ai';
+import { dummyResponseValues, MockLanguageModelV1 } from './mockModel';
 
-const mockToolDecision: AIAdapter['generateText'] = async (arg) => {
-  const keys = Object.keys(arg.tools!);
-
-  if (keys.length > 1) {
-    throw new Error('Expected only 1 choice');
-  }
-
-  if (keys.length === 0) {
-    return {
-      toolResults: [],
-    } as any as GenerateTextResult<any>;
-  }
+const doGenerate = async (params: LanguageModelV1CallOptions) => {
+  const keys =
+    params.mode.type === 'regular' ? params.mode.tools?.map((t) => t.name) : [];
 
   return {
-    toolResults: [
+    ...dummyResponseValues,
+    finishReason: 'tool-calls',
+    toolCalls: [
       {
-        result: {
-          type: keys[0],
-        },
+        toolCallType: 'function',
+        toolCallId: 'call-1',
+        toolName: keys![0],
+        args: `{ "type": "${keys?.[0]}" }`,
       },
     ],
-  } as any as GenerateTextResult<any>;
+  } as any;
 };
 
 test('fromDecision() makes a decision', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
   const agent = createAgent({
     name: 'test',
-    model: {} as any,
+    model,
     events: {
       doFirst: z.object({}),
       doSecond: 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 GenerateTextResult<any>;
-      },
-      streamText: {} as any,
-    },
   });
 
   const machine = createMachine({
@@ -91,17 +69,16 @@ test('fromDecision() makes a decision', async () => {
 });
 
 test('interacts with an actor', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
   const agent = createAgent({
     name: 'test',
-    model: {} as any,
+    model,
     events: {
       doFirst: z.object({}),
       doSecond: z.object({}),
     },
-    adapter: {
-      generateText: mockToolDecision,
-      streamText: {} as any,
-    },
   });
 
   const machine = createMachine({
@@ -135,17 +112,16 @@ test('interacts with an actor', async () => {
 });
 
 test('interacts with an actor (late interaction)', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
   const agent = createAgent({
     name: 'test',
-    model: {} as any,
+    model,
     events: {
       doFirst: z.object({}),
       doSecond: z.object({}),
     },
-    adapter: {
-      generateText: mockToolDecision,
-      streamText: {} as any,
-    },
   });
 
   const machine = createMachine({

package/src/decide.ts

@@ -0,0 +1,153 @@
+import { AnyActor, AnyMachineSnapshot, fromPromise } from 'xstate';
+import {
+  AnyAgent,
+  AgentDecideOptions,
+  AgentDecisionLogic,
+  AgentDecisionInput,
+  AgentPlanner,
+  AgentPlan,
+  EventsFromZodEventMapping,
+  AgentPlanInput,
+  TransitionData,
+} from './types';
+import { simplePlanner } from './planners/simplePlanner';
+import { getTransitions } from './utils';
+import { CoreTool, tool } from 'ai';
+
+export async function agentDecide<T extends AnyAgent>(
+  agent: T,
+  options: AgentDecideOptions
+): 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,
+    ...otherPlanInput
+  } = resolvedOptions;
+
+  const plan = await planner(agent, {
+    model,
+    goal,
+    events,
+    state,
+    machine,
+    ...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,103 @@
+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((p) => {
+        agent.addMessage({
+          id,
+          ...p,
+          timestamp: Date.now(),
+          correlationId: params.providerMetadata
+            ?.correlationId as unknown as string,
+          parentCorrelationId: params.providerMetadata
+            ?.parentCorrelationId as unknown as string,
+        });
+      });
+
+      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(),
+          correlationId: params.providerMetadata
+            ?.correlationId as unknown as string,
+          parentCorrelationId: params.providerMetadata
+            ?.parentCorrelationId as unknown as string,
+        });
+      });
+
+      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,
+            correlationId: params.providerMetadata
+              ?.correlationId as unknown as string,
+            parentCorrelationId: params.providerMetadata
+              ?.parentCorrelationId as unknown as string,
+          });
+        },
+      });
+
+      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/shortestPathPlanner.ts

@@ -1,22 +1,160 @@
-import { Agent, AgentPlan, AgentPlanInput, AnyAgent } from '../types';
+import { generateObject } from 'ai';
+import { getToolMap } from '../decide';
+import {
+  AgentPlan,
+  AgentPlanInput,
+  AgentStep,
+  AnyAgent,
+  CostFunction,
+  ObservedState,
+} from '../types';
 import { getShortestPaths } from '@xstate/graph';
+import { z } from 'zod';
+import { zodToJsonSchema } from 'zod-to-json-schema';
+import Ajv from 'ajv';
 
-export async function simplePlanner<T extends AnyAgent>(
+const ajv = new Ajv();
+
+function observedStatesEqual(state1: ObservedState, state2: ObservedState) {
+  // check state value && state context
+  return (
+    JSON.stringify(state1.value) === JSON.stringify(state2.value) &&
+    JSON.stringify(state1.context) === JSON.stringify(state2.context)
+  );
+}
+
+function trimSteps(steps: AgentStep<any>[], currentState: ObservedState) {
+  const index = steps.findIndex(
+    (step) => step.state && observedStatesEqual(step.state, currentState)
+  );
+
+  if (index === -1) {
+    return undefined;
+  }
+
+  return steps.slice(index + 1, steps.length);
+}
+
+export async function shortestPathPlanner<T extends AnyAgent>(
   agent: T,
   input: AgentPlanInput<any>
 ): Promise<AgentPlan<any> | undefined> {
-  // 1. Determine goal state criteria
-  // e.g. a state where the agent has won a game
-  void 0;
+  const costFunction: CostFunction<any> =
+    input.costFunction ?? ((path) => path.weight ?? Infinity);
+  const existingPlan = agent
+    .getPlans()
+    .find((p) => p.planner === 'shortestPath' && p.goal === input.goal);
+
+  let paths = existingPlan?.paths;
+
+  if (existingPlan) {
+    console.log('Existing plan found');
+  }
+
+  if (!input.machine && !existingPlan) {
+    return;
+  }
+
+  if (input.machine && !existingPlan) {
+    const contextSchema = zodToJsonSchema(z.object(agent.context));
+    const result = await generateObject({
+      model: agent.model,
+      prompt: `
+<goal>
+${input.goal}
+</goal>
+<contextSchema>
+${contextSchema}
+</contextSchema>
+
+
+Update the context JSON schema so that it validates the context to determine that it reaches the goal. Return the result as a diff.
+
+The contextSchema properties must not change. Do not add or remove properties, or modify the name of the properties.
+Use "const" for exact required values and define ranges/types for flexible conditions.
+
+Examples:
+1. For "user is logged in with admin role":
+{
+  "contextSchema": "{"type": "object", "properties": {"role": {"const": "admin"}, "lastLogin": {"type": "string"}}, "required": ["role"]}"
+}
+
+2. For "score is above 100":
+{
+  "contextSchema": "{"type": "object", "properties": {"score": {"type": "number", "minimum": 100}}, "required": ["score"]}"
+}
+    `.trim(),
+      schema: z.object({
+        // valueSchema: z
+        //   .string()
+        //   .describe('The JSON Schema representing the goal state value'),
+        contextSchema: z
+          .object({
+            type: z.literal('object'),
+            properties: z.object(
+              Object.keys((contextSchema as any).properties).reduce(
+                (acc, key) => {
+                  acc[key] = z.any();
+                  return acc;
+                },
+                {} as any
+              )
+            ),
+            required: z.array(z.string()).optional(),
+          })
+          .describe('The JSON Schema representing the goal state context'),
+      }),
+    });
+
+    console.log(result.object);
+    const validateContext = ajv.compile(result.object.contextSchema);
+
+    const resolvedState = input.machine.resolveState({
+      ...input.state,
+      context: input.state.context ?? {},
+    });
+
+    paths = getShortestPaths(input.machine, {
+      fromState: resolvedState,
+      toState: (state) => {
+        const v = validateContext(state.context);
+        return v;
+      },
+    });
+  }
+
+  if (!paths) {
+    return undefined;
+  }
+
+  const trimmedPaths = paths
+    .map((path) => {
+      const trimmedSteps = trimSteps(path.steps, input.state);
+      if (!trimmedSteps) {
+        return undefined;
+      }
+      return {
+        ...path,
+        steps: trimmedSteps,
+      };
+    })
+    .filter((p): p is NonNullable<typeof p> => p !== undefined);
 
-  // 2. Determine possible events that can occur
-  void 0;
+  // Sort paths from least weight to most weight
+  const sortedPaths = trimmedPaths.sort(
+    (a, b) => costFunction(a) - costFunction(b)
+  );
 
-  // 3. Get shortest paths from current state to
-  // a state matching the criteria, using
-  // possible events
-  void 0;
+  const leastWeightPath = sortedPaths[0];
+  const nextStep = leastWeightPath?.steps[0];
 
-  // 4. Return shortest path as a plan
-  return null as any;
+  return {
+    planner: 'shortestPath',
+    episodeId: agent.episodeId,
+    goal: input.goal,
+    goalState: paths[0]?.state,
+    nextEvent: nextStep?.event,
+    paths,
+    timestamp: Date.now(),
+  };
 }