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

package/src/agent.ts

@@ -22,13 +22,13 @@ import {
   AgentMessageInput,
   AgentFeedbackInput,
   AgentDecision,
-  AgentDecideOptions,
   AnyAgent,
   AgentInteractInput,
+  AgentDecideInput,
 } from './types';
-import { simpleStrategy } from './strategies/simple';
+import { simpleStrategy } from './strategies/simpleStrategy';
 import { agentDecide } from './decide';
-import { getMachineHash, isActorRef, isMachineActor, randomId } from './utils';
+import { isActorRef, isMachineActor, randomId } from './utils';
 import {
   experimental_wrapLanguageModel,
   LanguageModel,
@@ -184,7 +184,6 @@ export class Agent<
   // };
   public model: LanguageModel;
   public memory: AgentLongTermMemory<this> | undefined;
-  public defaultOptions: AgentDecideOptions<AnyAgent> | undefined; // todo
 
   constructor({
     logic = agentLogic as AgentLogic<any>,
@@ -262,7 +261,7 @@ export class Agent<
       comment: feedbackInput.comment ?? undefined,
       attributes: { ...feedbackInput.attributes },
       timestamp: feedbackInput.timestamp ?? Date.now(),
-      episodeId: this.episodeId,
+      episodeId: feedbackInput.episodeId ?? this.episodeId,
     } satisfies AgentFeedback;
     this.send({
       type: 'agent.feedback',
@@ -287,11 +286,12 @@ export class Agent<
       event,
       state,
       id: observationInput.id ?? randomId(),
-      episodeId: this.episodeId,
+      episodeId: observationInput.episodeId ?? this.episodeId,
       timestamp: observationInput.timestamp ?? Date.now(),
-      machineHash: observationInput.machine
-        ? getMachineHash(observationInput.machine)
-        : undefined,
+      decisionId: observationInput.decisionId,
+      // machineHash: observationInput.machine
+      //   ? getMachineHash(observationInput.machine)
+      //   : undefined,
     } satisfies AgentObservation<any>;
 
     this.send({
@@ -434,8 +434,8 @@ export class Agent<
               event: inspEvent.event,
               prevState,
               state: inspEvent.snapshot as any,
-              machine: (actorRef as any).src,
               goal: decision?.goal,
+              decisionId,
             } satisfies AgentObservationInput<any>;
 
             await handleObservation(observationInput);
@@ -446,11 +446,10 @@ export class Agent<
     // If actor already started, interact with current state
     if ((actorRef as any)._processingStatus === 1) {
       handleObservation({
+        decisionId: undefined,
         prevState: undefined,
         event: undefined,
         state: actorRef.getSnapshot(),
-        machine: (actorRef as any).src,
-        goal: undefined,
       });
     }
 
@@ -484,10 +483,10 @@ export class Agent<
               : undefined;
 
             const observationInput = {
+              decisionId,
               event: inspEvent.event,
               prevState,
               state: inspEvent.snapshot as any,
-              machine: (actorRef as any).src,
               goal: decision?.goal,
             } satisfies AgentObservationInput<this>;
 
@@ -517,8 +516,8 @@ export class Agent<
    * - Additional `context`
    */
   public async decide(
-    opts: AgentDecideOptions<this>
+    input: AgentDecideInput<this>
   ): Promise<AgentDecision<this> | undefined> {
-    return agentDecide(this, opts);
+    return agentDecide(this, input);
   }
 }

package/src/decide.test.ts

@@ -322,3 +322,25 @@ test.each([['MOVE'], ['FORFEIT']] as const)(
     expect(decision?.nextEvent?.type).toEqual(allowedEventType);
   }
 );
+
+test('agent.decide() accepts custom episodeId', async () => {
+  const model = new MockLanguageModelV1({
+    doGenerate,
+  });
+  const agent = createAgent({
+    id: 'test',
+    events: {
+      WIN: z.object({}),
+    },
+    model,
+  });
+
+  const customEpisodeId = 'custom-episode-123';
+  const decision = await agent.decide({
+    goal: 'Win the game',
+    state: { value: 'playing' },
+    episodeId: customEpisodeId,
+  });
+
+  expect(decision?.episodeId).toEqual(customEpisodeId);
+});

package/src/decide.ts

@@ -1,7 +1,6 @@
 import { AnyActor, AnyMachineSnapshot, fromPromise } from 'xstate';
 import {
   AnyAgent,
-  AgentDecideOptions,
   AgentDecisionLogic,
   AgentDecision,
   AgentDecideInput,
@@ -10,15 +9,13 @@ import {
 } from './types';
 import { getTransitions } from './utils';
 import { CoreMessage, CoreTool, tool } from 'ai';
+import { ZodEventMapping } from './schemas';
 
 export async function agentDecide<TAgent extends AnyAgent>(
   agent: TAgent,
-  options: AgentDecideOptions<TAgent>
+  options: AgentDecideInput<TAgent>
 ): Promise<AgentDecision<TAgent> | undefined> {
-  const resolvedOptions = {
-    ...agent.defaultOptions,
-    ...options,
-  };
+  const resolvedOptions = options;
   const {
     strategy = agent.strategy,
     goal,
@@ -28,6 +25,8 @@ export async function agentDecide<TAgent extends AnyAgent>(
     machine,
     model = agent.model,
     messages,
+    episodeId = agent.episodeId,
+    maxAttempts = 2,
     ...otherDecideInput
   } = resolvedOptions;
 
@@ -41,8 +40,6 @@ export async function agentDecide<TAgent extends AnyAgent>(
 
   let attempts = 0;
 
-  const maxAttempts = resolvedOptions.maxAttempts ?? 2;
-
   let decision: AgentDecision<any> | undefined;
 
   const minimalState = {
@@ -52,6 +49,7 @@ export async function agentDecide<TAgent extends AnyAgent>(
 
   while (attempts++ < maxAttempts) {
     decision = await strategy(agent, {
+      episodeId,
       model,
       goal,
       events: filteredEventSchemas,
@@ -63,7 +61,6 @@ export async function agentDecide<TAgent extends AnyAgent>(
 
     if (decision?.nextEvent) {
       agent.addDecision(decision);
-      await resolvedOptions.execute?.(decision.nextEvent);
       break;
     }
   }
@@ -88,37 +85,38 @@ export function fromDecision<T extends AnyAgent>(
       ...inputObject,
     };
 
-    const decision = await agentDecide(agent, {
+    const decision = await agentDecide<typeof agent>(agent, {
       machine: (parentRef as AnyActor).logic,
       state: snapshot,
-      execute: async (event) => {
-        parentRef.send(event);
-      },
       ...resolvedInput,
       // @ts-ignore
       messages: resolvedInput.messages,
     });
 
+    if (decision?.nextEvent) {
+      parentRef.send(decision.nextEvent);
+    }
+
     return decision;
   }) as AgentDecisionLogic<any>;
 }
 
-export function getToolMap<T extends AnyAgent>(
-  _agent: T,
+export function getToolMap<TAgent extends AnyAgent>(
+  agent: TAgent,
   input: AgentDecideInput<any>
 ): Record<string, CoreTool<any, any>> | undefined {
+  const events = input.events ?? (agent.events as ZodEventMapping);
   // Get all of the possible next transitions
   const transitions: TransitionData[] = input.machine
     ? getTransitions(input.state, input.machine)
-    : Object.entries(input.events).map(([eventType, { description }]) => ({
+    : Object.entries(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);
+  const filter = (eventType: string) => Object.keys(events).includes(eventType);
 
   // Mapping of each event type (e.g. "mouse.click")
   // to a valid function name (e.g. "mouse_click")

package/src/strategies/chainOfThought.ts

@@ -6,7 +6,7 @@ import {
   PromptTemplate,
 } from '../types';
 import { getMessages } from '../text';
-import { simpleStrategy } from './simple';
+import { simpleStrategy } from './simpleStrategy';
 import { convertToXml } from '../utils';
 
 const chainOfThoughtPromptTemplate: PromptTemplate<any> = ({

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

@@ -84,7 +84,7 @@ export async function simpleStrategy<T extends AnyAgent>(
     goal: input.goal,
     goalState: input.state,
     nextEvent: singleResult.result,
-    episodeId: agent.episodeId,
+    episodeId: input.episodeId ?? agent.episodeId,
     timestamp: Date.now(),
     paths: [
       {

package/src/text.ts

@@ -116,7 +116,6 @@ export function fromText<TAgent extends AnyAgent>(
   }
 > {
   const resolvedOptions = {
-    ...agent.defaultOptions,
     ...options,
   };
 

package/src/types.ts

@@ -33,8 +33,9 @@ export type CostFunction<TAgent extends AnyAgent> = (
 
 export type AgentDecideInput<TAgent extends AnyAgent> = Omit<
   AgentGenerateTextOptions<TAgent>,
-  'prompt' | 'tools'
+  'model' | 'prompt' | 'tools'
 > & {
+  episodeId?: string;
   /**
    * The currently observed state.
    */
@@ -52,7 +53,8 @@ export type AgentDecideInput<TAgent extends AnyAgent> = Omit<
    * The events that the agent can trigger. This is a mapping of
    * event types to Zod event schemas.
    */
-  events: ZodEventMapping;
+  events?: ZodEventMapping;
+  allowedEvents?: Array<EventFromAgent<TAgent>['type']>;
   /**
    * The state machine that represents the environment the agent
    * is interacting with.
@@ -69,6 +71,8 @@ export type AgentDecideInput<TAgent extends AnyAgent> = Omit<
    * Defaults to 2.
    */
   maxAttempts?: number;
+  strategy?: AgentStrategy<TAgent>;
+  model?: LanguageModel;
 };
 
 export type AgentStep<TAgent extends AnyAgent> = {
@@ -150,51 +154,33 @@ export type AgentStrategy<TAgent extends AnyAgent> = (
 ) => Promise<AgentDecision<TAgent> | undefined>;
 
 export type AgentInteractInput<T extends AnyAgent> = Omit<
-  AgentDecideOptions<T>,
+  AgentDecideInput<T>,
   'state'
 > & {
   state?: never;
 };
 
-export type AgentDecideOptions<TAgent extends AnyAgent> = {
-  goal: string;
-  state: ObservedState<TAgent>;
-  /**
-   * The context to provide in the prompt to the agent. This overrides the `state.context`.
-   */
-  context?: Record<string, any>;
-  machine?: AnyStateMachine;
-  model?: LanguageModel;
-  execute?: (event: AnyEventObject) => Promise<void>;
-  strategy?: AgentStrategy<TAgent>;
-  events?: ZodEventMapping;
-  allowedEvents?: Array<EventFromAgent<TAgent>['type']>;
-  /**
-   * The maximum number of times the agent will attempt to make a decision.
-   * Defaults to 2.
-   */
-  maxAttempts?: number;
-} & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt'>;
-
-export interface AgentFeedback {
-  observationId: string;
+export type AgentFeedback = {
   score: number;
   comment: string | undefined;
-  /**
-   * The message correlation that the feedback is relevant for
-   */
   attributes: Record<string, any>;
   timestamp: number;
   episodeId: string;
-}
+} & (
+  | { observationId: string; decisionId?: never }
+  | { decisionId: string; observationId?: never }
+);
 
-export interface AgentFeedbackInput {
-  observationId: string;
+export type AgentFeedbackInput = {
+  episodeId?: string;
   score: number;
   comment?: string;
   attributes?: Record<string, any>;
   timestamp?: number;
-}
+} & (
+  | { observationId: string; decisionId?: never }
+  | { decisionId: string; observationId?: never }
+);
 
 export type AgentMessage = CoreMessage & {
   timestamp: number;
@@ -336,23 +322,29 @@ export type AgentMessageInput = CoreMessage & {
 
 export interface AgentObservation<TActor extends ActorRefLike> {
   id: string;
+  decisionId?: string | undefined;
   goal?: string;
   prevState: SnapshotFrom<TActor> | undefined;
   event: EventFrom<TActor> | undefined;
   state: SnapshotFrom<TActor>;
-  machineHash: string | undefined;
+  // machineHash: string | undefined;
   episodeId: string;
   timestamp: number;
 }
 
 export interface AgentObservationInput<TAgent extends AnyAgent> {
   id?: string;
+  episodeId?: string;
+  /**
+   * The agent decision that the observation is relevant for
+   */
+  decisionId?: string | undefined;
   prevState?: ObservedState<TAgent>;
   event?: AnyEventObject;
   state: ObservedState<TAgent>;
-  machine?: AnyStateMachine;
+  // machine?: AnyStateMachine;
   timestamp?: number;
-  goal: string | undefined;
+  goal?: string | undefined;
 }
 
 export type AgentDecisionInput = {