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

package/src/decide.ts

@@ -3,55 +3,72 @@ import {
   AnyAgent,
   AgentDecideOptions,
   AgentDecisionLogic,
-  AgentDecisionInput,
-  AgentPlanner,
-  AgentPlan,
-  EventsFromZodEventMapping,
-  AgentPlanInput,
+  AgentDecision,
+  AgentDecideInput,
   TransitionData,
+  EventsFromAgent,
 } from './types';
-import { simplePlanner } from './planners/simplePlanner';
 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
-): Promise<AgentPlan<EventsFromZodEventMapping<T['events']>> | undefined> {
+  options: AgentDecideOptions<T>
+): Promise<AgentDecision<EventsFromAgent<T>> | undefined> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
   };
   const {
-    planner = simplePlanner as AgentPlanner<any>,
+    strategy = agent.strategy,
     goal,
+    allowedEvents,
     events = agent.events,
     state,
     machine,
     model = agent.model,
-    ...otherPlanInput
+    messages,
+    ...otherDecideInput
   } = resolvedOptions;
 
-  const plan = await planner(agent, {
-    model,
-    goal,
-    events,
-    state,
-    machine,
-    ...otherPlanInput,
-  });
+  const filteredEventSchemas = allowedEvents
+    ? Object.fromEntries(
+        Object.entries(events).filter(([key]) => {
+          return allowedEvents.includes(key);
+        })
+      )
+    : events;
+
+  let attempts = 0;
 
-  if (plan?.nextEvent) {
-    agent.addPlan(plan);
-    await resolvedOptions.execute?.(plan.nextEvent);
+  const maxAttempts = resolvedOptions.maxAttempts ?? 2;
+
+  let decision: AgentDecision<any> | undefined;
+
+  while (attempts++ < maxAttempts) {
+    decision = await strategy(agent, {
+      model,
+      goal,
+      events: filteredEventSchemas,
+      state,
+      machine,
+      messages: messages as CoreMessage[], // TODO: fix UIMessage thing
+      ...otherDecideInput,
+    });
+
+    if (decision?.nextEvent) {
+      agent.addDecision(decision);
+      await resolvedOptions.execute?.(decision.nextEvent);
+      break;
+    }
   }
 
-  return plan;
+  return decision;
 }
 
-export function fromDecision(
-  agent: AnyAgent,
-  defaultInput?: AgentDecisionInput
+export function fromDecision<T extends AnyAgent>(
+  agent: T,
+  defaultInput?: AgentDecideInput<EventsFromAgent<T>>
 ): AgentDecisionLogic<any> {
   return fromPromise(async ({ input, self }) => {
     const parentRef = self._parent;
@@ -70,22 +87,25 @@ export function fromDecision(
       context: resolvedInput.context,
     };
 
-    const plan = await agentDecide(agent, {
+    const decision = await agentDecide(agent, {
       machine: (parentRef as AnyActor).logic,
-      state,
+      state: snapshot,
+      context: resolvedInput.context,
       execute: async (event) => {
         parentRef.send(event);
       },
       ...resolvedInput,
+      // @ts-ignore
+      messages: resolvedInput.messages,
     });
 
-    return plan;
+    return decision;
   }) as AgentDecisionLogic<any>;
 }
 
 export function getToolMap<T extends AnyAgent>(
   _agent: T,
-  input: AgentPlanInput<any>
+  input: AgentDecideInput<any>
 ): Record<string, CoreTool<any, any>> | undefined {
   // Get all of the possible next transitions
   const transitions: TransitionData[] = input.machine

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/strategies/chainOfThought.ts

@@ -0,0 +1,48 @@
+import { generateText } from 'ai';
+import {
+  AnyAgent,
+  AgentDecideInput,
+  AgentDecision,
+  PromptTemplate,
+} from '../types';
+import { getMessages } from '../text';
+import { simpleStrategy } from './simple';
+import { convertToXml } from '../utils';
+
+const chainOfThoughtPromptTemplate: PromptTemplate<any> = ({
+  context,
+  goal,
+}) => {
+  return `
+${convertToXml({ context, goal })}
+
+How would you achieve the goal? Think step-by-step.
+`.trim();
+};
+
+export async function chainOfThoughtStrategy<T extends AnyAgent>(
+  agent: T,
+  input: AgentDecideInput<any>
+): Promise<AgentDecision<any> | undefined> {
+  const prompt = chainOfThoughtPromptTemplate({
+    context: input.state.context,
+    goal: input.goal,
+  });
+
+  const messages = await getMessages(agent, prompt, input);
+
+  const model = input.model ? agent.wrap(input.model) : agent.model;
+
+  const result = await generateText({
+    model,
+    system: input.system ?? agent.description,
+    messages,
+  });
+
+  const decision = await simpleStrategy(agent, {
+    ...input,
+    messages: messages.concat(result.response.messages),
+  });
+
+  return decision;
+}

package/src/strategies/shortestPath.test.ts

@@ -0,0 +1,91 @@
+import { createAgent, TypesFromAgent } from '..';
+import { assign, createActor, setup } from 'xstate';
+import { z } from 'zod';
+import { experimental_shortestPathStrategy } 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: {} as TypesFromAgent<typeof agent>,
+  }).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(),
+    strategy: experimental_shortestPathStrategy,
+  });
+
+  expect(decision?.nextEvent?.type).toBe('increment');
+});

package/src/{planners/shortestPathPlanner.ts → strategies/shortestPath.ts}

@@ -1,8 +1,8 @@
 import { generateObject } from 'ai';
-import { getToolMap } from '../decide';
 import {
-  AgentPlan,
-  AgentPlanInput,
+  AgentDecision,
+  AgentDecideInput,
+  AgentStrategy,
   AgentStep,
   AnyAgent,
   CostFunction,
@@ -12,6 +12,7 @@ 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();
 
@@ -35,30 +36,31 @@ function trimSteps(steps: AgentStep<any>[], currentState: ObservedState) {
   return steps.slice(index + 1, steps.length);
 }
 
-export async function shortestPathPlanner<T extends AnyAgent>(
+export async function experimental_shortestPathStrategy<T extends AnyAgent>(
   agent: T,
-  input: AgentPlanInput<any>
-): Promise<AgentPlan<any> | undefined> {
+  input: AgentDecideInput<any>
+): Promise<AgentDecision<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);
+  const existingDecision = agent
+    .getDecisions()
+    .find((p) => p.strategy === 'shortestPath' && p.goal === input.goal);
 
-  let paths = existingPlan?.paths;
+  let paths = existingDecision?.paths;
 
-  if (existingPlan) {
-    console.log('Existing plan found');
+  if (existingDecision) {
+    console.log('Existing decision found');
   }
 
-  if (!input.machine && !existingPlan) {
+  if (!input.machine && !existingDecision) {
     return;
   }
 
-  if (input.machine && !existingPlan) {
+  if (input.machine && !existingDecision) {
     const contextSchema = zodToJsonSchema(z.object(agent.context));
     const result = await generateObject({
       model: agent.model,
+      system: input.system ?? agent.description,
       prompt: `
 <goal>
 ${input.goal}
@@ -83,6 +85,16 @@ Examples:
 {
   "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
@@ -109,6 +121,10 @@ Examples:
     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 ?? {},
@@ -116,10 +132,7 @@ Examples:
 
     paths = getShortestPaths(input.machine, {
       fromState: resolvedState,
-      toState: (state) => {
-        const v = validateContext(state.context);
-        return v;
-      },
+      toState: stateFilter,
     });
   }
 
@@ -149,7 +162,7 @@ Examples:
   const nextStep = leastWeightPath?.steps[0];
 
   return {
-    planner: 'shortestPath',
+    strategy: 'shortestPath',
     episodeId: agent.episodeId,
     goal: input.goal,
     goalState: paths[0]?.state,

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

@@ -1,30 +1,27 @@
-import { CoreMessage, type CoreTool, generateText, tool } from 'ai';
+import { CoreMessage, generateText } from 'ai';
 import {
-  AgentPlan,
-  AgentPlanInput,
-  ObservedState,
+  AgentDecision,
+  AgentDecideInput,
   PromptTemplate,
-  TransitionData,
   AnyAgent,
 } from '../types';
-import { getAllTransitions, randomId } from '../utils';
-import { AnyStateMachine, getNextSnapshot } from 'xstate';
-import { defaultTextTemplate } from '../templates/defaultText';
+import { convertToXml, randomId } from '../utils';
+import { getNextSnapshot } from 'xstate';
 import { getMessages } from '../text';
 import { getToolMap } from '../decide';
 
-const simplePlannerPromptTemplate: PromptTemplate<any> = (data) => {
+const simpleStrategyPromptTemplate: PromptTemplate<any> = (data) => {
   return `
-${defaultTextTemplate(data)}
+${convertToXml(data)}
 
 Make at most one tool call to achieve the above goal. If the goal cannot be achieved with any tool calls, do not make any tool call.
   `.trim();
 };
 
-export async function simplePlanner<T extends AnyAgent>(
+export async function simpleStrategy<T extends AnyAgent>(
   agent: T,
-  input: AgentPlanInput<any>
-): Promise<AgentPlan<any> | undefined> {
+  input: AgentDecideInput<any>
+): Promise<AgentDecision<any> | undefined> {
   const toolMap = getToolMap(agent, input);
 
   if (!toolMap) {
@@ -34,8 +31,8 @@ export async function simplePlanner<T extends AnyAgent>(
 
   // Create a prompt with the given context and goal.
   // The template is used to ensure that a single tool call at most is made.
-  const prompt = simplePlannerPromptTemplate({
-    context: input.state.context,
+  const prompt = simpleStrategyPromptTemplate({
+    context: input.context,
     goal: input.goal,
   });
 
@@ -45,30 +42,33 @@ export async function simplePlanner<T extends AnyAgent>(
 
   const {
     state,
+    context,
     machine,
-    previousPlan,
+    prevDecision,
     events,
     goal,
     model: _,
     ...rest
   } = input;
 
-  const machineState = input.machine
-    ? input.machine.resolveState({
-        ...input.state,
-        context: input.state.context,
-      })
-    : undefined;
+  const machineState =
+    input.machine && input.state
+      ? input.machine.resolveState({
+          ...input.state,
+          context: input.state.context ?? {},
+        })
+      : undefined;
 
   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({
@@ -87,7 +87,7 @@ export async function simplePlanner<T extends AnyAgent>(
   }
 
   return {
-    planner: 'simple',
+    strategy: 'simple',
     goal: input.goal,
     goalState: input.state,
     nextEvent: singleResult.result,
@@ -109,3 +109,7 @@ export async function simplePlanner<T extends AnyAgent>(
     ],
   };
 }
+
+export function createSimpleStrategy<T extends AnyAgent>() {
+  return simpleStrategy;
+}