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

package/src/{strategies/chainOfThought.ts → policies/chainOfThoughtPolicy.ts}

@@ -5,8 +5,8 @@ import {
   AgentDecision,
   PromptTemplate,
 } from '../types';
-import { getMessages } from '../text';
-import { simpleStrategy } from './simple';
+import { combinePromptAndMessages } from '../text';
+import { toolPolicy } from './toolPolicy';
 import { convertToXml } from '../utils';
 
 const chainOfThoughtPromptTemplate: PromptTemplate<any> = ({
@@ -14,14 +14,12 @@ const chainOfThoughtPromptTemplate: PromptTemplate<any> = ({
   context,
   goal,
 }) => {
-  return `
-${convertToXml({ stateValue, context, goal })}
+  return `${convertToXml({ stateValue, context, goal })}
 
-How would you achieve the goal? Think step-by-step.
-`.trim();
+How would you achieve the goal? Think step-by-step.`;
 };
 
-export async function chainOfThoughtStrategy<T extends AnyAgent>(
+export async function chainOfThoughtPolicy<T extends AnyAgent>(
   agent: T,
   input: AgentDecideInput<any>
 ): Promise<AgentDecision<any> | undefined> {
@@ -31,7 +29,7 @@ export async function chainOfThoughtStrategy<T extends AnyAgent>(
     goal: input.goal,
   });
 
-  const messages = await getMessages(agent, prompt, input);
+  const messages = combinePromptAndMessages(prompt, input.messages);
 
   const model = input.model ? agent.wrap(input.model) : agent.model;
 
@@ -41,7 +39,7 @@ export async function chainOfThoughtStrategy<T extends AnyAgent>(
     messages,
   });
 
-  const decision = await simpleStrategy(agent, {
+  const decision = await toolPolicy(agent, {
     ...input,
     messages: messages.concat(result.response.messages),
   });

package/src/policies/index.ts

@@ -0,0 +1,3 @@
+export * from './chainOfThoughtPolicy';
+export * from './shortestPathPolicy';
+export * from './toolPolicy';

package/src/{strategies/shortestPath.test.ts → policies/shortestPathPolicy.test.ts}

@@ -1,7 +1,7 @@
 import { createAgent, TypesFromAgent } from '..';
 import { assign, createActor, setup } from 'xstate';
 import { z } from 'zod';
-import { experimental_shortestPathStrategy } from './shortestPath';
+import { experimental_shortestPathPolicy } from './shortestPathPolicy';
 import { test, expect } from 'vitest';
 import { dummyResponseValues, MockLanguageModelV1 } from '../mockModel';
 
@@ -84,7 +84,7 @@ test.skip('should find shortest path to goal', async () => {
     }),
     goal: 'Get the counter to exactly 3',
     state: counterActor.getSnapshot(),
-    strategy: experimental_shortestPathStrategy,
+    policy: experimental_shortestPathPolicy,
   });
 
   expect(decision?.nextEvent?.type).toBe('increment');

package/src/{strategies/shortestPath.ts → policies/shortestPathPolicy.ts}

@@ -2,7 +2,6 @@ import { generateObject } from 'ai';
 import {
   AgentDecision,
   AgentDecideInput,
-  AgentStrategy,
   AgentStep,
   AnyAgent,
   CostFunction,
@@ -40,15 +39,15 @@ function trimSteps(steps: AgentStep<any>[], currentState: ObservedState<any>) {
   return steps.slice(index + 1, steps.length);
 }
 
-export async function experimental_shortestPathStrategy<T extends AnyAgent>(
+export async function experimental_shortestPathPolicy<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);
+  const existingDecision = input.decisions?.find(
+    (p) => p.policy === 'shortestPath' && p.goal === input.goal
+  );
 
   let paths = existingDecision?.paths;
 
@@ -167,11 +166,12 @@ Examples:
 
   return {
     id: randomId(),
-    strategy: 'shortestPath',
+    decisionId: input.decisionId ?? null,
+    policy: 'shortestPath',
     episodeId: agent.episodeId,
     goal: input.goal,
-    goalState: paths[0]?.state,
-    nextEvent: nextStep?.event,
+    goalState: paths[0]?.state ?? null,
+    nextEvent: nextStep?.event ?? null,
     paths,
     timestamp: Date.now(),
   };

package/src/{strategies/simple.ts → policies/toolPolicy.ts}

@@ -1,4 +1,4 @@
-import { CoreMessage, generateText } from 'ai';
+import { CoreToolResult, generateText } from 'ai';
 import {
   AgentDecision,
   AgentDecideInput,
@@ -6,11 +6,11 @@ import {
   AnyAgent,
 } from '../types';
 import { convertToXml, randomId } from '../utils';
-import { getNextSnapshot } from 'xstate';
-import { getMessages } from '../text';
+import { transition } from 'xstate';
+import { combinePromptAndMessages } from '../text';
 import { getToolMap } from '../decide';
 
-const simpleStrategyPromptTemplate: PromptTemplate<any> = (data) => {
+const toolPolicyPromptTemplate: PromptTemplate<any> = (data) => {
   return `
 ${convertToXml(data)}
 
@@ -18,10 +18,10 @@ Make at most one tool call to achieve the above goal. If the goal cannot be achi
   `.trim();
 };
 
-export async function simpleStrategy<T extends AnyAgent>(
-  agent: T,
-  input: AgentDecideInput<any>
-): Promise<AgentDecision<any> | undefined> {
+export async function toolPolicy<TAgent extends AnyAgent>(
+  agent: TAgent,
+  input: AgentDecideInput<TAgent>
+): Promise<AgentDecision<TAgent> | undefined> {
   const toolMap = getToolMap(agent, input);
 
   if (!toolMap) {
@@ -31,13 +31,13 @@ export async function simpleStrategy<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 = simpleStrategyPromptTemplate({
+  const prompt = toolPolicyPromptTemplate({
     stateValue: input.state.value,
     context: input.context ?? input.state.context,
     goal: input.goal,
   });
 
-  const messages = await getMessages(agent, prompt, input);
+  const messages = combinePromptAndMessages(prompt, input.messages);
 
   const model = input.model ? agent.wrap(input.model) : agent.model;
 
@@ -56,21 +56,19 @@ export async function simpleStrategy<T extends AnyAgent>(
     system: input.system ?? agent.description,
     model,
     messages,
-    tools: toolMap as any,
+    tools: toolMap,
     toolChoice: input.toolChoice ?? 'required',
   });
 
   result.response.messages.forEach((m) => {
-    const message: CoreMessage = m;
-
-    agent.addMessage({
-      ...message,
-      id: randomId(),
-      timestamp: Date.now(),
-    });
+    agent.addMessage(m);
   });
 
-  const singleResult = result.toolResults[0];
+  const singleResult = result.toolResults[0] as unknown as CoreToolResult<
+    any,
+    any,
+    any
+  >;
 
   if (!singleResult) {
     // TODO: retries?
@@ -78,24 +76,27 @@ export async function simpleStrategy<T extends AnyAgent>(
     return undefined;
   }
 
+  const nextEvent = singleResult.result;
+
   return {
     id: randomId(),
-    strategy: 'simple',
+    decisionId: input.decisionId ?? null,
+    policy: 'simple',
     goal: input.goal,
     goalState: input.state,
-    nextEvent: singleResult.result,
-    episodeId: agent.episodeId,
+    nextEvent,
+    episodeId: input.episodeId ?? agent.episodeId,
     timestamp: Date.now(),
     paths: [
       {
-        state: undefined,
+        state: null,
         steps: [
           {
-            event: singleResult.result,
+            event: nextEvent,
             state:
               machine && machineState
-                ? getNextSnapshot(machine, machineState, singleResult.result)
-                : undefined,
+                ? transition(machine, machineState, nextEvent)[0]
+                : null,
           },
         ],
       },

package/src/text.ts

@@ -28,32 +28,22 @@ import {
  * @param options
  * @returns
  */
-export async function getMessages<TAgent extends AnyAgent>(
-  agent: TAgent,
+export function combinePromptAndMessages(
   prompt: string,
-  options: Omit<AgentGenerateTextOptions<TAgent>, 'prompt'>
-): Promise<CoreMessage[]> {
-  let messages: CoreMessage[] = [];
-  if (typeof options.messages === 'function') {
-    messages = await options.messages(agent);
-  } else if (options.messages) {
-    messages = options.messages;
-  }
-
-  messages = messages.concat({
+  messages?: CoreMessage[]
+): CoreMessage[] {
+  return (messages ?? []).concat({
     role: 'user',
     content: prompt,
   });
-
-  return messages;
 }
 
 export function fromTextStream<TAgent extends AnyAgent>(
   agent: TAgent,
-  options?: AgentStreamTextOptions<TAgent>
+  options?: AgentStreamTextOptions
 ): ObservableActorLogic<
   { textDelta: string },
-  Omit<AgentStreamTextOptions<TAgent>, 'context'> & {
+  Omit<AgentStreamTextOptions, 'context'> & {
     context?: Record<string, any>;
   }
 > {
@@ -73,7 +63,10 @@ export function fromTextStream<TAgent extends AnyAgent>(
         goal,
         context: input.context,
       });
-      const messages = await getMessages(agent, promptWithContext, input);
+      const messages = combinePromptAndMessages(
+        promptWithContext,
+        input.messages
+      );
       const result = await streamText({
         ...options,
         ...input,
@@ -108,15 +101,14 @@ export function fromTextStream<TAgent extends AnyAgent>(
 
 export function fromText<TAgent extends AnyAgent>(
   agent: TAgent,
-  options?: AgentGenerateTextOptions<TAgent>
+  options?: AgentGenerateTextOptions
 ): PromiseActorLogic<
-  GenerateTextResult<Record<string, CoreTool<any, any>>>,
-  Omit<AgentGenerateTextOptions<TAgent>, 'context'> & {
+  GenerateTextResult<Record<string, CoreTool<any, any>>, any>,
+  Omit<AgentGenerateTextOptions, 'context'> & {
     context?: Record<string, any>;
   }
 > {
   const resolvedOptions = {
-    ...agent.defaultOptions,
     ...options,
   };
 
@@ -133,7 +125,10 @@ export function fromText<TAgent extends AnyAgent>(
       context: input.context,
     });
 
-    const messages = await getMessages(agent, promptWithContext, input);
+    const messages = combinePromptAndMessages(
+      promptWithContext,
+      input.messages
+    );
 
     const model = input.model ? agent.wrap(input.model) : agent.model;