@statelyai/agent 0.0.8 → 0.1.0

package/src/agent.ts

@@ -1,8 +1,262 @@
-import { AnyStateMachine, createActor } from 'xstate';
+import {
+  AnyEventObject,
+  AnyStateMachine,
+  createActor,
+  EventObject,
+  fromTransition,
+  Observer,
+  toObserver,
+} from 'xstate';
+import { ZodEventMapping } from './schemas';
+import {
+  Agent,
+  AgentLogic,
+  AgentMessageHistory,
+  AgentPlanner,
+  EventsFromZodEventMapping,
+  GenerateTextOptions,
+  AgentLongTermMemory,
+  AIAdapter,
+  ObservedState,
+  AgentObservationInput,
+  AgentMemoryContext,
+} from './types';
+import { simplePlanner } from './planners/simplePlanner';
+import { randomUUID } from 'crypto';
+import { agentGenerateText, agentStreamText } from './text';
+import { agentDecide } from './decision';
+import { vercelAdapter } from './adapters/vercel';
 
-export function createAgent<T extends AnyStateMachine>(
-  ...args: Parameters<typeof createActor<T>>
-) {
-  const [machine, options] = args;
-  return createActor(machine, options);
+export const agentLogic: AgentLogic<AnyEventObject> = fromTransition(
+  (state, event, { emit }) => {
+    switch (event.type) {
+      case 'agent.feedback': {
+        state.feedback.push(event.feedback);
+        emit({
+          type: 'feedback',
+          // @ts-ignore TODO: fix types in XState
+          feedback: event.feedback,
+        });
+        break;
+      }
+      case 'agent.observe': {
+        state.observations.push(event.observation);
+        emit({
+          type: 'observation',
+          // @ts-ignore TODO: fix types in XState
+          observation: event.observation,
+        });
+        break;
+      }
+      case 'agent.message': {
+        state.messages.push(event.message);
+        emit({
+          type: 'message',
+          // @ts-ignore TODO: fix types in XState
+          message: event.message,
+        });
+        break;
+      }
+      case 'agent.plan': {
+        state.plans.push(event.plan);
+        emit({
+          type: 'plan',
+          // @ts-ignore TODO: fix types in XState
+          plan: event.plan,
+        });
+        break;
+      }
+      default:
+        break;
+    }
+    return state;
+  },
+  {
+    feedback: [],
+    messages: [],
+    observations: [],
+    plans: [],
+  } as AgentMemoryContext
+);
+
+export function createAgent<
+  const TEventSchemas extends ZodEventMapping,
+  TEvents extends EventObject = EventsFromZodEventMapping<TEventSchemas>
+>({
+  name,
+  description,
+  model,
+  events,
+  planner = simplePlanner as AgentPlanner<Agent<TEvents>>,
+  stringify = JSON.stringify,
+  getMemory,
+  logic = agentLogic as AgentLogic<TEvents>,
+  adapter = vercelAdapter,
+  ...generateTextOptions
+}: {
+  /**
+   * The name of the agent
+   */
+  name: string;
+  /**
+   * A description of the role of the agent
+   */
+  description?: string;
+  /**
+   * Events that the agent can cause (send) in an environment
+   * that the agent knows about.
+   */
+  events: TEventSchemas;
+  planner?: AgentPlanner<Agent<TEvents>>;
+  stringify?: typeof JSON.stringify;
+  /**
+   * A function that retrieves the agent's long term memory
+   */
+  getMemory?: (agent: Agent<any>) => AgentLongTermMemory;
+  /**
+   * Agent logic
+   */
+  logic?: AgentLogic<TEvents>;
+  adapter?: AIAdapter;
+} & GenerateTextOptions): Agent<TEvents> {
+  const messageHistoryListeners: Observer<AgentMessageHistory>[] = [];
+
+  const agent = createActor(logic) as unknown as Agent<TEvents>;
+  agent.events = events;
+  agent.model = model;
+  agent.name = name;
+  agent.description = description;
+  agent.adapter = adapter;
+  agent.defaultOptions = { ...generateTextOptions, model };
+  agent.select = (selector) => {
+    return selector(agent.getSnapshot().context);
+  };
+  agent.memory = getMemory ? getMemory(agent) : undefined;
+
+  agent.onMessage = (callback) => {
+    messageHistoryListeners.push(toObserver(callback));
+  };
+
+  agent.decide = (opts) => {
+    return agentDecide(agent, opts);
+  };
+
+  agent.addMessage = (messageInput) => {
+    const message = {
+      ...messageInput,
+      id: messageInput.id ?? randomUUID(),
+      timestamp: messageInput.timestamp ?? Date.now(),
+      sessionId: agent.sessionId,
+    };
+    agent.send({
+      type: 'agent.message',
+      message,
+    });
+
+    return message;
+  };
+
+  agent.generateText = (opts) => agentGenerateText(agent, opts);
+
+  agent.streamText = (opts) => agentStreamText(agent, opts);
+
+  agent.addFeedback = (feedbackInput) => {
+    const feedback = {
+      ...feedbackInput,
+      timestamp: feedbackInput.timestamp ?? Date.now(),
+      sessionId: agent.sessionId,
+    };
+    agent.send({
+      type: 'agent.feedback',
+      feedback,
+    });
+    return feedback;
+  };
+
+  agent.addObservation = (observationInput) => {
+    const observation = {
+      ...observationInput,
+      id: observationInput.id ?? randomUUID(),
+      sessionId: agent.sessionId,
+      timestamp: observationInput.timestamp ?? Date.now(),
+    };
+
+    agent.send({
+      type: 'agent.observe',
+      observation,
+    });
+
+    return observation;
+  };
+
+  agent.addPlan = (plan) => {
+    agent.send({
+      type: 'agent.plan',
+      plan,
+    });
+  };
+
+  agent.interact = (actorRef, getInput) => {
+    let prevState: ObservedState | undefined = undefined;
+    let subscribed = true;
+
+    async function handleObservation(observationInput: AgentObservationInput) {
+      const observation = agent.addObservation(observationInput);
+
+      const input = getInput?.(observation);
+
+      if (input) {
+        await agentDecide(agent, {
+          machine: actorRef.src as AnyStateMachine,
+          state: observation.state,
+          execute: async (event) => {
+            actorRef.send(event);
+          },
+          ...input,
+        });
+      }
+
+      prevState = observationInput.state;
+    }
+
+    // Inspect system, but only observe specified actor
+    actorRef.system.inspect({
+      next: async (inspEvent) => {
+        if (
+          !subscribed ||
+          inspEvent.actorRef !== actorRef ||
+          inspEvent.type !== '@xstate.snapshot'
+        ) {
+          return;
+        }
+
+        const observationInput = {
+          event: inspEvent.event,
+          prevState,
+          state: inspEvent.snapshot as any,
+        };
+
+        await handleObservation(observationInput);
+      },
+    });
+
+    // If actor already started, interact with current state
+    if ((actorRef as any)._processingStatus === 1) {
+      handleObservation({
+        prevState: undefined,
+        event: { type: '' }, // TODO: unknown events?
+        state: actorRef.getSnapshot(),
+      });
+    }
+
+    return {
+      unsubscribe: () => {
+        subscribed = false;
+      }, // TODO: make this actually unsubscribe
+    };
+  };
+
+  agent.start();
+
+  return agent;
 }

package/src/decision.test.ts

@@ -0,0 +1,179 @@
+import { test, expect } from 'vitest';
+import { createAgent, fromDecision, type AIAdapter } from './';
+import { createActor, createMachine, waitFor } from 'xstate';
+import { z } from 'zod';
+import { GenerateTextResult } from 'ai';
+
+const mockToolDecision: AIAdapter['generateText'] = async (arg) => {
+  const keys = Object.keys(arg.tools!);
+
+  if (keys.length > 1) {
+    throw new Error('Expected only 1 choice');
+  }
+
+  if (keys.length === 0) {
+    return {
+      toolResults: [],
+    } as any as GenerateTextResult<any>;
+  }
+
+  return {
+    toolResults: [
+      {
+        result: {
+          type: keys[0],
+        },
+      },
+    ],
+  } as any as GenerateTextResult<any>;
+};
+
+test('fromDecision() makes a decision', async () => {
+  const agent = createAgent({
+    name: 'test',
+    model: {} as any,
+    events: {
+      doFirst: z.object({}),
+      doSecond: z.object({}),
+    },
+    adapter: {
+      generateText: async (arg) => {
+        const keys = Object.keys(arg.tools!);
+
+        if (keys.length !== 1) {
+          throw new Error('Expected only 1 choice');
+        }
+
+        return {
+          toolResults: [
+            {
+              result: {
+                type: keys[0],
+              },
+            },
+          ],
+        } as any as GenerateTextResult<any>;
+      },
+      streamText: {} as any,
+    },
+  });
+
+  const machine = createMachine({
+    initial: 'first',
+    states: {
+      first: {
+        invoke: {
+          src: fromDecision(agent),
+        },
+        on: {
+          doFirst: 'second',
+        },
+      },
+      second: {
+        invoke: {
+          src: fromDecision(agent),
+        },
+        on: {
+          doSecond: 'third',
+        },
+      },
+      third: {},
+    },
+  });
+
+  const actor = createActor(machine);
+
+  actor.start();
+
+  await waitFor(actor, (s) => s.matches('third'));
+
+  expect(actor.getSnapshot().value).toBe('third');
+});
+
+test('interacts with an actor', async () => {
+  const agent = createAgent({
+    name: 'test',
+    model: {} as any,
+    events: {
+      doFirst: z.object({}),
+      doSecond: z.object({}),
+    },
+    adapter: {
+      generateText: mockToolDecision,
+      streamText: {} as any,
+    },
+  });
+
+  const machine = createMachine({
+    initial: 'first',
+    states: {
+      first: {
+        on: {
+          doFirst: 'second',
+        },
+      },
+      second: {
+        on: {
+          doSecond: 'third',
+        },
+      },
+      third: {},
+    },
+  });
+
+  const actor = createActor(machine);
+
+  agent.interact(actor, () => ({
+    goal: 'Some goal',
+  }));
+
+  actor.start();
+
+  await waitFor(actor, (s) => s.matches('third'));
+
+  expect(actor.getSnapshot().value).toBe('third');
+});
+
+test('interacts with an actor (late interaction)', async () => {
+  const agent = createAgent({
+    name: 'test',
+    model: {} as any,
+    events: {
+      doFirst: z.object({}),
+      doSecond: z.object({}),
+    },
+    adapter: {
+      generateText: mockToolDecision,
+      streamText: {} as any,
+    },
+  });
+
+  const machine = createMachine({
+    initial: 'first',
+    states: {
+      first: {
+        on: {
+          doFirst: 'second',
+        },
+      },
+      second: {
+        on: {
+          doSecond: 'third',
+        },
+      },
+      third: {},
+    },
+  });
+
+  const actor = createActor(machine);
+
+  actor.start();
+
+  agent.interact(actor, () => ({
+    goal: 'Some goal',
+  }));
+
+  await waitFor(actor, (s) => s.matches('third'));
+
+  expect(actor.getSnapshot().value).toBe('third');
+});

package/src/decision.ts

@@ -0,0 +1,83 @@
+import { AnyMachineSnapshot, fromPromise } from 'xstate';
+import {
+  Agent,
+  AgentDecideOptions,
+  AgentDecisionLogic,
+  AgentDecisionInput,
+  AgentPlanner,
+} from './types';
+import { simplePlanner } from './planners/simplePlanner';
+
+export async function agentDecide<T extends Agent<any>>(
+  agent: T,
+  options: AgentDecideOptions
+) {
+  const resolvedOptions = {
+    ...agent.defaultOptions,
+    ...options,
+  };
+  const {
+    planner = simplePlanner as AgentPlanner<any>,
+    goal,
+    events = agent.events,
+    state,
+    machine,
+    model = agent.model,
+    ...otherPlanInput
+  } = resolvedOptions;
+
+  const plan = await planner(agent, {
+    model,
+    goal,
+    events,
+    state,
+    machine,
+    ...otherPlanInput,
+  });
+
+  if (plan?.nextEvent) {
+    agent.addPlan(plan);
+    await resolvedOptions.execute?.(plan.nextEvent);
+  }
+
+  return plan;
+}
+
+export function fromDecision(
+  agent: Agent<any>,
+  defaultInput?: AgentDecisionInput
+) {
+  return fromPromise(async ({ input, self }) => {
+    const parentRef = self._parent;
+    if (!parentRef) {
+      return;
+    }
+
+    const snapshot = parentRef.getSnapshot() as AnyMachineSnapshot;
+    const inputObject = typeof input === 'string' ? { goal: input } : input;
+    const resolvedInput = {
+      ...defaultInput,
+      ...inputObject,
+    };
+    const contextToInclude =
+      resolvedInput.context === true
+        ? // include entire context
+          parentRef.getSnapshot().context
+        : resolvedInput.context;
+    const state = {
+      value: snapshot.value,
+      context: contextToInclude,
+    };
+
+    const plan = await agentDecide(agent, {
+      machine: parentRef.src as any,
+      state,
+      execute: async (event) => {
+        parentRef.send(event);
+      },
+      ...resolvedInput,
+    });
+
+    return plan;
+  }) as AgentDecisionLogic<any>;
+}

package/src/index.ts

@@ -1,3 +1,4 @@
-export { defineEvents as defineEvents } from './schemas';
 export { createAgent } from './agent';
-export { createOpenAIAdapter } from './adapters/openai';
+export { fromText, fromTextStream, agentGenerateText } from './text';
+export { fromDecision, agentDecide } from './decision';
+export * from './types';

package/src/memory.ts

@@ -0,0 +1,25 @@
+import { AgentMemory, AgentMemoryContext } from './types';
+
+export function createAgentMemory(): AgentMemory {
+  const storage = {
+    sessions: {} as Record<string, AgentMemoryContext>,
+  };
+
+  return {
+    append: async (sessionId, key, item) => {
+      storage.sessions[sessionId] =
+        storage.sessions[sessionId] ||
+        ({
+          observations: [],
+          messages: [],
+          plans: [],
+          feedback: [],
+        } satisfies AgentMemoryContext);
+
+      storage.sessions[sessionId]![key].push(item as any);
+    },
+    getAll: async (sessionId, key) => {
+      return storage.sessions[sessionId]?.[key];
+    },
+  };
+}

package/src/planners/shortestPathPlanner.ts

@@ -0,0 +1,22 @@
+import { Agent, AgentPlan, AgentPlanInput } from '../types';
+import { getShortestPaths } from '@xstate/graph';
+
+export async function simplePlanner<T extends Agent<any>>(
+  agent: T,
+  input: AgentPlanInput<any>
+): Promise<AgentPlan<any> | undefined> {
+  // 1. Determine goal state criteria
+  // e.g. a state where the agent has won a game
+  void 0;
+
+  // 2. Determine possible events that can occur
+  void 0;
+
+  // 3. Get shortest paths from current state to
+  // a state matching the criteria, using
+  // possible events
+  void 0;
+
+  // 4. Return shortest path as a plan
+  return null as any;
+}

package/src/planners/simplePlanner.ts

@@ -0,0 +1,126 @@
+import { CoreTool, tool } from 'ai';
+import {
+  Agent,
+  AgentPlan,
+  AgentPlanInput,
+  ObservedState,
+  PromptTemplate,
+  TransitionData,
+} from '../types';
+import { getAllTransitions } from '../utils';
+import { AnyStateMachine } from 'xstate';
+import { z } from 'zod';
+import { defaultTextTemplate } from '../templates/defaultText';
+
+function getTransitions(
+  state: ObservedState,
+  machine: AnyStateMachine
+): TransitionData[] {
+  if (!machine) {
+    return [];
+  }
+
+  const resolvedState = machine.resolveState(state);
+  return getAllTransitions(resolvedState);
+}
+
+const simplePlannerPromptTemplate: PromptTemplate<any> = (data) => {
+  return `
+${defaultTextTemplate(data)}
+
+Only make a single tool call to achieve the above goal.
+  `.trim();
+};
+
+export async function simplePlanner<T extends Agent<any>>(
+  agent: T,
+  input: AgentPlanInput<any>
+): Promise<AgentPlan<any> | undefined> {
+  // Get all of the possible next transitions
+  const transitions: TransitionData[] = input.machine
+    ? getTransitions(input.state, input.machine)
+    : Object.entries(input.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);
+
+  // Mapping of each event type (e.g. "mouse.click")
+  // to a valid function name (e.g. "mouse_click")
+  const functionNameMapping: Record<string, string> = {};
+
+  const toolTransitions = transitions
+    .filter((t) => {
+      return filter(t.eventType);
+    })
+    .map((t) => {
+      const name = t.eventType.replace(/\./g, '_');
+      functionNameMapping[name] = t.eventType;
+
+      return {
+        type: 'function',
+        eventType: t.eventType,
+        description: t.description,
+        name,
+      } as const;
+    });
+
+  // Convert the transition data to a tool map that the
+  // Vercel AI SDK can use
+  const toolMap: Record<string, CoreTool<any, any>> = {};
+  for (const toolTransitionData of toolTransitions) {
+    const toolZodType = input.events?.[toolTransitionData.eventType];
+
+    toolMap[toolTransitionData.name] = tool({
+      description: toolZodType?.description ?? toolTransitionData.description,
+      parameters: toolZodType ?? z.object({}),
+      execute: async (params) => {
+        const event = {
+          type: toolTransitionData.eventType,
+          ...params,
+        };
+
+        return event;
+      },
+    });
+  }
+
+  // Create a prompt with the given context and goal.
+  // The template is used to ensure that a single tool call is made.
+  const prompt = simplePlannerPromptTemplate({
+    context: input.state.context,
+    goal: input.goal,
+  });
+
+  const result = await agent.generateText({
+    prompt,
+    tools: toolMap,
+    toolChoice: 'required',
+    ...input,
+  });
+
+  const singleResult = result.toolResults[0];
+
+  if (!singleResult) {
+    // TODO: retries?
+    console.warn('No tool call results returned');
+    return undefined;
+  }
+
+  return {
+    goal: input.goal,
+    state: input.state,
+    steps: [
+      {
+        event: singleResult.result,
+      },
+    ],
+    nextEvent: singleResult.result,
+    sessionId: agent.sessionId,
+    timestamp: Date.now(),
+  };
+}