@statelyai/agent 2.0.0-next.0 → 2.0.0-next.1

package/src/decide.ts

@@ -10,13 +10,13 @@ import {
   AgentPlanInput,
   TransitionData,
 } from './types';
-import { simplePlanner } from './planners/simplePlanner';
+import { simplePlanner } from './planners/simple';
 import { getTransitions } from './utils';
-import { CoreTool, tool } from 'ai';
+import { CoreMessage, CoreTool, tool } from 'ai';
 
 export async function agentDecide<T extends AnyAgent>(
   agent: T,
-  options: AgentDecideOptions
+  options: AgentDecideOptions<T>
 ): Promise<AgentPlan<EventsFromZodEventMapping<T['events']>> | undefined> {
   const resolvedOptions = {
     ...agent.defaultOptions,
@@ -29,21 +29,31 @@ export async function agentDecide<T extends AnyAgent>(
     state,
     machine,
     model = agent.model,
+    messages,
     ...otherPlanInput
   } = resolvedOptions;
 
-  const plan = await planner(agent, {
-    model,
-    goal,
-    events,
-    state,
-    machine,
-    ...otherPlanInput,
-  });
+  let attempts = 0;
+
+  const maxAttempts = resolvedOptions.maxAttempts ?? 2;
+
+  let plan;
 
-  if (plan?.nextEvent) {
-    agent.addPlan(plan);
-    await resolvedOptions.execute?.(plan.nextEvent);
+  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;

package/src/middleware.ts

@@ -17,15 +17,11 @@ export function createAgentMiddleware(agent: AnyAgent) {
     wrapGenerate: async ({ doGenerate, params }) => {
       const id = randomId();
 
-      params.prompt.forEach((p) => {
+      params.prompt.forEach((message) => {
         agent.addMessage({
           id,
-          ...p,
+          ...message,
           timestamp: Date.now(),
-          correlationId: params.providerMetadata
-            ?.correlationId as unknown as string,
-          parentCorrelationId: params.providerMetadata
-            ?.parentCorrelationId as unknown as string,
         });
       });
 
@@ -43,10 +39,6 @@ export function createAgentMiddleware(agent: AnyAgent) {
           id,
           ...message,
           timestamp: Date.now(),
-          correlationId: params.providerMetadata
-            ?.correlationId as unknown as string,
-          parentCorrelationId: params.providerMetadata
-            ?.parentCorrelationId as unknown as string,
         });
       });
 
@@ -85,10 +77,6 @@ export function createAgentMiddleware(agent: AnyAgent) {
             role: 'assistant',
             content,
             responseId: id,
-            correlationId: params.providerMetadata
-              ?.correlationId as unknown as string,
-            parentCorrelationId: params.providerMetadata
-              ?.parentCorrelationId as unknown as string,
           });
         },
       });

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

package/src/planners/shortestPath.ts

@@ -0,0 +1,177 @@
+import { generateObject } from 'ai';
+import {
+  AgentPlan,
+  AgentPlanInput,
+  AgentPlanner,
+  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';
+import { AnyMachineSnapshot } from 'xstate';
+
+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 function experimental_createShortestPathPlanner<
+  T extends AnyAgent
+>(): AgentPlanner<T> {
+  return async function shortestPathPlanner<T extends AnyAgent>(
+    agent: T,
+    input: AgentPlanInput<any>
+  ): Promise<AgentPlan<any> | undefined> {
+    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,
+        system: input.system ?? agent.description,
+        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"]}"
+  }
+
+  3. For "fruits contain apple, orange, banana":
+  {
+    "type": "array",
+    "allOf": [
+      { "contains": { "const": "apple" } },
+      { "contains": { "const": "orange" } },
+      { "contains": { "const": "banana" } }
+    ]
+  }
+      `.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 stateFilter = (state: AnyMachineSnapshot) => {
+        return validateContext(state.context);
+      };
+
+      const resolvedState = input.machine.resolveState({
+        ...input.state,
+        context: input.state.context ?? {},
+      });
+
+      paths = getShortestPaths(input.machine, {
+        fromState: resolvedState,
+        toState: stateFilter,
+      });
+    }
+
+    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);
+
+    // Sort paths from least weight to most weight
+    const sortedPaths = trimmedPaths.sort(
+      (a, b) => costFunction(a) - costFunction(b)
+    );
+
+    const leastWeightPath = sortedPaths[0];
+    const nextStep = leastWeightPath?.steps[0];
+
+    return {
+      planner: 'shortestPath',
+      episodeId: agent.episodeId,
+      goal: input.goal,
+      goalState: paths[0]?.state,
+      nextEvent: nextStep?.event,
+      paths,
+      timestamp: Date.now(),
+    };
+  };
+}

package/src/planners/{simplePlanner.ts → simple.ts}

@@ -1,14 +1,7 @@
-import { CoreMessage, type CoreTool, generateText, tool } from 'ai';
-import {
-  AgentPlan,
-  AgentPlanInput,
-  ObservedState,
-  PromptTemplate,
-  TransitionData,
-  AnyAgent,
-} from '../types';
-import { getAllTransitions, randomId } from '../utils';
-import { AnyStateMachine, getNextSnapshot } from 'xstate';
+import { CoreMessage, generateText } from 'ai';
+import { AgentPlan, AgentPlanInput, PromptTemplate, AnyAgent } from '../types';
+import { randomId } from '../utils';
+import { getNextSnapshot } from 'xstate';
 import { defaultTextTemplate } from '../templates/defaultText';
 import { getMessages } from '../text';
 import { getToolMap } from '../decide';
@@ -62,13 +55,14 @@ export async function simplePlanner<T extends AnyAgent>(
 
   const result = await generateText({
     ...rest,
+    system: input.system ?? agent.description,
     model,
     messages,
     tools: toolMap as any,
     toolChoice: input.toolChoice ?? 'required',
   });
 
-  result.responseMessages.forEach((m) => {
+  result.response.messages.forEach((m) => {
     const message: CoreMessage = m;
 
     agent.addMessage({

package/src/types.ts

@@ -32,7 +32,7 @@ export type CostFunction<TEvent extends EventObject> = (
 ) => number;
 
 export type AgentPlanInput<TEvent extends EventObject> = Omit<
-  GenerateTextOptions,
+  AgentGenerateTextOptions,
   'prompt' | 'tools'
 > & {
   /**
@@ -63,6 +63,12 @@ export type AgentPlanInput<TEvent extends EventObject> = Omit<
    * The total cost of the path to the goal state.
    */
   costFunction?: CostFunction<TEvent>;
+
+  /**
+   * The maximum number of attempts to generate a plan.
+   * Defaults to 2.
+   */
+  maxAttempts?: number;
 };
 
 export type AgentStep<TEvent extends EventObject> = {
@@ -146,14 +152,19 @@ export type AgentPlanner<T extends AnyAgent> = (
   input: AgentPlanInput<T['types']['events']>
 ) => Promise<AgentPlan<T['types']['events']> | undefined>;
 
-export type AgentDecideOptions = {
+export type AgentDecideOptions<T extends AnyAgent> = {
   goal: string;
   model?: LanguageModel;
   state: ObservedState;
   machine?: AnyStateMachine;
   execute?: (event: AnyEventObject) => Promise<void>;
-  planner?: AgentPlanner<any>;
+  planner?: AgentPlanner<T>;
   events?: ZodEventMapping;
+  /**
+   * The maximum number of times the agent will attempt to make a decision.
+   * Defaults to 2.
+   */
+  maxAttempts?: number;
 } & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt'>;
 
 export interface AgentFeedback {
@@ -162,7 +173,6 @@ export interface AgentFeedback {
   /**
    * The message correlation that the feedback is relevant for
    */
-  correlationId?: string;
   attributes: Record<string, any>;
   reward: number;
   timestamp: number;
@@ -172,7 +182,6 @@ export interface AgentFeedback {
 export interface AgentFeedbackInput {
   goal?: string;
   observationId?: string;
-  correlationId?: string;
   attributes?: Record<string, any>;
   timestamp?: number;
   reward?: number;
@@ -313,8 +322,6 @@ export type AgentMessageInput = CoreMessage & {
    * which message this message is responding to, if any.
    */
   responseId?: string;
-  correlationId?: string;
-  parentCorrelationId?: string;
   result?: GenerateTextResult<any>;
 };
 
@@ -402,7 +409,7 @@ export type ContextFromZodContextMapping<
   [K in keyof TContextSchema & string]: TypeOf<TContextSchema[K]>;
 };
 
-export type AnyAgent = Agent<any, any>;
+export type AnyAgent = Agent<any, any, any, any>;
 
 export type FromAgent<T> = T | ((agent: AnyAgent) => T | Promise<T>);
 

package/src/utils.ts

@@ -103,3 +103,14 @@ export function getTransitions(
   });
   return getAllTransitions(resolvedState);
 }
+
+export function isMachineActor(
+  actor: ActorRefLike
+): actor is typeof actor & { src: AnyStateMachine } {
+  return (
+    'src' in actor &&
+    typeof actor.src === 'object' &&
+    actor.src !== null &&
+    'definition' in actor.src
+  );
+}

package/vitest.config.ts

@@ -1,9 +1,15 @@
-// vitest.config.ts
+import { defineConfig } from 'vitest/config';
 import dotenv from 'dotenv';
 dotenv.config();
 
-export default {
+export default defineConfig({
   test: {
     testTimeout: 10000, // Global timeout of 10000ms for all tests
+    coverage: {
+      provider: 'v8',
+      reporter: ['text', 'json', 'html'],
+      exclude: ['**.test.ts'],
+      include: ['src'],
+    },
   },
-};
+});

package/src/planners/shortestPathPlanner.ts

@@ -1,160 +0,0 @@
-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';
-
-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> {
-  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);
-
-  // Sort paths from least weight to most weight
-  const sortedPaths = trimmedPaths.sort(
-    (a, b) => costFunction(a) - costFunction(b)
-  );
-
-  const leastWeightPath = sortedPaths[0];
-  const nextStep = leastWeightPath?.steps[0];
-
-  return {
-    planner: 'shortestPath',
-    episodeId: agent.episodeId,
-    goal: input.goal,
-    goalState: paths[0]?.state,
-    nextEvent: nextStep?.event,
-    paths,
-    timestamp: Date.now(),
-  };
-}