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

package/src/agent.test.ts

@@ -70,34 +70,24 @@ test('agent.addFeedback() adds to feedback', () => {
     model: {} as any,
   });
 
-  const feedback = agent.addFeedback({
-    attributes: {
-      score: -1,
-    },
+  const obs = agent.addObservation({
+    prevState: { value: 'playing' },
+    state: { value: 'lost' },
+    event: { type: 'play', position: 3 },
     goal: 'Win the game',
-    observationId: 'obs-1',
+  });
+
+  const feedback = agent.addFeedback({
+    score: 0,
+    observationId: obs.id,
   });
 
   expect(feedback.episodeId).toEqual(agent.episodeId);
 
   expect(agent.getFeedback()).toContainEqual(
     expect.objectContaining({
-      attributes: {
-        score: -1,
-      },
-      goal: 'Win the game',
-      observationId: 'obs-1',
-      episodeId: expect.any(String),
-      timestamp: expect.any(Number),
-    })
-  );
-  expect(agent.getFeedback()).toContainEqual(
-    expect.objectContaining({
-      attributes: {
-        score: -1,
-      },
-      goal: 'Win the game',
-      observationId: 'obs-1',
+      score: 0,
+      observationId: obs.id,
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
     })
@@ -115,6 +105,7 @@ test('agent.addObservation() adds to observations', () => {
     prevState: { value: 'playing', context: {} },
     event: { type: 'play', position: 3 },
     state: { value: 'lost', context: {} },
+    goal: 'Win the game',
   });
 
   expect(observation.episodeId).toEqual(agent.episodeId);
@@ -139,6 +130,7 @@ test('agent.addObservation() adds to observations (initial state)', () => {
 
   const observation = agent.addObservation({
     state: { value: 'lost' },
+    goal: 'Win the game',
   });
 
   expect(observation.episodeId).toEqual(agent.episodeId);
@@ -176,6 +168,7 @@ test('agent.addObservation() adds to observations with machine hash', () => {
     event: { type: 'play', position: 3 },
     state: { value: 'lost', context: {} },
     machine,
+    goal: 'Win the game',
   });
 
   expect(observation.episodeId).toEqual(agent.episodeId);
@@ -203,13 +196,11 @@ test('agent.addFeedback() adds to feedback (with observation)', () => {
     state: {
       value: 'playing',
     },
+    goal: 'Win the game',
   });
 
   const feedback = agent.addFeedback({
-    attributes: {
-      score: -1,
-    },
-    goal: 'Win the game',
+    score: 0,
     observationId: observation.id,
   });
 
@@ -217,10 +208,7 @@ test('agent.addFeedback() adds to feedback (with observation)', () => {
 
   expect(agent.getFeedback()).toContainEqual(
     expect.objectContaining({
-      attributes: {
-        score: -1,
-      },
-      goal: 'Win the game',
+      score: 0,
       observationId: observation.id,
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
@@ -228,10 +216,7 @@ test('agent.addFeedback() adds to feedback (with observation)', () => {
   );
   expect(agent.getFeedback()).toContainEqual(
     expect.objectContaining({
-      attributes: {
-        score: -1,
-      },
-      goal: 'Win the game',
+      score: 0,
       observationId: observation.id,
       episodeId: expect.any(String),
       timestamp: expect.any(Number),
@@ -297,10 +282,7 @@ test('You can listen for feedback events', () => {
   agent.on('feedback', fn);
 
   agent.addFeedback({
-    attributes: {
-      score: -1,
-    },
-    goal: 'Win the game',
+    score: -1,
     observationId: 'obs-1',
   });
 
@@ -445,6 +427,7 @@ test('agent.getDecisions() returns decisions from context', () => {
     model: {} as any,
     strategy: async (agent) => {
       return {
+        id: Date.now().toString(),
         episodeId: agent.episodeId,
         strategy: 'test-strategy',
         goal: '',

package/src/agent.ts

@@ -1,7 +1,6 @@
 import {
   Actor,
   ActorRefLike,
-  AnyEventObject,
   EventObject,
   fromTransition,
   Subscription,
@@ -23,11 +22,8 @@ import {
   AgentMessageInput,
   AgentFeedbackInput,
   AgentDecision,
-  Compute,
-  AgentDecisionInput,
   AgentDecideOptions,
   AnyAgent,
-  EventsFromAgent,
   AgentInteractInput,
 } from './types';
 import { simpleStrategy } from './strategies/simple';
@@ -40,7 +36,7 @@ import {
 } from 'ai';
 import { createAgentMiddleware } from './middleware';
 
-export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
+export const agentLogic: AgentLogic<any> = fromTransition(
   (state, event, { emit }) => {
     switch (event.type) {
       case 'agent.feedback': {
@@ -92,14 +88,15 @@ export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
       messages: [],
       observations: [],
       decisions: [],
-    } as AgentMemoryContext)
+    } as AgentMemoryContext<any>)
 );
 
 export function createAgent<
   const TContextSchema extends ZodContextMapping,
   const TEventSchemas extends ZodEventMapping,
   TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
-  TContext = ContextFromZodContextMapping<TContextSchema>
+  TContext = ContextFromZodContextMapping<TContextSchema>,
+  TAgent extends AnyAgent = Agent<TContextSchema, TEventSchemas>
 >({
   id,
   description: description,
@@ -108,7 +105,7 @@ export function createAgent<
   context,
   episodeId,
   strategy = simpleStrategy,
-  logic = agentLogic as AgentLogic<TEvents>,
+  logic = agentLogic as AgentLogic<any>,
 }: {
   /**
    * The unique identifier for the agent.
@@ -142,11 +139,11 @@ export function createAgent<
    */
   getMemory?: (
     agent: Agent<TContextSchema, TEventSchemas>
-  ) => AgentLongTermMemory;
+  ) => AgentLongTermMemory<TAgent>;
   /**
    * Agent logic
    */
-  logic?: AgentLogic<TEvents>;
+  logic?: AgentLogic<TAgent>;
   model: LanguageModel;
   episodeId?: string;
 }): Agent<TContextSchema, TEventSchemas> {
@@ -167,7 +164,7 @@ export class Agent<
   const TEventSchemas extends ZodEventMapping,
   TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>,
   TContext = ContextFromZodContextMapping<TContextSchema>
-> extends Actor<AgentLogic<TEvents>> {
+> extends Actor<AgentLogic<any>> {
   /**
    * The name of the agent. All agents with the same name are related and
    * able to share experiences (observations, feedback) with each other.
@@ -186,11 +183,11 @@ export class Agent<
   //   context: Compute<TContext>;
   // };
   public model: LanguageModel;
-  public memory: AgentLongTermMemory | undefined;
+  public memory: AgentLongTermMemory<this> | undefined;
   public defaultOptions: AgentDecideOptions<AnyAgent> | undefined; // todo
 
   constructor({
-    logic = agentLogic as AgentLogic<TEvents>,
+    logic = agentLogic as AgentLogic<any>,
     id,
     name,
     description,
@@ -200,7 +197,7 @@ export class Agent<
     episodeId,
     strategy = simpleStrategy,
   }: {
-    logic: AgentLogic<TEvents>;
+    logic: AgentLogic<any>;
     id?: string;
     name?: string;
     description?: string;
@@ -262,6 +259,7 @@ export class Agent<
   public addFeedback(feedbackInput: AgentFeedbackInput) {
     const feedback = {
       ...feedbackInput,
+      comment: feedbackInput.comment ?? undefined,
       attributes: { ...feedbackInput.attributes },
       timestamp: feedbackInput.timestamp ?? Date.now(),
       episodeId: this.episodeId,
@@ -281,7 +279,7 @@ export class Agent<
   }
 
   public addObservation(
-    observationInput: AgentObservationInput
+    observationInput: AgentObservationInput<this>
   ): AgentObservation<any> {
     const { prevState, event, state } = observationInput;
     const observation = {
@@ -311,7 +309,7 @@ export class Agent<
     return this.getSnapshot().context.observations;
   }
 
-  public addDecision(decision: AgentDecision<TEvents>) {
+  public addDecision(decision: AgentDecision<this>) {
     this.send({
       type: 'agent.decision',
       decision,
@@ -383,25 +381,29 @@ export class Agent<
     const actorRefCheck = isActorRef(actorRef) && actorRef.src;
     const machine = isMachineActor(actorRef) ? actorRef.src : undefined;
 
-    let prevState: ObservedState | undefined = undefined;
+    let prevState: ObservedState<this> | undefined = undefined;
     let subscribed = true;
 
     const agent = this;
 
-    async function handleObservation(observationInput: AgentObservationInput) {
+    async function handleObservation(
+      observationInput: AgentObservationInput<any>
+    ) {
       const observation = agent.addObservation(observationInput);
 
-      const input = getInput?.(observation);
+      const interactInput = getInput?.(observation);
 
-      if (input) {
-        const res = await agentDecide(agent, {
+      if (interactInput) {
+        const decision = await agentDecide(agent, {
           machine,
           state: observation.state,
-          ...input,
+          ...interactInput,
         });
 
-        if (res?.nextEvent) {
-          actorRef.send(res.nextEvent);
+        if (decision?.nextEvent) {
+          // @ts-ignore
+          decision.nextEvent['_decision'] = decision.id;
+          actorRef.send(decision.nextEvent);
         }
       }
 
@@ -420,12 +422,21 @@ export class Agent<
               return;
             }
 
+            const decisionId = inspEvent.event['_decision'] as
+              | string
+              | undefined;
+
+            const decision = decisionId
+              ? agent.getDecisions().find((d) => d.id === decisionId)
+              : undefined;
+
             const observationInput = {
               event: inspEvent.event,
               prevState,
               state: inspEvent.snapshot as any,
               machine: (actorRef as any).src,
-            } satisfies AgentObservationInput;
+              goal: decision?.goal,
+            } satisfies AgentObservationInput<any>;
 
             await handleObservation(observationInput);
           },
@@ -439,6 +450,7 @@ export class Agent<
         event: undefined,
         state: actorRef.getSnapshot(),
         machine: (actorRef as any).src,
+        goal: undefined,
       });
     }
 
@@ -451,7 +463,7 @@ export class Agent<
   }
 
   public observe<TActor extends ActorRefLike>(actorRef: TActor): Subscription {
-    let prevState: ObservedState = actorRef.getSnapshot();
+    let prevState: ObservedState<this> = actorRef.getSnapshot();
     const actorRefCheck = isActorRef(actorRef);
 
     const sub = actorRefCheck
@@ -464,12 +476,20 @@ export class Agent<
               return;
             }
 
+            const decisionId = inspEvent.event['_decision'] as
+              | string
+              | undefined;
+            const decision = decisionId
+              ? this.getDecisions().find((d) => d.id === decisionId)
+              : undefined;
+
             const observationInput = {
               event: inspEvent.event,
               prevState,
               state: inspEvent.snapshot as any,
               machine: (actorRef as any).src,
-            } satisfies AgentObservationInput;
+              goal: decision?.goal,
+            } satisfies AgentObservationInput<this>;
 
             prevState = observationInput.state;
 
@@ -498,7 +518,7 @@ export class Agent<
    */
   public async decide(
     opts: AgentDecideOptions<this>
-  ): Promise<AgentDecision<EventsFromAgent<this>> | undefined> {
+  ): Promise<AgentDecision<this> | undefined> {
     return agentDecide(this, opts);
   }
 }

package/src/decide.ts

@@ -6,15 +6,15 @@ import {
   AgentDecision,
   AgentDecideInput,
   TransitionData,
-  EventsFromAgent,
+  EventFromAgent,
 } from './types';
 import { getTransitions } from './utils';
 import { CoreMessage, CoreTool, tool } from 'ai';
 
-export async function agentDecide<T extends AnyAgent>(
-  agent: T,
-  options: AgentDecideOptions<T>
-): Promise<AgentDecision<EventsFromAgent<T>> | undefined> {
+export async function agentDecide<TAgent extends AnyAgent>(
+  agent: TAgent,
+  options: AgentDecideOptions<TAgent>
+): Promise<AgentDecision<TAgent> | undefined> {
   const resolvedOptions = {
     ...agent.defaultOptions,
     ...options,
@@ -45,12 +45,17 @@ export async function agentDecide<T extends AnyAgent>(
 
   let decision: AgentDecision<any> | undefined;
 
+  const minimalState = {
+    value: state.value,
+    context: state.context,
+  };
+
   while (attempts++ < maxAttempts) {
     decision = await strategy(agent, {
       model,
       goal,
       events: filteredEventSchemas,
-      state,
+      state: minimalState,
       machine,
       messages: messages as CoreMessage[], // TODO: fix UIMessage thing
       ...otherDecideInput,
@@ -68,7 +73,7 @@ export async function agentDecide<T extends AnyAgent>(
 
 export function fromDecision<T extends AnyAgent>(
   agent: T,
-  defaultInput?: AgentDecideInput<EventsFromAgent<T>>
+  defaultInput?: AgentDecideInput<EventFromAgent<T>>
 ): AgentDecisionLogic<any> {
   return fromPromise(async ({ input, self }) => {
     const parentRef = self._parent;
@@ -82,15 +87,10 @@ export function fromDecision<T extends AnyAgent>(
       ...defaultInput,
       ...inputObject,
     };
-    const state = {
-      value: snapshot.value,
-      context: resolvedInput.context,
-    };
 
     const decision = await agentDecide(agent, {
       machine: (parentRef as AnyActor).logic,
       state: snapshot,
-      context: resolvedInput.context,
       execute: async (event) => {
         parentRef.send(event);
       },

package/src/strategies/chainOfThought.ts

@@ -10,11 +10,12 @@ import { simpleStrategy } from './simple';
 import { convertToXml } from '../utils';
 
 const chainOfThoughtPromptTemplate: PromptTemplate<any> = ({
+  stateValue,
   context,
   goal,
 }) => {
   return `
-${convertToXml({ context, goal })}
+${convertToXml({ stateValue, context, goal })}
 
 How would you achieve the goal? Think step-by-step.
 `.trim();
@@ -25,7 +26,8 @@ export async function chainOfThoughtStrategy<T extends AnyAgent>(
   input: AgentDecideInput<any>
 ): Promise<AgentDecision<any> | undefined> {
   const prompt = chainOfThoughtPromptTemplate({
-    context: input.state.context,
+    stateValue: input.state.value,
+    context: input.context ?? input.state.context,
     goal: input.goal,
   });
 

package/src/strategies/shortestPath.ts

@@ -13,10 +13,14 @@ 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, state2: ObservedState) {
+function observedStatesEqual(
+  state1: ObservedState<any>,
+  state2: ObservedState<any>
+) {
   // check state value && state context
   return (
     JSON.stringify(state1.value) === JSON.stringify(state2.value) &&
@@ -24,7 +28,7 @@ function observedStatesEqual(state1: ObservedState, state2: ObservedState) {
   );
 }
 
-function trimSteps(steps: AgentStep<any>[], currentState: ObservedState) {
+function trimSteps(steps: AgentStep<any>[], currentState: ObservedState<any>) {
   const index = steps.findIndex(
     (step) => step.state && observedStatesEqual(step.state, currentState)
   );
@@ -162,6 +166,7 @@ Examples:
   const nextStep = leastWeightPath?.steps[0];
 
   return {
+    id: randomId(),
     strategy: 'shortestPath',
     episodeId: agent.episodeId,
     goal: input.goal,

package/src/strategies/simple.ts

@@ -32,7 +32,8 @@ 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({
-    context: input.context,
+    stateValue: input.state.value,
+    context: input.context ?? input.state.context,
     goal: input.goal,
   });
 
@@ -40,16 +41,7 @@ export async function simpleStrategy<T extends AnyAgent>(
 
   const model = input.model ? agent.wrap(input.model) : agent.model;
 
-  const {
-    state,
-    context,
-    machine,
-    prevDecision,
-    events,
-    goal,
-    model: _,
-    ...rest
-  } = input;
+  const { state, machine, events, goal, model: _, ...rest } = input;
 
   const machineState =
     input.machine && input.state
@@ -87,6 +79,7 @@ export async function simpleStrategy<T extends AnyAgent>(
   }
 
   return {
+    id: randomId(),
     strategy: 'simple',
     goal: input.goal,
     goalState: input.state,
@@ -109,7 +102,3 @@ export async function simpleStrategy<T extends AnyAgent>(
     ],
   };
 }
-
-export function createSimpleStrategy<T extends AnyAgent>() {
-  return simpleStrategy;
-}

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 = {