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

package/src/strategies/chainOfThought.ts

@@ -6,15 +6,16 @@ import {
   PromptTemplate,
 } from '../types';
 import { getMessages } from '../text';
-import { simpleStrategy } from './simple';
+import { simpleStrategy } from './simpleStrategy';
 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 → simpleStrategy.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,11 +79,12 @@ export async function simpleStrategy<T extends AnyAgent>(
   }
 
   return {
+    id: randomId(),
     strategy: 'simple',
     goal: input.goal,
     goalState: input.state,
     nextEvent: singleResult.result,
-    episodeId: agent.episodeId,
+    episodeId: input.episodeId ?? agent.episodeId,
     timestamp: Date.now(),
     paths: [
       {
@@ -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,17 +106,16 @@ 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 = {
-    ...agent.defaultOptions,
     ...options,
   };
 

package/src/types.ts

@@ -27,18 +27,23 @@ export type GenerateTextOptions = Parameters<typeof generateText>[0];
 
 export type StreamTextOptions = Parameters<typeof streamText>[0];
 
-export type CostFunction<TEvent extends EventObject> = (
-  path: AgentPath<TEvent>
+export type CostFunction<TAgent extends AnyAgent> = (
+  path: AgentPath<TAgent>
 ) => number;
 
-export type AgentDecideInput<TEvent extends EventObject> = Omit<
-  AgentGenerateTextOptions,
-  'prompt' | 'tools'
+export type AgentDecideInput<TAgent extends AnyAgent> = Omit<
+  AgentGenerateTextOptions<TAgent>,
+  'model' | 'prompt' | 'tools'
 > & {
+  episodeId?: string;
   /**
    * The currently observed state.
    */
-  state: ObservedState;
+  state: ObservedState<TAgent>;
+  /**
+   * The context to provide in the prompt to the agent. This overrides the `state.context`.
+   */
+  context?: Record<string, any>;
   /**
    * The goal for the agent to accomplish.
    * The agent will make a decision based on this goal.
@@ -48,45 +53,45 @@ export type AgentDecideInput<TEvent extends EventObject> = 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.
    */
   machine?: AnyStateMachine;
-  /**
-   * The previous decision made by the agent.
-   */
-  prevDecision?: AgentDecision<TEvent>;
 
   /**
    * The total cost of the path to the goal state.
    */
-  costFunction?: CostFunction<TEvent>;
+  costFunction?: CostFunction<TAgent>;
 
   /**
    * The maximum number of attempts to make a decision.
    * Defaults to 2.
    */
   maxAttempts?: number;
+  strategy?: AgentStrategy<TAgent>;
+  model?: LanguageModel;
 };
 
-export type AgentStep<TEvent extends EventObject> = {
+export type AgentStep<TAgent extends AnyAgent> = {
   /** The event to take */
-  event: TEvent;
+  event: EventFromAgent<TAgent>;
   /** The next expected state after taking the event */
-  state: ObservedState | undefined;
+  state: ObservedState<TAgent> | undefined;
 };
 
-export type AgentPath<TEvent extends EventObject> = {
+export type AgentPath<TAgent extends AnyAgent> = {
   /** The expected ending state of the path */
-  state: ObservedState | undefined;
+  state: ObservedState<TAgent> | undefined;
   /** The steps to reach the ending state */
-  steps: Array<AgentStep<TEvent>>;
+  steps: Array<AgentStep<TAgent>>;
   weight?: number;
 };
 
-export type AgentDecision<TEvent extends EventObject> = {
+export type AgentDecision<TAgent extends AnyAgent> = {
+  id: string;
   /**
    * The strategy used to generate the decision
    */
@@ -95,17 +100,17 @@ export type AgentDecision<TEvent extends EventObject> = {
   /**
    * The ending state of the decision.
    */
-  goalState: ObservedState | undefined;
+  goalState: ObservedState<TAgent> | undefined;
   /**
    * The next event that the agent decided needs to occur to achieve the `goal`.
    *
    * This next event is chosen from the
    */
-  nextEvent: TEvent | undefined;
+  nextEvent: EventFromAgent<TAgent> | undefined;
   /**
    * The paths that the agent can take to achieve the goal.
    */
-  paths: AgentPath<TEvent>[];
+  paths: AgentPath<TAgent>[];
   episodeId: string;
   timestamp: number;
   // result: GenerateObjectResult<any>;
@@ -118,17 +123,13 @@ export interface TransitionData {
   target?: any;
 }
 
-export type PromptTemplate<TEvents extends EventObject> = (data: {
+export type PromptTemplate<TAgent extends AnyAgent> = (data: {
   goal: string;
   /**
    * The observed state
    */
-  state?: ObservedState;
-  /**
-   * The context to provide.
-   * This overrides the observed state.context, if provided.
-   */
-  context?: any;
+  stateValue?: any;
+  context?: Record<string, any>;
   /**
    * The state machine model of the observed environment
    */
@@ -144,53 +145,42 @@ export type PromptTemplate<TEvents extends EventObject> = (data: {
   observations?: AgentObservation<any>[]; // TODO
   feedback?: AgentFeedback[];
   messages?: AgentMessage[];
-  decisions?: AgentDecision<TEvents>[];
+  decisions?: AgentDecision<TAgent>[];
 }) => string;
 
-export type AgentStrategy<T extends AnyAgent> = (
-  agent: T,
-  input: AgentDecideInput<EventsFromAgent<T>>
-) => Promise<AgentDecision<EventsFromAgent<T>> | undefined>;
+export type AgentStrategy<TAgent extends AnyAgent> = (
+  agent: TAgent,
+  input: AgentDecideInput<EventFromAgent<TAgent>>
+) => Promise<AgentDecision<TAgent> | undefined>;
 
 export type AgentInteractInput<T extends AnyAgent> = Omit<
-  AgentDecideOptions<T>,
+  AgentDecideInput<T>,
   'state'
->;
-
-export type AgentDecideOptions<T extends AnyAgent> = {
-  goal: string;
-  state: ObservedState;
-  context?: Record<string, any>;
-  machine?: AnyStateMachine;
-  model?: LanguageModel;
-  execute?: (event: AnyEventObject) => Promise<void>;
-  strategy?: AgentStrategy<T>;
-  events?: ZodEventMapping;
-  allowedEvents?: Array<EventsFromAgent<T>['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'>;
+> & {
+  state?: never;
+};
 
-export interface AgentFeedback {
-  goal: string;
-  observationId: string;
-  /**
-   * The message correlation that the feedback is relevant for
-   */
+export type AgentFeedback = {
+  score: number;
+  comment: string | undefined;
   attributes: Record<string, any>;
   timestamp: number;
   episodeId: string;
-}
-
-export interface AgentFeedbackInput {
-  goal: string;
-  observationId: string;
-  attributes: Record<string, any>;
+} & (
+  | { observationId: string; decisionId?: never }
+  | { decisionId: string; observationId?: never }
+);
+
+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;
@@ -332,22 +322,29 @@ export type AgentMessageInput = CoreMessage & {
 
 export interface AgentObservation<TActor extends ActorRefLike> {
   id: string;
-  // TODO: goal
+  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 {
+export interface AgentObservationInput<TAgent extends AnyAgent> {
   id?: string;
-  prevState?: ObservedState;
+  episodeId?: string;
+  /**
+   * The agent decision that the observation is relevant for
+   */
+  decisionId?: string | undefined;
+  prevState?: ObservedState<TAgent>;
   event?: AnyEventObject;
-  state: ObservedState;
-  machine?: AnyStateMachine;
+  state: ObservedState<TAgent>;
+  // machine?: AnyStateMachine;
   timestamp?: number;
+  goal?: string | undefined;
 }
 
 export type AgentDecisionInput = {
@@ -356,12 +353,12 @@ export type AgentDecisionInput = {
   context?: Record<string, any>;
 } & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt'>;
 
-export type AgentDecisionLogic<TEvents extends EventObject> = PromiseActorLogic<
-  AgentDecision<TEvents> | undefined,
+export type AgentDecisionLogic<TAgent extends AnyAgent> = PromiseActorLogic<
+  AgentDecision<TAgent> | undefined,
   AgentDecisionInput | string
 >;
 
-export type AgentEmitted<TEvents extends EventObject> =
+export type AgentEmitted<TAgent extends AnyAgent> =
   | {
       type: 'feedback';
       feedback: AgentFeedback;
@@ -376,11 +373,11 @@ export type AgentEmitted<TEvents extends EventObject> =
     }
   | {
       type: 'decision';
-      decision: AgentDecision<TEvents>;
+      decision: AgentDecision<TAgent>;
     };
 
-export type AgentLogic<TEvents extends EventObject> = ActorLogic<
-  TransitionSnapshot<AgentMemoryContext>,
+export type AgentLogic<TAgent extends AnyAgent> = ActorLogic<
+  TransitionSnapshot<AgentMemoryContext<TAgent>>,
   | {
       type: 'agent.feedback';
       feedback: AgentFeedback;
@@ -395,19 +392,21 @@ export type AgentLogic<TEvents extends EventObject> = ActorLogic<
     }
   | {
       type: 'agent.decision';
-      decision: AgentDecision<TEvents>;
+      decision: AgentDecision<TAgent>;
     },
   any, // TODO: input
   any,
-  AgentEmitted<TEvents>
+  AgentEmitted<TAgent>
 >;
 
 export type EventsFromZodEventMapping<TEventSchemas extends ZodEventMapping> =
-  Values<{
-    [K in keyof TEventSchemas & string]: {
-      type: K;
-    } & TypeOf<TEventSchemas[K]>;
-  }>;
+  Compute<
+    Values<{
+      [K in keyof TEventSchemas & string]: {
+        type: K;
+      } & TypeOf<TEventSchemas[K]>;
+    }>
+  >;
 
 export type ContextFromZodContextMapping<
   TContextSchema extends ZodContextMapping
@@ -419,27 +418,27 @@ export type AnyAgent = Agent<any, any, any, any>;
 
 export type FromAgent<T> = T | ((agent: AnyAgent) => T | Promise<T>);
 
-export type CommonTextOptions = {
+export type CommonTextOptions<TAgent extends AnyAgent> = {
   prompt: FromAgent<string>;
   model?: LanguageModel;
-  context?: Record<string, any>;
   messages?: FromAgent<CoreMessage[]>;
   template?: PromptTemplate<any>;
+  context?: Record<string, any>;
 };
 
-export type AgentGenerateTextOptions = Omit<
+export type AgentGenerateTextOptions<TAgent extends AnyAgent> = Omit<
   GenerateTextOptions,
   'model' | 'prompt' | 'messages'
 > &
-  CommonTextOptions;
+  CommonTextOptions<TAgent>;
 
-export type AgentStreamTextOptions = Omit<
+export type AgentStreamTextOptions<TAgent extends AnyAgent> = Omit<
   StreamTextOptions,
   'model' | 'prompt' | 'messages'
 > &
-  CommonTextOptions;
+  CommonTextOptions<TAgent>;
 
-export interface ObservedState {
+export interface ObservedState<TAgent extends AnyAgent> {
   /**
    * The current state value of the state machine, e.g.
    * `"loading"` or `"processing"` or `"ready"`
@@ -448,7 +447,7 @@ export interface ObservedState {
   /**
    * Additional contextual data related to the current state
    */
-  context?: Record<string, unknown>;
+  context?: ContextFromAgent<TAgent>;
 }
 
 export type ObservedStateFrom<TActor extends ActorRefLike> = Pick<
@@ -456,24 +455,24 @@ export type ObservedStateFrom<TActor extends ActorRefLike> = Pick<
   'value' | 'context'
 >;
 
-export type AgentMemoryContext = {
-  observations: AgentObservation<any>[]; // TODO
+export type AgentMemoryContext<TAgent extends AnyAgent> = {
+  observations: AgentObservation<TAgent>[]; // TODO
   messages: AgentMessage[];
-  decisions: AgentDecision<any>[];
+  decisions: AgentDecision<TAgent>[];
   feedback: AgentFeedback[];
 };
 
-export interface AgentLongTermMemory {
-  get<K extends keyof AgentMemoryContext>(
+export interface AgentLongTermMemory<TAgent extends AnyAgent> {
+  get<K extends keyof AgentMemoryContext<TAgent>>(
     key: K
-  ): Promise<AgentMemoryContext[K]>;
-  append<K extends keyof AgentMemoryContext>(
+  ): Promise<AgentMemoryContext<TAgent>[K]>;
+  append<K extends keyof AgentMemoryContext<TAgent>>(
     key: K,
-    item: AgentMemoryContext[K][0]
+    item: AgentMemoryContext<TAgent>[K][0]
   ): Promise<void>;
-  set<K extends keyof AgentMemoryContext>(
+  set<K extends keyof AgentMemoryContext<TAgent>>(
     key: K,
-    items: AgentMemoryContext[K]
+    items: AgentMemoryContext<TAgent>[K]
   ): Promise<void>;
 }
 
@@ -481,7 +480,7 @@ export type Compute<A extends any> = { [K in keyof A]: A[K] } & unknown;
 
 export type MaybePromise<T> = T | Promise<T>;
 
-export type EventsFromAgent<T extends AnyAgent> = T extends Agent<
+export type EventFromAgent<T extends AnyAgent> = T extends Agent<
   infer _,
   infer __,
   infer TEvents,

package/src/utils.ts

@@ -101,7 +101,7 @@ export function isActorRef(
 }
 
 export function getTransitions(
-  state: ObservedState,
+  state: ObservedState<any>,
   machine: AnyStateMachine
 ): TransitionData[] {
   if (!machine) {