@statelyai/agent 0.0.8 → 0.1.1

package/src/types.ts

@@ -1,61 +1,378 @@
-import OpenAI from 'openai';
-import {
-  ChatCompletionCreateParamsNonStreaming,
-  ChatCompletionCreateParamsStreaming,
-} from 'openai/resources';
 import {
+  ActorLogic,
+  ActorRefFrom,
+  AnyActorRef,
   AnyEventObject,
-  ObservableActorLogic,
+  AnyStateMachine,
+  EventFrom,
+  EventObject,
   PromiseActorLogic,
+  SnapshotFrom,
+  StateValue,
+  Subscription,
+  TransitionSnapshot,
+  Values,
 } from 'xstate';
+import {
+  CoreMessage,
+  CoreTool,
+  generateText,
+  GenerateTextResult,
+  LanguageModel,
+  streamText,
+  StreamTextResult,
+} from 'ai';
+import { ZodEventMapping } from './schemas';
+import { TypeOf } from 'zod';
+
+export type GenerateTextOptions = Parameters<typeof generateText>[0];
+
+export type StreamTextOptions = Parameters<typeof streamText>[0];
+
+export type AgentPlanInput<TEvent extends EventObject> = {
+  model: LanguageModel;
+  state: ObservedState;
+  goal: string;
+  events: ZodEventMapping;
+  machine?: AnyStateMachine;
+  /**
+   * The previous plan
+   */
+  previousPlan?: AgentPlan<TEvent>;
+};
+
+export type AgentPlan<TEvent extends EventObject> = {
+  goal: string;
+  state: ObservedState;
+  content?: string;
+  steps?: Array<{
+    event: TEvent;
+    state?: ObservedState;
+  }>;
+  nextEvent: TEvent | undefined;
+  sessionId: string;
+  timestamp: number;
+};
+
+export interface TransitionData {
+  eventType: string;
+  description?: string;
+  guard?: { type: string };
+  target?: any;
+}
 
-export interface StatelyAgentAdapter {
-  model: string;
+export type PromptTemplate<TEvents extends EventObject> = (data: {
+  goal: string;
   /**
-   * Creates actor logic that chooses an event from all of the
-   * possible next events of the parent state machine
-   * and sends it to the parent actor.
+   * The observed state
    */
-  fromEvent: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming
-  ) => PromiseActorLogic<AnyEventObject[] | undefined, TInput>;
+  state?: ObservedState;
   /**
-   * Creates actor logic that resolves with a chat completion.
+   * The context to provide.
+   * This overrides the observed state.context, if provided.
    */
-  fromChat: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming
-  ) => PromiseActorLogic<OpenAI.Chat.Completions.ChatCompletion, TInput>;
+  context?: any;
   /**
-   * Creates actor logic that emits a chat completion stream.
+   * The state machine model of the observed environment
    */
-  fromChatStream: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsStreaming
-  ) => ObservableActorLogic<
-    OpenAI.Chat.Completions.ChatCompletionChunk,
-    TInput
-  >;
+  machine?: unknown;
   /**
-   * Creates actor logic that chooses a tool from the provided
-   * tools and runs that tool.
+   * The potential next transitions that can be taken
+   * in the state machine
    */
-  fromTool: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming,
-    tools: {
-      [key: string]: Tool<any, any>;
+  transitions?: TransitionData[];
+  /**
+   * Past observations
+   */
+  observations?: AgentObservation<any>[]; // TODO
+  feedback?: AgentFeedback[];
+  messages?: AgentMessageHistory[];
+  plans?: AgentPlan<TEvents>[];
+}) => string;
+
+export type AgentPlanner<T extends Agent<any>> = (
+  agent: T['eventTypes'],
+  options: AgentPlanInput<T['eventTypes']>
+) => Promise<AgentPlan<T['eventTypes']> | undefined>;
+
+export type AgentDecideOptions = {
+  goal: string;
+  model?: LanguageModel;
+  context?: any;
+  state: ObservedState;
+  machine: AnyStateMachine;
+  execute?: (event: AnyEventObject) => Promise<void>;
+  planner?: AgentPlanner<any>;
+  events?: ZodEventMapping;
+} & Omit<
+  Parameters<typeof generateText>[0],
+  'model' | 'tools' | 'prompt' | 'messages'
+>;
+
+export interface AgentFeedback {
+  goal: string;
+  observationId: string;
+  attributes: Record<string, any>;
+  timestamp: number;
+  sessionId: string;
+}
+
+export interface AgentFeedbackInput {
+  goal: string;
+  observationId: string; // Observation ID;
+  attributes: Record<string, any>;
+  timestamp?: number;
+}
+
+export type AgentMessageHistory = CoreMessage & {
+  timestamp: number;
+  id: string;
+  /**
+   * The response ID of the message, which references
+   * which message this message is responding to, if any.
+   */
+  responseId?: string;
+  result?: GenerateTextResult<any>;
+  sessionId: string;
+};
+
+export type AgentMessageHistoryInput = CoreMessage & {
+  timestamp?: number;
+  id?: string;
+  /**
+   * The response ID of the message, which references
+   * which message this message is responding to, if any.
+   */
+  responseId?: string;
+  result?: GenerateTextResult<any>;
+};
+
+export interface AgentObservation<TActor extends AnyActorRef> {
+  id: string;
+  prevState: SnapshotFrom<TActor> | undefined;
+  event: EventFrom<TActor>;
+  state: SnapshotFrom<TActor>;
+  sessionId: string;
+  timestamp: number;
+}
+
+export interface AgentObservationInput {
+  id?: string;
+  prevState: ObservedState | undefined;
+  event: AnyEventObject;
+  state: ObservedState;
+  timestamp?: number;
+}
+
+export type AgentDecisionInput = {
+  goal: string;
+  model?: LanguageModel;
+  context?: any;
+} & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt'>;
+
+export type AgentDecisionLogic<TEvents extends EventObject> = PromiseActorLogic<
+  AgentPlan<TEvents> | undefined,
+  AgentDecisionInput | string
+>;
+
+export type AgentEmitted<TEvents extends EventObject> =
+  | {
+      type: 'feedback';
+      feedback: AgentFeedback;
+    }
+  | {
+      type: 'observation';
+      observation: AgentObservation<any>; // TODO
+    }
+  | {
+      type: 'message';
+      message: AgentMessageHistory;
+    }
+  | {
+      type: 'plan';
+      plan: AgentPlan<TEvents>;
+    };
+
+export type AgentLogic<TEvents extends EventObject> = ActorLogic<
+  TransitionSnapshot<AgentMemoryContext>,
+  | {
+      type: 'agent.feedback';
+      feedback: AgentFeedback;
     }
-  ) => PromiseActorLogic<
-    | {
-        result: any;
-        tool: string;
-        toolCall: OpenAI.Chat.Completions.ChatCompletionMessageToolCall;
-      }
-    | undefined,
-    TInput
-  >;
-}
-
-export interface Tool<TInput, TOutput> {
-  description: string;
-  inputSchema: any;
-  run: (input: TInput) => TOutput;
+  | {
+      type: 'agent.observe';
+      observation: AgentObservation<any>; // TODO
+    }
+  | {
+      type: 'agent.message';
+      message: AgentMessageHistory;
+    }
+  | {
+      type: 'agent.plan';
+      plan: AgentPlan<TEvents>;
+    },
+  any, // TODO: input
+  any,
+  AgentEmitted<TEvents>
+>;
+
+export type EventsFromZodEventMapping<TEventSchemas extends ZodEventMapping> =
+  Values<{
+    [K in keyof TEventSchemas & string]: {
+      type: K;
+    } & TypeOf<TEventSchemas[K]>;
+  }>;
+
+export type Agent<TEvents extends EventObject> = ActorRefFrom<
+  AgentLogic<TEvents>
+> & {
+  /**
+   * The general name of the agent. All agents with the same name are related and
+   * able to share experiences (observations, feedback) with each other.
+   */
+  name: string;
+  /**
+   * The unique id of the agent. This is used to partition message history.
+   */
+  id?: string;
+  description?: string;
+  events: ZodEventMapping;
+  eventTypes: TEvents;
+  model: LanguageModel;
+  defaultOptions: GenerateTextOptions;
+  memory: AgentLongTermMemory | undefined;
+  /**
+   * The adapter used to perform LLM actions such as
+   * `.generateText(…)` and `.streamText(…)`.
+   *
+   * Defaults to the Vercel AI SDK.
+   */
+  adapter: AIAdapter;
+
+  /**
+   * Resolves with an `AgentPlan` based on the information provided in the `options`, including:
+   *
+   * - The `goal` for the agent to achieve
+   * - The observed current `state`
+   * - The `logic` (e.g. a state machine) that specifies what can happen next
+   * - Additional `context`
+   */
+  decide: (
+    options: AgentDecideOptions
+  ) => Promise<AgentPlan<TEvents> | undefined>;
+
+  // Generate text
+  generateText: (
+    options: AgentGenerateTextOptions
+  ) => Promise<GenerateTextResult<Record<string, any>>>;
+
+  // Stream text
+  streamText: (
+    options: AgentStreamTextOptions
+  ) => Promise<StreamTextResult<Record<string, CoreTool<any, any>>>>;
+
+  addObservation: (observation: AgentObservationInput) => AgentObservation<any>; // TODO
+  addMessage: (history: AgentMessageHistoryInput) => AgentMessageHistory;
+  addFeedback: (feedbackItem: AgentFeedbackInput) => AgentFeedback;
+  addPlan: (plan: AgentPlan<TEvents>) => void;
+  /**
+   * Called whenever the agent (LLM assistant) receives or sends a message.
+   */
+  onMessage: (callback: (message: AgentMessageHistory) => void) => void;
+  /**
+   * Selects agent data from its context.
+   */
+  select: <T>(selector: (context: AgentMemoryContext) => T) => T;
+
+  /**
+   * Inspects state machine actor transitions and automatically observes
+   * (prevState, event, state) tuples.
+   */
+  interact: <TActor extends AnyActorRef>(
+    actorRef: TActor,
+    getInput?: (
+      observation: AgentObservation<TActor>
+    ) => AgentDecisionInput | undefined
+  ) => Subscription;
+};
+
+export type AnyAgent = Agent<any>;
+
+export type FromAgent<T> = T | ((self: AnyAgent) => T | Promise<T>);
+
+export interface CommonTextOptions {
+  prompt: FromAgent<string>;
+  model?: LanguageModel;
+  context?: Record<string, any>;
+  messages?: FromAgent<CoreMessage[]> | true;
+  template?: PromptTemplate<any>;
+}
+
+export type AgentGenerateTextOptions = Omit<
+  GenerateTextOptions,
+  'model' | 'prompt' | 'messages'
+> &
+  CommonTextOptions;
+
+export type AgentStreamTextOptions = Omit<
+  StreamTextOptions,
+  'model' | 'prompt' | 'messages'
+> &
+  CommonTextOptions;
+
+export interface ObservedState {
+  /**
+   * The current state value of the state machine, e.g.
+   * `"loading"` or `"processing"` or `"ready"`
+   */
+  value: StateValue;
+  /**
+   * Additional contextual data related to the current state
+   */
+  context: Record<string, unknown>;
+}
+
+export type ObservedStateFrom<TActor extends AnyActorRef> = Pick<
+  SnapshotFrom<TActor>,
+  'value' | 'context'
+>;
+
+export type AgentMemoryContext = {
+  observations: AgentObservation<any>[]; // TODO
+  messages: AgentMessageHistory[];
+  plans: AgentPlan<any>[];
+  feedback: AgentFeedback[];
+};
+
+export type AgentMemory = AppendOnlyStorage<AgentMemoryContext>;
+
+export interface AppendOnlyStorage<T extends Record<string, any[]>> {
+  append<K extends keyof T>(
+    sessionId: string,
+    key: K,
+    item: T[K][0]
+  ): Promise<void>;
+  getAll<K extends keyof T>(
+    sessionId: string,
+    key: K
+  ): Promise<T[K] | undefined>;
+}
+
+export interface AgentLongTermMemory {
+  get<K extends keyof AgentMemoryContext>(
+    key: K
+  ): Promise<AgentMemoryContext[K]>;
+  append<K extends keyof AgentMemoryContext>(
+    key: K,
+    item: AgentMemoryContext[K][0]
+  ): Promise<void>;
+  set<K extends keyof AgentMemoryContext>(
+    key: K,
+    items: AgentMemoryContext[K]
+  ): Promise<void>;
+}
+
+export interface AIAdapter {
+  generateText: typeof generateText;
+  streamText: typeof streamText;
 }

package/src/utils.ts

@@ -1,81 +1,22 @@
-import { AnyMachineSnapshot, AnyStateNode, Prop, Values } from 'xstate';
-import { FromSchema } from 'json-schema-to-ts';
-import { JSONSchema7 } from 'json-schema-to-ts/lib/types/definitions';
-import zodToJsonSchema, { JsonSchema7Type } from 'zod-to-json-schema';
-import { ZodEventTypes } from './schemas';
-import { z } from 'zod';
+import { AnyMachineSnapshot, AnyStateNode } from 'xstate';
+import { TransitionData } from './types';
 
-export function getAllTransitions(state: AnyMachineSnapshot) {
+export function getAllTransitions(state: AnyMachineSnapshot): TransitionData[] {
   const nodes = state._nodes;
   const transitions = (nodes as AnyStateNode[])
     .map((node) => [...(node as AnyStateNode).transitions.values()])
-    .flat(2);
+    .flat(2)
+    .map((transition) => ({
+      ...transition,
+      guard:
+        typeof transition.guard === 'string'
+          ? { type: transition.guard }
+          : (transition.guard as any), // TODO: fix
+    }));
 
   return transitions;
 }
 
-export type EventSchemas = {
-  [key: string]: {
-    description?: string;
-    properties?: {
-      [key: string]: JsonSchema7Type;
-    };
-  };
-};
-
-export type ContextSchema = JSONSchema7 & { type: 'object' };
-
-export type ConvertToJSONSchemas<T> = {
-  [K in keyof T]: {
-    properties: { type: { const: K } } & Prop<T[K], 'properties'>;
-    type: 'object';
-    required: Array<(keyof Prop<T[K], 'properties'> & string) | 'type'>;
-    additionalProperties: false;
-  };
-} & {};
-
-export function createEventSchemas<T extends EventSchemas>(
-  eventSchemaMap: T
-): ConvertToJSONSchemas<T> {
-  const resolvedEventSchemaMap = {};
-
-  for (const [key, schema] of Object.entries(eventSchemaMap)) {
-    // @ts-ignore
-    resolvedEventSchemaMap[key] = {
-      type: 'object',
-      required: ['type'],
-      properties: {
-        type: {
-          const: key,
-        },
-        ...schema.properties,
-      },
-      additionalProperties: false,
-      ...schema,
-    } as JSONSchema7;
-  }
-
-  return resolvedEventSchemaMap as ConvertToJSONSchemas<T>;
+export function wrapInXml(tagName: string, content: string): string {
+  return `<${tagName}>${content}</${tagName}>`;
 }
-
-export function createZodEventSchemas<T extends ZodEventTypes>(
-  eventSchemaMap: T
-): {
-  [K in keyof T]: unknown;
-} {
-  const resolvedEventSchemaMap = {};
-
-  for (const [eventType, zodType] of Object.entries(eventSchemaMap)) {
-    // @ts-ignore
-    resolvedEventSchemaMap[eventType] = zodToJsonSchema(
-      zodType.extend({
-        type: z.literal(eventType),
-      })
-    );
-  }
-
-  return resolvedEventSchemaMap as ConvertToJSONSchemas<T>;
-}
-
-export type InferEventsFromSchemas<T extends ConvertToJSONSchemas<any>> =
-  FromSchema<Values<T>>;

package/tsconfig.json

@@ -52,7 +52,7 @@
     // "declaration": true,                              /* Generate .d.ts files from TypeScript and JavaScript files in your project. */
     // "declarationMap": true,                           /* Create sourcemaps for d.ts files. */
     // "emitDeclarationOnly": true,                      /* Only output d.ts files and not JavaScript files. */
-    // "sourceMap": true,                                /* Create source map files for emitted JavaScript files. */
+    "sourceMap": true /* Create source map files for emitted JavaScript files. */,
     // "inlineSourceMap": true,                          /* Include sourcemap files inside the emitted JavaScript. */
     // "outFile": "./",                                  /* Specify a file that bundles all outputs into one JavaScript file. If 'declaration' is true, also designates a file that bundles all .d.ts output. */
     // "outDir": "./",                                   /* Specify an output folder for all emitted files. */

package/examples/numberGuesser.ts

@@ -1,101 +0,0 @@
-import OpenAI from 'openai';
-import { createAgent, createOpenAIAdapter, defineEvents } from '../src';
-import { assign, setup } from 'xstate';
-import { z } from 'zod';
-const openai = new OpenAI({
-  apiKey: process.env.OPENAI_API_KEY,
-});
-
-const adapter = createOpenAIAdapter(openai, {
-  model: 'gpt-3.5-turbo-1106',
-});
-
-const guessLogic = adapter.fromEvent(
-  ({
-    previousGuesses,
-    lastResult,
-  }: {
-    previousGuesses: number[];
-    lastResult: string;
-  }) => `
-  Guess the number between 1 and 10. The previous guesses were ${
-    previousGuesses.length ? previousGuesses.join(', ') : 'not made yet'
-  } and the last result was ${lastResult}.
-`
-);
-
-const events = defineEvents({
-  guess: z.object({
-    number: z.number().min(1).max(10).describe('The number guessed'),
-  }),
-});
-
-const machine = setup({
-  types: {
-    context: {} as {
-      previousGuesses: number[];
-      answer: number;
-    },
-    input: {} as { answer: number },
-    events: events.types,
-  },
-  schemas: {
-    events: events.schemas,
-  },
-  actors: {
-    guessLogic,
-  },
-}).createMachine({
-  context: ({ input }) => ({
-    answer: input.answer,
-    previousGuesses: [],
-  }),
-  initial: 'guessing',
-  states: {
-    guessing: {
-      always: {
-        guard: ({ context }) =>
-          context.answer === context.previousGuesses.at(-1),
-        target: 'winner',
-      },
-      invoke: {
-        src: 'guessLogic',
-        input: ({ context }) => ({
-          previousGuesses: context.previousGuesses,
-          lastResult:
-            context.previousGuesses.length === 0
-              ? 'not given yet'
-              : context.previousGuesses.at(-1)! - context.answer > 0
-              ? 'too high'
-              : 'too low',
-        }),
-      },
-      on: {
-        guess: {
-          actions: assign({
-            previousGuesses: ({ context, event }) => [
-              ...context.previousGuesses,
-              event.number,
-            ],
-          }),
-          target: 'guessing',
-          reenter: true,
-        },
-      },
-    },
-    winner: {
-      type: 'final',
-    },
-  },
-});
-
-const agent = createAgent(machine, {
-  input: { answer: 4 },
-  inspect: (ev) => {
-    if (ev.type === '@xstate.event') {
-      console.log(ev.event);
-    }
-  },
-});
-
-agent.start();