@statelyai/agent 1.1.6 → 2.0.0-next.1

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/simple.ts

@@ -0,0 +1,105 @@
+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';
+
+const simplePlannerPromptTemplate: PromptTemplate<any> = (data) => {
+  return `
+${defaultTextTemplate(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>(
+  agent: T,
+  input: AgentPlanInput<any>
+): Promise<AgentPlan<any> | undefined> {
+  const toolMap = getToolMap(agent, input);
+
+  if (!toolMap) {
+    // No valid transitions for the specified tools
+    return undefined;
+  }
+
+  // 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,
+    goal: input.goal,
+  });
+
+  const messages = await getMessages(agent, prompt, input);
+
+  const model = input.model ? agent.wrap(input.model) : agent.model;
+
+  const {
+    state,
+    machine,
+    previousPlan,
+    events,
+    goal,
+    model: _,
+    ...rest
+  } = input;
+
+  const machineState = input.machine
+    ? 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.response.messages.forEach((m) => {
+    const message: CoreMessage = m;
+
+    agent.addMessage({
+      ...message,
+      id: randomId(),
+      timestamp: Date.now(),
+    });
+  });
+
+  const singleResult = result.toolResults[0];
+
+  if (!singleResult) {
+    // TODO: retries?
+    console.warn('No tool call results returned');
+    return undefined;
+  }
+
+  return {
+    planner: 'simple',
+    goal: input.goal,
+    goalState: input.state,
+    nextEvent: singleResult.result,
+    episodeId: agent.episodeId,
+    timestamp: Date.now(),
+    paths: [
+      {
+        state: undefined,
+        steps: [
+          {
+            event: singleResult.result,
+            state:
+              machine && machineState
+                ? getNextSnapshot(machine, machineState, singleResult.result)
+                : undefined,
+          },
+        ],
+      },
+    ],
+  };
+}

package/src/strategies/chain-of-note.ts

@@ -67,8 +67,8 @@ export const chainOfNote = setup({
   },
 }).createMachine({
   initial: 'searching',
-  context: (x) => ({
-    ...x.input,
+  context: ({ input }) => ({
+    ...input,
     searchResults: null,
     summaries: null,
   }),
@@ -76,8 +76,8 @@ export const chainOfNote = setup({
     searching: {
       invoke: {
         src: 'searchWiki',
-        input: (x) => ({
-          query: x.context.prompt,
+        input: ({ context }) => ({
+          query: context.prompt,
         }),
         onDone: {
           actions: assign({
@@ -90,8 +90,8 @@ export const chainOfNote = setup({
     extracting: {
       invoke: {
         src: 'extractSummaries',
-        input: (x) => ({
-          searchResult: x.context.searchResults!,
+        input: ({ context }) => ({
+          searchResult: context.searchResults!,
         }),
         onDone: {
           actions: assign({
@@ -104,52 +104,3 @@ export const chainOfNote = setup({
     generating: {},
   },
 });
-
-// export function chainOfNote() {
-//   return {
-//     generateText: async (x) => {
-//       const passages = await wiki.search(x.prompt!, {
-//         limit: 5,
-//       });
-
-//       const extracts = await Promise.all(
-//         passages.results.map(async (p) => {
-//           const summary = await wiki.summary(p.title);
-//           return summary.extract;
-//         })
-//       );
-//       x.agent?.addMessage({
-//         content: x.prompt!,
-//         id: Date.now() + '',
-//         role: 'user',
-//         timestamp: Date.now(),
-//       });
-//       const result = await generateText({
-//         model: x.model,
-//         system: `Task Description:
-
-// 1. Read the given question and five Wikipedia passages to gather relevant information.
-
-// 2. Write reading notes summarizing the key points from these passages.
-
-// 3. Discuss the relevance of the given question and Wikipedia passages.
-
-// 4. If some passages are relevant to the given question, provide a brief answer based on the passages.
-
-// 5. If no passage is relevant, direcly provide answer without considering the passages.
-
-// Passages: \n${extracts.join('\n')}`,
-//         prompt: `${x.prompt!}`,
-//       });
-
-//       x.agent?.addMessage({
-//         content: result.text,
-//         id: Date.now() + '',
-//         role: 'user',
-//         timestamp: Date.now(),
-//       });
-
-//       return result;
-//     },
-//   } satisfies AgentStrategy;
-// }

package/src/text.ts

@@ -1,9 +1,13 @@
-import type { CoreMessage, CoreTool, GenerateTextResult } from 'ai';
+import {
+  generateText,
+  streamText,
+  type CoreMessage,
+  type CoreTool,
+  type GenerateTextResult,
+} from 'ai';
 import {
   AgentGenerateTextOptions,
-  AgentGenerateTextResult,
   AgentStreamTextOptions,
-  AgentStreamTextResult,
   AnyAgent,
 } from './types';
 import { defaultTextTemplate } from './templates/defaultText';
@@ -15,7 +19,6 @@ import {
   fromPromise,
   toObserver,
 } from 'xstate';
-import { randomId } from './utils';
 
 /**
  * Gets an array of messages from the given prompt, based on the agent and options.
@@ -45,158 +48,39 @@ export async function getMessages(
   return messages;
 }
 
-export async function agentGenerateText<T extends AnyAgent>(
-  agent: T,
-  options: AgentGenerateTextOptions
-): Promise<AgentGenerateTextResult> {
-  const resolvedOptions = {
-    ...agent.defaultOptions,
-    ...options,
-    correlationId: options.correlationId ?? randomId(),
-  };
-  // Generate a correlation ID if one is not provided
-  const template = resolvedOptions.template ?? defaultTextTemplate;
-  // TODO: check if messages was provided instead
-  const id = randomId();
-  const goal =
-    typeof resolvedOptions.prompt === 'string'
-      ? resolvedOptions.prompt
-      : await resolvedOptions.prompt(agent);
-
-  const promptWithContext = template({
-    goal,
-    context: resolvedOptions.context,
-  });
-
-  const messages = await getMessages(agent, promptWithContext, resolvedOptions);
-
-  agent.addMessage({
-    id,
-    role: 'user',
-    content: promptWithContext,
-    timestamp: Date.now(),
-    correlationId: resolvedOptions.correlationId,
-    parentCorrelationId: resolvedOptions.parentCorrelationId,
-  });
-
-  const result = await agent.adapter.generateText({
-    ...resolvedOptions,
-    prompt: undefined,
-    messages,
-  });
-
-  agent.addMessage({
-    content: result.text,
-    id,
-    role: 'assistant',
-    timestamp: Date.now(),
-    responseId: id,
-    result,
-    correlationId: resolvedOptions.correlationId,
-    parentCorrelationId: resolvedOptions.parentCorrelationId,
-  });
-
-  return {
-    ...result,
-    parentCorrelationId: resolvedOptions.parentCorrelationId,
-    correlationId: resolvedOptions.correlationId,
-  };
-}
-
-export async function agentStreamText(
-  agent: AnyAgent,
-  options: AgentStreamTextOptions
-): Promise<AgentStreamTextResult> {
-  const resolvedOptions = {
-    ...agent.defaultOptions,
-    ...options,
-    correlationId: options.correlationId ?? randomId(),
-  };
-  const template = resolvedOptions.template ?? defaultTextTemplate;
-
-  const id = randomId();
-  const goal =
-    typeof resolvedOptions.prompt === 'string'
-      ? resolvedOptions.prompt
-      : await resolvedOptions.prompt(agent);
-
-  const promptWithContext = template({
-    goal,
-    context: resolvedOptions.context,
-  });
-
-  const messages = await getMessages(agent, promptWithContext, resolvedOptions);
-
-  agent.addMessage({
-    role: 'user',
-    content: promptWithContext,
-    id,
-    timestamp: Date.now(),
-    correlationId: resolvedOptions.correlationId,
-    parentCorrelationId: resolvedOptions.parentCorrelationId,
-  });
-
-  const result = await agent.adapter.streamText({
-    ...resolvedOptions,
-    prompt: undefined,
-    messages,
-    onFinish: async (res) => {
-      agent.addMessage({
-        role: 'assistant',
-        result: {
-          text: res.text,
-          finishReason: res.finishReason,
-          logprobs: undefined,
-          responseMessages: [],
-          toolCalls: [],
-          toolResults: [],
-          usage: res.usage,
-          warnings: res.warnings,
-          rawResponse: res.rawResponse,
-          roundtrips: [], // TODO: how do we get this information?,
-          steps: res.steps,
-          response: res.response,
-          experimental_providerMetadata: res.experimental_providerMetadata,
-        },
-        content: res.text,
-        id: randomId(),
-        timestamp: Date.now(),
-        responseId: id,
-        correlationId: resolvedOptions.correlationId,
-        parentCorrelationId: resolvedOptions.parentCorrelationId,
-      });
-    },
-  });
-
-  return {
-    ...result,
-    textStream: result.textStream,
-    fullStream: result.fullStream,
-    parentCorrelationId: resolvedOptions.parentCorrelationId,
-    correlationId: resolvedOptions.correlationId,
-  } as unknown as AgentStreamTextResult; // TODO: fix
-}
-
 export function fromTextStream<T extends AnyAgent>(
   agent: T,
-  defaultOptions?: AgentStreamTextOptions
+  options?: AgentStreamTextOptions
 ): ObservableActorLogic<
   { textDelta: string },
   Omit<AgentStreamTextOptions, 'context'> & {
     context?: AgentStreamTextOptions['context'];
   }
 > {
+  const template = options?.template ?? defaultTextTemplate;
   return fromObservable(({ input }) => {
     const observers = new Set<Observer<{ textDelta: string }>>();
 
     // TODO: check if messages was provided instead
 
     (async () => {
-      const result = await agentStreamText(agent, {
-        ...defaultOptions,
-        ...input,
+      const model = input.model ? agent.wrap(input.model) : agent.model;
+      const goal =
+        typeof input.prompt === 'string'
+          ? input.prompt
+          : await input.prompt(agent);
+      const promptWithContext = template({
+        goal,
         context: input.context,
       });
+      const messages = await getMessages(agent, promptWithContext, input);
+      const result = await streamText({
+        ...options,
+        ...input,
+        prompt: undefined, // overwritten by messages
+        model,
+        messages,
+      });
 
       for await (const part of result.fullStream) {
         if (part.type === 'text-delta') {
@@ -224,18 +108,41 @@ export function fromTextStream<T extends AnyAgent>(
 
 export function fromText<T extends AnyAgent>(
   agent: T,
-  defaultOptions?: AgentGenerateTextOptions
+  options?: AgentGenerateTextOptions
 ): PromiseActorLogic<
   GenerateTextResult<Record<string, CoreTool<any, any>>>,
   Omit<AgentGenerateTextOptions, 'context'> & {
     context?: AgentGenerateTextOptions['context'];
   }
 > {
+  const resolvedOptions = {
+    ...agent.defaultOptions,
+    ...options,
+  };
+
+  const template = resolvedOptions.template ?? defaultTextTemplate;
+
   return fromPromise(async ({ input }) => {
-    return await agentGenerateText(agent, {
-      ...input,
-      ...defaultOptions,
+    const goal =
+      typeof input.prompt === 'string'
+        ? input.prompt
+        : await input.prompt(agent);
+
+    const promptWithContext = template({
+      goal,
       context: input.context,
     });
+
+    const messages = await getMessages(agent, promptWithContext, input);
+
+    const model = input.model ? agent.wrap(input.model) : agent.model;
+
+    return await generateText({
+      ...input,
+      ...options,
+      prompt: undefined,
+      messages,
+      model,
+    });
   });
 }