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

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

@@ -1,7 +1,7 @@
-import { createAgent } from '../';
+import { createAgent, TypesFromAgent } from '..';
 import { assign, createActor, setup } from 'xstate';
 import { z } from 'zod';
-import { experimental_createShortestPathPlanner } from './shortestPath';
+import { experimental_shortestPathStrategy } from './shortestPath';
 import { test, expect } from 'vitest';
 import { dummyResponseValues, MockLanguageModelV1 } from '../mockModel';
 
@@ -35,10 +35,7 @@ test.skip('should find shortest path to goal', async () => {
   });
 
   const counterMachine = setup({
-    types: {
-      context: agent.types.context,
-      events: agent.types.events,
-    },
+    types: {} as TypesFromAgent<typeof agent>,
   }).createMachine({
     initial: 'counting',
     context: { count: 0 },
@@ -87,7 +84,7 @@ test.skip('should find shortest path to goal', async () => {
     }),
     goal: 'Get the counter to exactly 3',
     state: counterActor.getSnapshot(),
-    planner: experimental_createShortestPathPlanner(),
+    strategy: experimental_shortestPathStrategy,
   });
 
   expect(decision?.nextEvent?.type).toBe('increment');

package/src/strategies/shortestPath.ts

@@ -0,0 +1,178 @@
+import { generateObject } from 'ai';
+import {
+  AgentDecision,
+  AgentDecideInput,
+  AgentStrategy,
+  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';
+import { randomId } from '../utils';
+
+const ajv = new Ajv();
+
+function observedStatesEqual(
+  state1: ObservedState<any>,
+  state2: ObservedState<any>
+) {
+  // 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<any>) {
+  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 experimental_shortestPathStrategy<T extends AnyAgent>(
+  agent: T,
+  input: AgentDecideInput<any>
+): Promise<AgentDecision<any> | undefined> {
+  const costFunction: CostFunction<any> =
+    input.costFunction ?? ((path) => path.weight ?? Infinity);
+  const existingDecision = agent
+    .getDecisions()
+    .find((p) => p.strategy === 'shortestPath' && p.goal === input.goal);
+
+  let paths = existingDecision?.paths;
+
+  if (existingDecision) {
+    console.log('Existing decision found');
+  }
+
+  if (!input.machine && !existingDecision) {
+    return;
+  }
+
+  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}
+</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 {
+    id: randomId(),
+    strategy: 'shortestPath',
+    episodeId: agent.episodeId,
+    goal: input.goal,
+    goalState: paths[0]?.state,
+    nextEvent: nextStep?.event,
+    paths,
+    timestamp: Date.now(),
+  };
+}

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

@@ -1,23 +1,27 @@
 import { CoreMessage, generateText } from 'ai';
-import { AgentPlan, AgentPlanInput, PromptTemplate, AnyAgent } from '../types';
-import { randomId } from '../utils';
+import {
+  AgentDecision,
+  AgentDecideInput,
+  PromptTemplate,
+  AnyAgent,
+} from '../types';
+import { convertToXml, randomId } from '../utils';
 import { getNextSnapshot } from 'xstate';
-import { defaultTextTemplate } from '../templates/defaultText';
 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) {
@@ -27,8 +31,9 @@ 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({
+    stateValue: input.state.value,
+    context: input.context ?? input.state.context,
     goal: input.goal,
   });
 
@@ -36,22 +41,15 @@ export async function simplePlanner<T extends AnyAgent>(
 
   const model = input.model ? agent.wrap(input.model) : agent.model;
 
-  const {
-    state,
-    machine,
-    previousPlan,
-    events,
-    goal,
-    model: _,
-    ...rest
-  } = input;
+  const { state, machine, 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,
@@ -81,7 +79,8 @@ export async function simplePlanner<T extends AnyAgent>(
   }
 
   return {
-    planner: 'simple',
+    id: randomId(),
+    strategy: 'simple',
     goal: input.goal,
     goalState: input.state,
     nextEvent: singleResult.result,

package/src/templates/defaultText.ts

@@ -3,6 +3,9 @@ import { wrapInXml } from '../utils';
 
 export const defaultTextTemplate: PromptTemplate<any> = (data) => {
   const preamble = [
+    data.stateValue
+      ? wrapInXml('stateValue', JSON.stringify(data.stateValue))
+      : undefined,
     data.context
       ? wrapInXml('context', JSON.stringify(data.context))
       : undefined,

package/src/text.ts

@@ -28,10 +28,10 @@ import {
  * @param options
  * @returns
  */
-export async function getMessages(
-  agent: AnyAgent,
+export async function getMessages<TAgent extends AnyAgent>(
+  agent: TAgent,
   prompt: string,
-  options: Omit<AgentGenerateTextOptions, 'prompt'>
+  options: Omit<AgentGenerateTextOptions<TAgent>, 'prompt'>
 ): Promise<CoreMessage[]> {
   let messages: CoreMessage[] = [];
   if (typeof options.messages === 'function') {
@@ -48,13 +48,13 @@ export async function getMessages(
   return messages;
 }
 
-export function fromTextStream<T extends AnyAgent>(
-  agent: T,
-  options?: AgentStreamTextOptions
+export function fromTextStream<TAgent extends AnyAgent>(
+  agent: TAgent,
+  options?: AgentStreamTextOptions<TAgent>
 ): ObservableActorLogic<
   { textDelta: string },
-  Omit<AgentStreamTextOptions, 'context'> & {
-    context?: AgentStreamTextOptions['context'];
+  Omit<AgentStreamTextOptions<TAgent>, 'context'> & {
+    context?: Record<string, any>;
   }
 > {
   const template = options?.template ?? defaultTextTemplate;
@@ -106,13 +106,13 @@ export function fromTextStream<T extends AnyAgent>(
   });
 }
 
-export function fromText<T extends AnyAgent>(
-  agent: T,
-  options?: AgentGenerateTextOptions
+export function fromText<TAgent extends AnyAgent>(
+  agent: TAgent,
+  options?: AgentGenerateTextOptions<TAgent>
 ): PromiseActorLogic<
   GenerateTextResult<Record<string, CoreTool<any, any>>>,
-  Omit<AgentGenerateTextOptions, 'context'> & {
-    context?: AgentGenerateTextOptions['context'];
+  Omit<AgentGenerateTextOptions<TAgent>, 'context'> & {
+    context?: Record<string, any>;
   }
 > {
   const resolvedOptions = {