@statelyai/agent 1.1.6 → 2.0.0-next.0

package/src/planners/simplePlanner.ts

@@ -1,4 +1,4 @@
-import { type CoreTool, tool } from 'ai';
+import { CoreMessage, type CoreTool, generateText, tool } from 'ai';
 import {
   AgentPlan,
   AgentPlanInput,
@@ -7,22 +7,11 @@ import {
   TransitionData,
   AnyAgent,
 } from '../types';
-import { getAllTransitions } from '../utils';
-import { AnyStateMachine } from 'xstate';
+import { getAllTransitions, randomId } from '../utils';
+import { AnyStateMachine, getNextSnapshot } from 'xstate';
 import { defaultTextTemplate } from '../templates/defaultText';
 import { getMessages } from '../text';
-
-function getTransitions(
-  state: ObservedState,
-  machine: AnyStateMachine
-): TransitionData[] {
-  if (!machine) {
-    return [];
-  }
-
-  const resolvedState = machine.resolveState(state);
-  return getAllTransitions(resolvedState);
-}
+import { getToolMap } from '../decide';
 
 const simplePlannerPromptTemplate: PromptTemplate<any> = (data) => {
   return `
@@ -36,64 +25,9 @@ export async function simplePlanner<T extends AnyAgent>(
   agent: T,
   input: AgentPlanInput<any>
 ): Promise<AgentPlan<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;
-    }
+  const toolMap = getToolMap(agent, input);
 
-    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) {
+  if (!toolMap) {
     // No valid transitions for the specified tools
     return undefined;
   }
@@ -107,12 +41,41 @@ export async function simplePlanner<T extends AnyAgent>(
 
   const messages = await getMessages(agent, prompt, input);
 
-  const result = await agent.generateText({
-    toolChoice: 'required',
-    ...input,
-    prompt,
+  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,
+    model,
     messages,
-    tools: toolMap,
+    tools: toolMap as any,
+    toolChoice: input.toolChoice ?? 'required',
+  });
+
+  result.responseMessages.forEach((m) => {
+    const message: CoreMessage = m;
+
+    agent.addMessage({
+      ...message,
+      id: randomId(),
+      timestamp: Date.now(),
+    });
   });
 
   const singleResult = result.toolResults[0];
@@ -124,16 +87,25 @@ export async function simplePlanner<T extends AnyAgent>(
   }
 
   return {
+    planner: 'simple',
     goal: input.goal,
-    state: input.state,
-    execute: async (state) => {
-      if (JSON.stringify(state) === JSON.stringify(input.state)) {
-        return singleResult.result;
-      }
-      return undefined;
-    },
+    goalState: input.state,
     nextEvent: singleResult.result,
-    sessionId: agent.sessionId,
+    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,
+    });
   });
 }