@statelyai/agent 0.0.8 → 0.1.0

package/src/schemas.ts

@@ -1,26 +1,15 @@
-import { Values } from 'xstate';
-import { createZodEventSchemas } from './utils';
-import { SomeZodObject, TypeOf } from 'zod';
+import { SomeZodObject } from 'zod';
+import { AnyEventObject } from 'xstate';
+import { ObservedState } from './types';
 
-export type ZodEventTypes = {
+export type ZodEventMapping = {
   // map event types to Zod types
   [eventType: string]: SomeZodObject;
 };
 
-export function defineEvents<const TEventSchemas extends ZodEventTypes>(
-  events: TEventSchemas
-): {
-  types: Values<{
-    [K in keyof TEventSchemas]: {
-      type: K;
-    } & TypeOf<TEventSchemas[K]>;
-  }>;
-  schemas: {
-    [K in keyof TEventSchemas]: unknown;
+export type ZodActionMapping = {
+  [eventType: string]: {
+    schema: SomeZodObject;
+    action: (state: ObservedState, event: AnyEventObject) => Promise<void>;
   };
-} {
-  return {
-    types: {} as any,
-    schemas: createZodEventSchemas(events),
-  };
-}
+};

package/src/strategies/chain-of-note.ts

@@ -0,0 +1,155 @@
+import { GenerateTextResult, LanguageModel } from 'ai';
+import wiki, { wikiSearchResult, wikiSummary } from 'wikipedia';
+import { assign, fromPromise, setup } from 'xstate';
+import { Agent } from '../types';
+
+const searchWiki = fromPromise(
+  async ({
+    input,
+  }: {
+    input: {
+      query: string;
+      limit?: number;
+    };
+  }) => {
+    const passages = await wiki.search(input.query, {
+      limit: input.limit ?? 5,
+    });
+    return passages;
+  }
+);
+
+const extractSummaries = fromPromise(
+  async ({
+    input,
+  }: {
+    input: {
+      searchResult: wikiSearchResult;
+    };
+  }) => {
+    const summaries = await Promise.all(
+      input.searchResult.results.map(async (result) => {
+        const summary = await wiki.summary(result.title);
+        return {
+          title: result.title,
+          summary,
+        };
+      })
+    );
+    return summaries;
+  }
+);
+
+export const chainOfNote = setup({
+  types: {
+    input: {} as {
+      model: LanguageModel;
+      agent: Agent<any>;
+      prompt: string;
+    },
+    context: {} as {
+      searchResults: wikiSearchResult | null;
+      summaries:
+        | {
+            title: any;
+            summary: wikiSummary;
+          }[]
+        | null;
+      model: LanguageModel;
+      agent: Agent<any>;
+      prompt: string;
+    },
+    output: {} as GenerateTextResult<any>,
+  },
+  actors: {
+    searchWiki,
+    extractSummaries,
+  },
+}).createMachine({
+  initial: 'searching',
+  context: (x) => ({
+    ...x.input,
+    searchResults: null,
+    summaries: null,
+  }),
+  states: {
+    searching: {
+      invoke: {
+        src: 'searchWiki',
+        input: (x) => ({
+          query: x.context.prompt,
+        }),
+        onDone: {
+          actions: assign({
+            searchResults: ({ event }) => event.output,
+          }),
+          target: 'extracting',
+        },
+      },
+    },
+    extracting: {
+      invoke: {
+        src: 'extractSummaries',
+        input: (x) => ({
+          searchResult: x.context.searchResults!,
+        }),
+        onDone: {
+          actions: assign({
+            summaries: ({ event }) => event.output,
+          }),
+          target: 'generating',
+        },
+      },
+    },
+    generating: {},
+  },
+});
+
+// export function chainOfNote() {
+//   return {
+//     generateText: async (x) => {
+//       const passages = await wiki.search(x.prompt!, {
+//         limit: 5,
+//       });
+
+//       const extracts = await Promise.all(
+//         passages.results.map(async (p) => {
+//           const summary = await wiki.summary(p.title);
+//           return summary.extract;
+//         })
+//       );
+//       x.agent?.addMessage({
+//         content: x.prompt!,
+//         id: Date.now() + '',
+//         role: 'user',
+//         timestamp: Date.now(),
+//       });
+//       const result = await generateText({
+//         model: x.model,
+//         system: `Task Description:
+
+// 1. Read the given question and five Wikipedia passages to gather relevant information.
+
+// 2. Write reading notes summarizing the key points from these passages.
+
+// 3. Discuss the relevance of the given question and Wikipedia passages.
+
+// 4. If some passages are relevant to the given question, provide a brief answer based on the passages.
+
+// 5. If no passage is relevant, direcly provide answer without considering the passages.
+
+// Passages: \n${extracts.join('\n')}`,
+//         prompt: `${x.prompt!}`,
+//       });
+
+//       x.agent?.addMessage({
+//         content: result.text,
+//         id: Date.now() + '',
+//         role: 'user',
+//         timestamp: Date.now(),
+//       });
+
+//       return result;
+//     },
+//   } satisfies AgentStrategy;
+// }

package/src/templates/defaultText.ts

@@ -0,0 +1,18 @@
+import { PromptTemplate } from '../types';
+import { wrapInXml } from '../utils';
+
+export const defaultTextTemplate: PromptTemplate<any> = (data) => {
+  const preamble = [
+    data.context
+      ? wrapInXml('context', JSON.stringify(data.context))
+      : undefined,
+  ]
+    .filter(Boolean)
+    .join('\n');
+
+  return `
+${preamble}
+
+${data.goal}
+  `.trim();
+};

package/src/templates/defaultToolCall.ts

@@ -0,0 +1,10 @@
+import { PromptTemplate } from '../types';
+import { defaultTextTemplate } from './defaultText';
+
+export const defaultToolCallTemplate: PromptTemplate<any> = (data) => {
+  return `
+${defaultTextTemplate(data)}
+
+Only make a single tool call to achieve the above goal.
+  `.trim();
+};

package/src/text.ts

@@ -0,0 +1,232 @@
+import type {
+  CoreMessage,
+  CoreTool,
+  GenerateTextResult,
+  StreamTextResult,
+} from 'ai';
+import {
+  Agent,
+  AgentGenerateTextOptions,
+  AgentStreamTextOptions,
+} from './types';
+import { randomUUID } from 'crypto';
+import { defaultTextTemplate } from './templates/defaultText';
+import {
+  AnyMachineSnapshot,
+  ObservableActorLogic,
+  Observer,
+  PromiseActorLogic,
+  fromObservable,
+  fromPromise,
+  toObserver,
+} from 'xstate';
+import { vercelAdapter } from './adapters/vercel';
+
+/**
+ * Gets an array of messages from the given prompt, based on the agent and options.
+ *
+ * @param agent
+ * @param prompt
+ * @param options
+ * @returns
+ */
+async function getMessages(
+  agent: Agent<any>,
+  prompt: string,
+  options: AgentStreamTextOptions
+): Promise<CoreMessage[]> {
+  let messages: CoreMessage[] = [];
+  if (options.messages === true) {
+    messages = agent.select((s) => s.messages);
+  } else if (typeof options.messages === 'function') {
+    messages = await options.messages(agent);
+  } else if (options.messages) {
+    messages = options.messages;
+  }
+
+  messages = messages.concat({
+    role: 'user',
+    content: prompt,
+  });
+
+  return messages;
+}
+
+export async function agentGenerateText<T extends Agent<any>>(
+  agent: T,
+  options: AgentGenerateTextOptions
+) {
+  const resolvedOptions = {
+    ...agent.defaultOptions,
+    ...options,
+  };
+  const template = resolvedOptions.template ?? defaultTextTemplate;
+  // TODO: check if messages was provided instead
+  const id = randomUUID();
+  const goal =
+    typeof resolvedOptions.prompt === 'string'
+      ? resolvedOptions.prompt
+      : await resolvedOptions.prompt(agent);
+
+  const promptWithContext = template({
+    goal,
+    context: resolvedOptions.context,
+  });
+
+  const messages = await getMessages(agent, promptWithContext, resolvedOptions);
+
+  agent.addMessage({
+    id,
+    role: 'user',
+    content: promptWithContext,
+    timestamp: Date.now(),
+  });
+
+  const result = await agent.adapter.generateText({
+    ...resolvedOptions,
+    prompt: undefined,
+    messages,
+  });
+
+  agent.addMessage({
+    content: result.text,
+    id,
+    role: 'assistant',
+    timestamp: Date.now(),
+    responseId: id,
+    result,
+  });
+
+  return result;
+}
+
+export async function agentStreamText(
+  agent: Agent<any>,
+  options: AgentStreamTextOptions
+): Promise<StreamTextResult<any>> {
+  const resolvedOptions = {
+    ...agent.defaultOptions,
+    ...options,
+  };
+  const template = resolvedOptions.template ?? defaultTextTemplate;
+
+  const id = randomUUID();
+  const goal =
+    typeof resolvedOptions.prompt === 'string'
+      ? resolvedOptions.prompt
+      : await resolvedOptions.prompt(agent);
+
+  const promptWithContext = template({
+    goal,
+    context: resolvedOptions.context,
+  });
+
+  const messages = await getMessages(agent, promptWithContext, resolvedOptions);
+
+  agent.addMessage({
+    role: 'user',
+    content: promptWithContext,
+    id,
+    timestamp: Date.now(),
+  });
+
+  const result = await agent.adapter.streamText({
+    ...resolvedOptions,
+    prompt: undefined,
+    messages,
+    onFinish: async (res) => {
+      agent.addMessage({
+        role: 'assistant',
+        result: {
+          text: res.text,
+          finishReason: res.finishReason,
+          logprobs: undefined,
+          responseMessages: [],
+          toolCalls: [],
+          toolResults: [],
+          usage: res.usage,
+          warnings: res.warnings,
+          rawResponse: res.rawResponse,
+        },
+        content: res.text,
+        id: randomUUID(),
+        timestamp: Date.now(),
+        responseId: id,
+      });
+    },
+  });
+
+  return result;
+}
+
+export function fromTextStream<T extends Agent<any>>(
+  agent: T,
+  defaultOptions?: AgentStreamTextOptions
+): ObservableActorLogic<
+  { textDelta: string },
+  Omit<AgentStreamTextOptions, 'context'> & {
+    context?: AgentStreamTextOptions['context'] | boolean;
+  }
+> {
+  return fromObservable(({ input, self }) => {
+    const context =
+      input.context === true
+        ? (self._parent?.getSnapshot() as AnyMachineSnapshot).context
+        : input.context;
+
+    const observers = new Set<Observer<{ textDelta: string }>>();
+
+    // TODO: check if messages was provided instead
+
+    (async () => {
+      const result = await agentStreamText(agent, {
+        ...defaultOptions,
+        ...input,
+        context,
+      });
+
+      for await (const part of result.fullStream) {
+        if (part.type === 'text-delta') {
+          observers.forEach((observer) => {
+            observer.next?.(part);
+          });
+        }
+      }
+    })();
+
+    return {
+      subscribe: (...args: any[]) => {
+        const observer = toObserver(...args);
+        observers.add(observer);
+
+        return {
+          unsubscribe: () => {
+            observers.delete(observer);
+          },
+        };
+      },
+    };
+  });
+}
+
+export function fromText<T extends Agent<any>>(
+  agent: T,
+  defaultOptions?: AgentGenerateTextOptions
+): PromiseActorLogic<
+  GenerateTextResult<Record<string, CoreTool<any, any>>>,
+  Omit<AgentGenerateTextOptions, 'context'> & {
+    context?: AgentGenerateTextOptions['context'] | boolean;
+  }
+> {
+  return fromPromise(async ({ input, self }) => {
+    const context =
+      input.context === true
+        ? (self._parent?.getSnapshot() as AnyMachineSnapshot).context
+        : input.context;
+    return await agentGenerateText(agent, {
+      ...input,
+      ...defaultOptions,
+      context,
+    });
+  });
+}