Cleanup

statelyai · May 13, 2024 · 8f0a25b · 8f0a25b
1 parent e9f3ef9
commit 8f0a25b
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 245 deletions.
diff --git a/src/agent-experimental.ts b/src/agent-experimental.ts
@@ -20,7 +20,7 @@ import {
   EventFrom,
 } from 'xstate';
 import { getToolCalls } from './adapters/openai';
-import { ZodEventTypes } from './schemas';
+import { ZodEventMapping } from './schemas';
 
 // >;
 export interface AgentExperience<TState, TEvent extends AnyEventObject> {
@@ -116,7 +116,7 @@ export function createAgent2<TEnvironment extends AnyActorRef>(
   // logic: AnyActorLogic,
   // input: InputFrom<TLogic>,
   getGoals: (state: SnapshotFrom<TEnvironment>) => string | string[],
-  schemas: ZodEventTypes
+  schemas: ZodEventMapping
 ): Agent<SnapshotFrom<TEnvironment>, EventFrom<TEnvironment>> {
   const experiences: Array<AgentExperience<any, any>> = [];
 

diff --git a/src/agent.ts b/src/agent.ts
@@ -11,9 +11,9 @@ import {
   toObserver,
   Values,
 } from 'xstate';
-import { AgentHistoryItem, getAllTransitions, PromptTemplate } from './utils';
+import { getAllTransitions, PromptTemplate } from './utils';
 import { ChatCompletionCreateParamsBase } from 'openai/resources/chat/completions';
-import { ZodEventTypes, EventSchemas } from './schemas';
+import { ZodEventMapping, EventSchemas } from './schemas';
 import { createZodEventSchemas } from './utils';
 import { TypeOf, z } from 'zod';
 import {
@@ -25,59 +25,63 @@ import {
   tool,
 } from 'ai';
 
-export type AgentLogic<TEventSchemas extends ZodEventTypes> = PromiseActorLogic<
-  void,
-  | ({
+export type AgentLogic<TEventSchemas extends ZodEventMapping> =
+  PromiseActorLogic<
+    void,
+    | ({
+        goal: string;
+        model?: ChatCompletionCreateParamsBase['model'];
+        /**
+         * Context to include
+         */
+        context?: any;
+      } & Omit<
+        Parameters<typeof generateText>[0],
+        'model' | 'tools' | 'prompt'
+      >)
+    | string
+  > & {
+    eventTypes: Values<{
+      [K in keyof TEventSchemas]: {
+        type: K;
+      } & TypeOf<TEventSchemas[K]>;
+    }>;
+    eventSchemas: EventSchemas<keyof TEventSchemas & string>;
+    fromText: () => PromiseActorLogic<
+      GenerateTextResult<Record<string, CoreTool<any, any>>>,
+      AgentTextStreamLogicInput
+    >;
+    fromTextStream: () => ObservableActorLogic<
+      { textDelta: string },
+      AgentTextStreamLogicInput
+    >;
+    inspect: (inspectionEvent: InspectionEvent) => void;
+    observe: ({
+      state,
+      event,
+    }: {
+      state: ObservedState;
+      event: AnyEventObject;
+      timestamp: number;
+      eventOrigin: 'environment' | 'agent';
+    }) => void;
+    reward: ({
+      goal,
+      reward,
+      timestamp,
+    }: {
       goal: string;
-      model?: ChatCompletionCreateParamsBase['model'];
-      /**
-       * Context to include
-       */
-      context?: any;
-    } & Omit<Parameters<typeof generateText>[0], 'model' | 'tools' | 'prompt'>)
-  | string
-> & {
-  eventTypes: Values<{
-    [K in keyof TEventSchemas]: {
-      type: K;
-    } & TypeOf<TEventSchemas[K]>;
-  }>;
-  eventSchemas: EventSchemas<keyof TEventSchemas & string>;
-  fromText: () => PromiseActorLogic<
-    GenerateTextResult<Record<string, CoreTool<any, any>>>,
-    AgentTextStreamLogicInput
-  >;
-  fromTextStream: () => ObservableActorLogic<
-    { textDelta: string },
-    AgentTextStreamLogicInput
-  >;
-  inspect: (inspectionEvent: InspectionEvent) => void;
-  observe: ({
-    state,
-    event,
-  }: {
-    state: ObservedState;
-    event: AnyEventObject;
-    timestamp: number;
-    eventOrigin: 'environment' | 'agent';
-  }) => void;
-  reward: ({
-    goal,
-    reward,
-    timestamp,
-  }: {
-    goal: string;
-    reward: number;
-    timestamp: number;
-  }) => void;
-  decide: ({}: {
-    model: LanguageModel;
-    goal: string;
-    state: ObservedState;
-    events: ZodEventTypes;
-    logic: AnyStateMachine;
-  }) => void;
-};
+      reward: number;
+      timestamp: number;
+    }) => void;
+    decide: ({}: {
+      goal: string;
+      state: ObservedState;
+      events: ZodEventMapping;
+      logic: AnyStateMachine;
+      promptTemplate?: PromptTemplate;
+    }) => Promise<AnyEventObject | undefined>;
+  };
 
 export type AgentTextStreamLogicInput = Omit<
   Parameters<typeof streamText>[0],
@@ -105,7 +109,6 @@ type GenerateTextOptions = Omit<
 
 export interface AgentState {
   state: ObservedState;
-  history: Array<AgentHistoryItem>;
 }
 
 const getTransitions = (state: ObservedState, logic: AnyStateMachine) => {
@@ -116,7 +119,7 @@ const getTransitions = (state: ObservedState, logic: AnyStateMachine) => {
   const resolvedState = logic.resolveState(state);
   return getAllTransitions(resolvedState);
 };
-export function createAgent<const TEventSchemas extends ZodEventTypes>({
+export function createAgent<const TEventSchemas extends ZodEventMapping>({
   model,
   events,
   stringify = JSON.stringify,
@@ -129,26 +132,13 @@ export function createAgent<const TEventSchemas extends ZodEventTypes>({
   promptTemplate?: PromptTemplate;
 } & GenerateTextOptions): AgentLogic<TEventSchemas> {
   const eventSchemas = events ? createZodEventSchemas(events) : undefined;
-  let agentState: AgentState | undefined;
 
   const observe: AgentLogic<any>['observe'] = ({
     state,
     event,
     timestamp,
     eventOrigin: eventOrigin,
-  }) => {
-    agentState = agentState ?? {
-      state,
-      history: [],
-    };
-
-    agentState.history.push({
-      state: agentState.state,
-      event: event,
-      timestamp: timestamp,
-      eventOrigin: eventOrigin,
-    });
-  };
+  }) => {};
 
   const agentLogic: AgentLogic<TEventSchemas> = fromPromise(
     async ({ input, self }) => {
@@ -168,7 +158,7 @@ export function createAgent<const TEventSchemas extends ZodEventTypes>({
         context: contextToInclude,
       };
 
-      const event = await decideFromMachine({
+      const event = await decide({
         model,
         goal: resolvedInput.goal,
         events: events ?? {}, // TODO: events should be required
@@ -262,6 +252,12 @@ export function createAgent<const TEventSchemas extends ZodEventTypes>({
   agentLogic.fromTextStream = fromTextStream;
   agentLogic.inspect = (inspectionEvent) => {};
   agentLogic.observe = observe;
+  agentLogic.decide = (stuff) =>
+    decide({
+      promptTemplate,
+      model,
+      ...stuff,
+    });
 
   return agentLogic as AgentLogic<TEventSchemas>;
 }
@@ -271,7 +267,7 @@ export interface ObservedState {
   context: Record<string, unknown>;
 }
 
-export async function decideFromMachine({
+export async function decide({
   model,
   goal,
   events,
@@ -283,15 +279,14 @@ export async function decideFromMachine({
   model: LanguageModel;
   goal: string;
   state: ObservedState;
-  events: ZodEventTypes;
+  events: ZodEventMapping;
   sessionId?: string;
   history?: Array<{
     snapshot: any;
     event: AnyEventObject;
     reward?: number;
   }>;
   logic: AnyStateMachine;
-  // transitions: TransitionData[];
   promptTemplate: PromptTemplate;
 } & GenerateTextOptions): Promise<AnyEventObject | undefined> {
   const transitions = getTransitions(state, logic);
@@ -349,7 +344,6 @@ export async function decideFromMachine({
     context: state.context,
     logic,
     transitions,
-    plan: undefined,
   });
 
   const result = await generateText({

diff --git a/src/schemas.ts b/src/schemas.ts
@@ -1,7 +1,7 @@
 import { SomeZodObject } from 'zod';
 import { JsonSchema7Type } from 'zod-to-json-schema';
 
-export type ZodEventTypes = {
+export type ZodEventMapping = {
   // map event types to Zod types
   [eventType: string]: SomeZodObject;
 };

diff --git a/src/types.ts b/src/types.ts
@@ -1,61 +1,15 @@
-import OpenAI from 'openai';
-import {
-  ChatCompletionCreateParamsNonStreaming,
-  ChatCompletionCreateParamsStreaming,
-} from 'openai/resources';
-import {
-  AnyEventObject,
-  ObservableActorLogic,
-  PromiseActorLogic,
-} from 'xstate';
+import { AnyStateMachine } from 'xstate';
+import { ObservedState } from './agent';
+import { ZodEventMapping } from './schemas';
+import { AgentPlan } from './utils';
 
-export interface StatelyAgentAdapter {
-  model: 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.
-   */
-  fromEvent: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming
-  ) => PromiseActorLogic<AnyEventObject[] | undefined, TInput>;
-  /**
-   * Creates actor logic that resolves with a chat completion.
-   */
-  fromChat: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming
-  ) => PromiseActorLogic<OpenAI.Chat.Completions.ChatCompletion, TInput>;
-  /**
-   * Creates actor logic that emits a chat completion stream.
-   */
-  fromChatStream: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsStreaming
-  ) => ObservableActorLogic<
-    OpenAI.Chat.Completions.ChatCompletionChunk,
-    TInput
-  >;
-  /**
-   * Creates actor logic that chooses a tool from the provided
-   * tools and runs that tool.
-   */
-  fromTool: <TInput>(
-    inputFn: (input: TInput) => string | ChatCompletionCreateParamsNonStreaming,
-    tools: {
-      [key: string]: Tool<any, any>;
-    }
-  ) => PromiseActorLogic<
-    | {
-        result: any;
-        tool: string;
-        toolCall: OpenAI.Chat.Completions.ChatCompletionMessageToolCall;
-      }
-    | undefined,
-    TInput
-  >;
+export interface Planner {
+  plan: (stuff: {
+    goal: string;
+    state: ObservedState;
+    events: ZodEventMapping;
+    logic: AnyStateMachine;
+  }) => Promise<AgentPlan | undefined>;
 }
 
-export interface Tool<TInput, TOutput> {
-  description: string;
-  inputSchema: any;
-  run: (input: TInput) => TOutput;
-}
+// A planner returns a plan for how to achieve a goal.