Add runner for function/tool messages (#13)

Co-authored-by: Travis Fischer <fisch0920@gmail.com>

docs/pages/guide/examples.md

@@ -37,6 +37,28 @@ npx tsx examples/caching-redis.ts
 
 [source](https://github.com/dexaai/dexter/tree/master/examples/caching-redis.ts)
 
+### AI Function
+
+This example shows how to use `createAIFunction` to handle `function` and `tool_calls` with the OpenAI chat completions API and Zod.
+
+```bash
+npx tsx examples/ai-function.ts
+```
+
+[source](https://github.com/dexaai/dexter/tree/master/examples/ai-function.ts)
+
+### AI Runner
+
+This example shows how to use `createAIRunner` to easily invoke a chain of OpenAI chat completion calls, resolving tool / function calls, retrying when necessary, and optionally validating the resulting output via Zod.
+
+Note that `createAIRunner` takes in a `functions` array of `AIFunction` objects created by `createAIFunction`, as the two utility functions are meant to used together.
+
+```bash
+npx tsx examples/ai-runner.ts
+```
+
+[source](https://github.com/dexaai/dexter/tree/master/examples/ai-runner.ts)
+
 ### Chatbot
 
 This is a more involved example of a chatbot using RAG. It indexes 100 transcript chunks from the [Huberman Lab Podcast](https://hubermanlab.com) into a [hybrid Pinecone datastore](https://docs.pinecone.io/docs/hybrid-search) using [OpenAI ada-002 embeddings](https://platform.openai.com/docs/guides/embeddings) for the dense vectors and a [HuggingFace SPLADE model](https://huggingface.co/naver/splade-cocondenser-ensembledistil) for the sparse embeddings.

examples/ai-function.ts

@@ -0,0 +1,98 @@
+import 'dotenv/config';
+import { z } from 'zod';
+import { ChatModel, createAIFunction, Msg, type Prompt } from '@dexaai/dexter';
+
+/**
+ * npx tsx examples/ai-function.ts
+ */
+async function main() {
+  const getWeather = createAIFunction(
+    {
+      name: 'get_weather',
+      description: 'Gets the weather for a given location',
+      argsSchema: z.object({
+        location: z
+          .string()
+          .describe('The city and state e.g. San Francisco, CA'),
+        unit: z
+          .enum(['c', 'f'])
+          .optional()
+          .default('f')
+          .describe('The unit of temperature to use'),
+      }),
+    },
+    // Fake weather API implementation which returns a random temperature
+    // after a short delay
+    async (args: { location: string; unit?: string }) => {
+      await new Promise((resolve) => setTimeout(resolve, 500));
+
+      return {
+        location: args.location,
+        unit: args.unit,
+        temperature: (30 + Math.random() * 70) | 0,
+      };
+    }
+  );
+
+  const chatModel = new ChatModel({
+    debug: true,
+    params: {
+      model: 'gpt-4-1106-preview',
+      temperature: 0.5,
+      max_tokens: 500,
+      tools: [
+        {
+          type: 'function',
+          function: getWeather.spec,
+        },
+      ],
+    },
+  });
+
+  const messages: Prompt.Msg[] = [
+    Msg.user('What is the weather in San Francisco?'),
+  ];
+
+  {
+    // Invoke the chat model and have it create the args for the `get_weather` function
+    const { message } = await chatModel.run({
+      messages,
+      tool_choice: {
+        type: 'function',
+        function: {
+          name: 'get_weather',
+        },
+      },
+    });
+
+    if (!Msg.isToolCall(message)) {
+      throw new Error('Expected tool call');
+    }
+    messages.push(message);
+
+    for (const toolCall of message.tool_calls) {
+      if (toolCall.function.name !== 'get_weather') {
+        throw new Error(`Invalid function name: ${toolCall.function.name}`);
+      }
+
+      const result = await getWeather(toolCall.function.arguments);
+      const toolResult = Msg.toolResult(result, toolCall.id);
+      messages.push(toolResult);
+    }
+  }
+
+  {
+    // Invoke the chat model with the result
+    const { message } = await chatModel.run({
+      messages,
+      tool_choice: 'none',
+    });
+    if (!Msg.isAssistant(message)) {
+      throw new Error('Expected assistant message');
+    }
+
+    console.log(message.content);
+  }
+}
+
+main();

examples/ai-runner.ts

@@ -0,0 +1,113 @@
+import 'dotenv/config';
+import { z } from 'zod';
+import {
+  ChatModel,
+  Msg,
+  createAIFunction,
+  createAIExtractFunction,
+  createAIRunner,
+} from '@dexaai/dexter';
+
+/** Get the weather for a given location. */
+const getWeather = createAIFunction(
+  {
+    name: 'get_weather',
+    description: 'Gets the weather for a given location',
+    argsSchema: z.object({
+      location: z
+        .string()
+        .describe('The city and state e.g. San Francisco, CA'),
+      unit: z
+        .enum(['c', 'f'])
+        .optional()
+        .default('f')
+        .describe('The unit of temperature to use'),
+    }),
+  },
+  async ({ location, unit }) => {
+    await new Promise((resolve) => setTimeout(resolve, 500));
+    const temperature = (30 + Math.random() * 70) | 0;
+    return { location, unit, temperature };
+  }
+);
+
+/** Get the capital city for a given state. */
+const getCapitalCity = createAIFunction(
+  {
+    name: 'get_capital_city',
+    description: 'Use this to get the the capital city for a given state',
+    argsSchema: z.object({
+      state: z
+        .string()
+        .length(2)
+        .describe(
+          'The state to get the capital city for, using the two letter abbreviation e.g. CA'
+        ),
+    }),
+  },
+  async ({ state }) => {
+    await new Promise((resolve) => setTimeout(resolve, 500));
+    let capitalCity = '';
+    switch (state) {
+      case 'CA':
+        capitalCity = 'Sacramento';
+        break;
+      case 'NY':
+        capitalCity = 'Albany';
+        break;
+      default:
+        capitalCity = 'Unknown';
+    }
+    return { capitalCity };
+  }
+);
+
+/** A runner that uses the weather and capital city functions. */
+const weatherCapitalRunner = createAIRunner({
+  chatModel: new ChatModel({ params: { model: 'gpt-4-1106-preview' } }),
+  functions: [getWeather, getCapitalCity],
+  systemMessage: `You use functions to answer questions about the weather and capital cities.`,
+});
+
+/** A function to extract people names from a message. */
+const extractPeopleNamesRunner = createAIExtractFunction({
+  chatModel: new ChatModel({ params: { model: 'gpt-4-1106-preview' } }),
+  systemMessage: `You use functions to extract people names from a message.`,
+  name: 'log_people_names',
+  description: `Use this to log the full names of people from a message. Don't include duplicate names.`,
+  schema: z.object({
+    names: z.array(
+      z
+        .string()
+        .describe(
+          `The name of a person from the message. Normalize the name by removing suffixes, prefixes, and fixing capitalization`
+        )
+    ),
+  }),
+});
+
+async function main() {
+  // Use OpenAI functions to extract data adhering to a Zod schema
+  const peopleNames = await extractPeopleNamesRunner(
+    `Dr. Andrew Huberman interviewed Tony Hawk, an idol of Andrew Hubermans.`
+  );
+  console.log('peopleNames', peopleNames);
+
+  // Run with a string input
+  const rString = await weatherCapitalRunner(
+    `Whats the capital of California and NY and the weather for both`
+  );
+  console.log('rString', rString);
+
+  // Run with a message input
+  const rMessage = await weatherCapitalRunner({
+    messages: [
+      Msg.user(
+        `Whats the capital of California and NY and the weather for both`
+      ),
+    ],
+  });
+  console.log('rMessage', rMessage);
+}
+
+main().catch(console.error);

package.json

@@ -66,35 +66,35 @@
   "dependencies": {
     "@fastify/deepmerge": "^1.3.0",
     "dedent": "^1.5.1",
-    "hash-object": "^5.0.0",
-    "jsonrepair": "^3.2.0",
+    "hash-object": "^5.0.1",
+    "jsonrepair": "^3.4.1",
     "ky": "^1.1.0",
     "openai-fetch": "2.0.1",
     "p-map": "^6.0.0",
-    "p-throttle": "^5.1.0",
-    "parse-json": "^7.0.0",
+    "p-throttle": "^6.0.0",
+    "parse-json": "^8.0.1",
     "pinecone-client": "^2.0.0",
-    "tiktoken": "^1.0.10",
+    "tiktoken": "^1.0.11",
     "zod": "^3.21.4",
-    "zod-to-json-schema": "^3.21.4",
-    "zod-validation-error": "^1.3.1"
+    "zod-to-json-schema": "^3.22.0",
+    "zod-validation-error": "^2.1.0"
   },
   "devDependencies": {
     "@dexaai/eslint-config": "^0.4.0",
-    "@types/node": "^20.5.9",
+    "@types/node": "^20.9.3",
     "dotenv-cli": "^7.3.0",
-    "eslint": "^8.48.0",
-    "globby": "^13.2.2",
+    "eslint": "^8.54.0",
+    "globby": "^14.0.0",
     "np": "^8.0.4",
     "npm-run-all": "^4.1.5",
     "prettier": "^2.8.8",
     "rimraf": "^5.0.5",
-    "tsx": "^3.14.0",
-    "type-fest": "4.3.1",
+    "tsx": "^4.1.4",
+    "type-fest": "4.8.1",
     "typedoc": "^0.25.3",
     "typedoc-plugin-markdown": "^4.0.0-next.25",
-    "typescript": "^5.2.2",
-    "vite": "^4.3.9",
+    "typescript": "^5.3.2",
+    "vite": "^5.0.0",
     "vitest": "^0.34.3"
   },
   "prettier": {

readme.md

@@ -19,6 +19,8 @@ _If you're a TypeScript AI engineer, check it out!_ 😊
   - [Basic](#basic)
   - [Caching](#caching)
   - [Redis Caching](#redis-caching)
+  - [AI Function](#ai-function)
+  - [AI Runner](#ai-runner)
   - [Chatbot](#chatbot)
 - [License](#license)
 
@@ -121,6 +123,28 @@ npx tsx examples/caching-redis.ts
 
 [source](./examples/caching-redis.ts)
 
+### AI Function
+
+This example shows how to use `createAIFunction` to handle `function` and `tool_calls` with the OpenAI chat completions API and Zod.
+
+```bash
+npx tsx examples/ai-function.ts
+```
+
+[source](./examples/ai-function.ts)
+
+### AI Runner
+
+This example shows how to use `createAIRunner` to easily invoke a chain of OpenAI chat completion calls, resolving tool / function calls, retrying when necessary, and optionally validating the resulting output via Zod.
+
+Note that `createAIRunner` takes in a `functions` array of `AIFunction` objects created by `createAIFunction`, as the two utility functions are meant to used together.
+
+```bash
+npx tsx examples/ai-runner.ts
+```
+
+[source](./examples/ai-runner.ts)
+
 ### Chatbot
 
 This is a more involved example of a chatbot using RAG. It indexes 100 transcript chunks from the [Huberman Lab Podcast](https://hubermanlab.com) into a [hybrid Pinecone datastore](https://docs.pinecone.io/docs/hybrid-search) using [OpenAI ada-002 embeddings](https://platform.openai.com/docs/guides/embeddings) for the dense vectors and a [HuggingFace SPLADE model](https://huggingface.co/naver/splade-cocondenser-ensembledistil) for the sparse embeddings.

src/datastore/datastore.ts

@@ -42,7 +42,7 @@ export abstract class AbstractDatastore<
     this.context = args.context ?? {};
     this.events = args.events ?? {};
     if (args.debug) {
-      this.mergeEvents(args.events ?? {}, {
+      this.addEvents({
         onQueryStart: [console.debug],
         onQueryComplete: [console.debug],
         onError: [console.error],
@@ -144,6 +144,26 @@ export abstract class AbstractDatastore<
     }
   }
 
+  /** Get the current event handlers */
+  getEvents() {
+    return this.events;
+  }
+
+  /** Add event handlers to the datastore. */
+  addEvents(events: typeof this.events): this {
+    this.events = this.mergeEvents(this.events, events);
+    return this;
+  }
+
+  /**
+   * Set the event handlers to a new set of events. Removes all existing event handlers.
+   * Set to empty object `{}` to remove all events.
+   */
+  setEvents(events: typeof this.events): this {
+    this.events = events;
+    return this;
+  }
+
   protected mergeEvents(
     existingEvents: typeof this.events,
     newEvents: typeof this.events

src/model/chat.test.ts

@@ -8,6 +8,7 @@ const FAKE_RESPONSE: Model.Chat.Response = {
     role: 'assistant',
   },
   cached: false,
+  latency: 0,
   cost: 0,
   created: 0,
   id: 'fake-id',

src/model/chat.ts

@@ -34,7 +34,7 @@ export class ChatModel extends AbstractModel<
     params = params ?? { model: 'gpt-3.5-turbo' };
     super({ client, params, ...rest });
     if (args?.debug) {
-      this.mergeEvents(args.events || {}, {
+      this.addEvents({
         onStart: [logInput],
         onComplete: [logResponse],
       });
@@ -72,6 +72,7 @@ export class ChatModel extends AbstractModel<
         ...response,
         message: response.choices[0].message,
         cached: false,
+        latency: Date.now() - start,
         cost: calculateCost({ model: params.model, tokens: response.usage }),
       };
 
@@ -174,6 +175,7 @@ export class ChatModel extends AbstractModel<
         ...response,
         message: response.choices[0].message,
         cached: false,
+        latency: Date.now() - start,
         cost: calculateCost({ model: params.model, tokens: response.usage }),
       };
 
@@ -214,14 +216,14 @@ function logResponse(args: {
       completion_tokens: number;
       prompt_tokens: number;
     };
-    latency?: number;
     cached: boolean;
+    latency?: number;
     choices: { message: Model.Message }[];
     cost?: number;
   };
   params: { messages: Model.Message[] };
 }) {
-  const { usage, latency, cost, choices } = args.response;
+  const { usage, cost, latency, choices } = args.response;
   const tokens = {
     prompt: usage?.prompt_tokens ?? 0,
     completion: usage?.completion_tokens ?? 0,
@@ -229,7 +231,7 @@ function logResponse(args: {
   };
   const message = choices[0].message;
   const tokensStr = `[Tokens: ${tokens.prompt} + ${tokens.completion} = ${tokens.total}]`;
-  const latencyStr = `[Latency: ${latency}ms]`;
+  const latencyStr = latency ? `[Latency: ${latency}ms]` : '';
   const costStr =
     typeof cost === 'number'
       ? `[Cost: $${(cost / 100).toFixed(5)}]`

src/model/types.ts

@@ -34,6 +34,7 @@ export namespace Model {
     export interface Params extends Config, Run {}
     export interface Response {
       cached: boolean;
+      latency?: number;
       cost?: number;
     }
     export type Model = AbstractModel<Client, Config, Run, Response, any>;

src/model/utils/tokenizer.ts

@@ -24,6 +24,9 @@ const GPT_4_MODELS = [
   'gpt-4-0314',
   'gpt-4-32k',
   'gpt-4-32k-0314',
+  'gpt-4-turbo',
+  'gpt-4-1106',
+  'gpt-4-1106-preview',
 ] as const;
 type Gpt4ModelName = (typeof GPT_4_MODELS)[number];
 

src/prompt/functions/ai-extract-function.ts

@@ -0,0 +1,64 @@
+import type { z } from 'zod';
+import type { Model } from '../../model/index.js';
+import { Msg, type Prompt } from '../index.js';
+import { createAIFunction } from './ai-function.js';
+import { createAIRunner } from './ai-runner.js';
+
+/**
+ * Use OpenAI function calling to extract data from a message.
+ */
+export function createAIExtractFunction<Schema extends z.ZodObject<any>>({
+  chatModel,
+  name,
+  description,
+  schema,
+  maxRetries = 0,
+  systemMessage,
+}: {
+  /** The ChatModel used to make API calls. */
+  chatModel: Model.Chat.Model;
+  /** The name of the extractor function. */
+  name: string;
+  /** The description of the extractor function. */
+  description?: string;
+  /** The Zod schema for the data to extract. */
+  schema: Schema;
+  /** The maximum number of times to retry the function call. */
+  maxRetries?: number;
+  /** Add a system message to the beginning of the messages array. */
+  systemMessage?: string;
+}): Prompt.ExtractFunction<Schema> {
+  // The AIFunction that will be used to extract the data
+  const extractFunction = createAIFunction(
+    {
+      name,
+      description,
+      argsSchema: schema,
+    },
+    async (args) => args
+  );
+
+  // Create a runner that will call the function, validate the args and retry
+  // if necessary, and return the result.
+  const runner = createAIRunner({
+    chatModel,
+    systemMessage,
+    functions: [extractFunction],
+    mode: 'functions',
+    maxIterations: maxRetries + 1,
+    params: {
+      function_call: { name },
+    },
+    shouldBreakLoop: (message) => Msg.isFuncResult(message),
+    validateContent: (content) => {
+      return extractFunction.parseArgs(content || '');
+    },
+  });
+
+  // Execute the runner and return the extracted data.
+  return async function run(params, context) {
+    const response = await runner(params, context);
+    if (response.status === 'error') throw response.error;
+    return response.content;
+  };
+}

src/prompt/functions/ai-function.ts

@@ -5,7 +5,7 @@ import type { Prompt } from '../types.js';
 import { cleanString } from '../utils/message.js';
 
 /**
- * Create a function meant to be used with OpenAI function calling.
+ * Create a function meant to be used with OpenAI tool or function calling.
  *
  * The returned function will parse the arguments string and call the
  * implementation function with the parsed arguments.

src/prompt/functions/ai-runner.ts

@@ -0,0 +1,220 @@
+import pMap from 'p-map';
+import { Msg, getErrorMsg } from '../index.js';
+import type { Prompt } from '../types.js';
+import type { Model } from '../../index.js';
+
+type RunnerModelParams = Partial<
+  Omit<Model.Chat.Run & Model.Chat.Config, 'messages' | 'functions' | 'tools'>
+>;
+
+/**
+ * Creates a function to run a chat model in a loop
+ * - Handles parsing, running, and inserting responses for function & tool call messages
+ * - Handles errors by adding a message with the error and rerunning the model
+ * - Optionally validates the content of the last message
+ */
+export function createAIRunner<Content extends any = string>(args: {
+  /** The ChatModel used to make API calls. */
+  chatModel: Model.Chat.Model;
+  /** The functions the model can call. */
+  functions?: Prompt.AIFunction[];
+  /** Use this to control when the runner should stop. */
+  shouldBreakLoop?: (msg: Prompt.Msg) => boolean;
+  /** The maximum number of iterations before the runner throws an error. An iteration is a single call to the model/API. */
+  maxIterations?: number;
+  /** The number of function calls to make concurrently. */
+  functionCallConcurrency?: number;
+  /** Parse and validate the content of the last message. */
+  validateContent?: (content: string | null) => Content | Promise<Content>;
+  /** Controls whether functions or tool_calls are used. */
+  mode?: Prompt.Runner.Mode;
+  /** Add a system message to the beginning of the messages array. */
+  systemMessage?: string;
+  /** Model params to use for each API call (optional). */
+  params?: RunnerModelParams;
+}): Prompt.Runner<Content> {
+  /** Return the content string or an empty string if null. */
+  function defaultValidateContent(content: string | null): Content {
+    return (content ?? '') as Content;
+  }
+
+  /** Break when an assistant message with content is received. */
+  function defaultShouldBreakLoop(msg: Prompt.Msg): boolean {
+    return msg.role === 'assistant' && msg.content !== null;
+  }
+
+  /** Execute the runner and return the messages and content of the last message. */
+  return async function run(params, context) {
+    const {
+      chatModel,
+      functions,
+      mode = 'tools',
+      maxIterations = 5,
+      functionCallConcurrency,
+      systemMessage,
+      params: runnerModelParams,
+      validateContent = defaultValidateContent,
+      shouldBreakLoop = defaultShouldBreakLoop,
+    } = args;
+
+    // Add the functions/tools to the model params
+    const additonalParams = getParams({ functions, mode });
+
+    // Create a message from the input if it's a string
+    const { messages, ...modelParams }: Model.Chat.Run =
+      typeof params === 'string'
+        ? {
+            messages: [Msg.user(params)],
+          }
+        : params;
+    if (systemMessage) {
+      messages.unshift(Msg.system(systemMessage));
+    }
+
+    let iterations = 0;
+
+    // Iterate until the shouldBreakLoop function returns true or the maxIterations
+    // is reached
+    while (iterations < maxIterations) {
+      iterations++;
+
+      try {
+        // Run the model with the current messages and functions/tools
+        const runParams = {
+          ...runnerModelParams,
+          ...modelParams,
+          ...additonalParams,
+          messages,
+        };
+        const { message } = await chatModel.run(runParams, context);
+        messages.push(message);
+
+        // Run functions from tool/function call messages and append the new messages
+        const newMessages = await handleFunctionCallMessage({
+          message,
+          functions,
+          functionCallConcurrency,
+        });
+        messages.push(...newMessages);
+
+        // Check if the last message should break the loop
+        const lastMessage = messages[messages.length - 1];
+        if (shouldBreakLoop(lastMessage)) {
+          const content = await Promise.resolve(
+            validateContent(lastMessage.content)
+          );
+          return { status: 'success', messages, content };
+        }
+      } catch (error: any) {
+        // Halt the runner and return an error if the error is an AbortError
+        if (error.name === 'AbortError') {
+          return { status: 'error', messages, error };
+        }
+
+        // Otherwise, create a message with the error and continue iterating
+        const errMessage = getErrorMsg(error);
+        messages.push(
+          Msg.user(
+            `There was an error validating the response. Please check the error message and try again.\nError:\n${errMessage}`
+          )
+        );
+      }
+    }
+
+    // Return an error if the maxIterations is reached
+    const error = new Error(
+      `Failed to get a valid response from the model after ${maxIterations} iterations.`
+    );
+    return { status: 'error', messages, error };
+  };
+}
+
+/** Get the chat model params for the tools or functions. */
+function getParams(args: {
+  functions?: Prompt.AIFunction[];
+  mode: Prompt.Runner.Mode;
+}): Pick<Model.Chat.Config, 'functions' | 'tools'> {
+  const { functions } = args;
+  // Return an empty object if there are no functions
+  if (!functions?.length) {
+    return {};
+  }
+
+  // Use the functions mode if explicitly set
+  if (args.mode === 'functions') {
+    return { functions: functions.map((func) => func.spec) };
+  }
+
+  // Otherwise, default to the tools mode
+  return {
+    tools: functions.map((func) => ({
+      type: 'function' as const,
+      function: func.spec,
+    })),
+  };
+}
+
+/**
+ * Handle messages that require calling functions.
+ * @returns An array of the new messages from the function calls
+ * Note: Does not include args.message in the returned array
+ */
+export async function handleFunctionCallMessage(args: {
+  message: Prompt.Msg;
+  functions?: Prompt.AIFunction[];
+  functionCallConcurrency?: number;
+}): Promise<Prompt.Msg[]> {
+  const { message, functions = [], functionCallConcurrency = 8 } = args;
+  const messages: Prompt.Msg[] = [message];
+  const funcMap = getFuncMap(functions);
+
+  /** Call a function and return the result. */
+  async function callFunction(args: { name: string; arguments: string }) {
+    const { arguments: funcArgs, name } = args;
+    const func = funcMap.get(name);
+    if (!func) {
+      throw new Error(`No function found with name: ${name}`);
+    }
+
+    try {
+      return await func(funcArgs);
+    } catch (err: any) {
+      if (err.name === 'AbortError') {
+        throw err;
+      }
+
+      // Augment any function errors with the name of the function
+      throw new Error(`Error running function "${name}": ${err.message}`, {
+        cause: err,
+      });
+    }
+  }
+
+  // Run the function with the given name and arguments and add the result messages.
+  if (Msg.isFuncCall(message)) {
+    const result = await callFunction(message.function_call);
+    messages.push(Msg.funcResult(result, message.function_call.name));
+  }
+
+  // Run all the tool_calls functions and add the result messages.
+  if (Msg.isToolCall(message)) {
+    await pMap(
+      message.tool_calls,
+      async (toolCall) => {
+        const result = await callFunction(toolCall.function);
+        messages.push(Msg.toolResult(result, toolCall.id));
+      },
+      { concurrency: functionCallConcurrency }
+    );
+  }
+
+  return messages.slice(1);
+}
+
+/** Create a map of function names to functions for easy lookup. */
+function getFuncMap(functions: Prompt.AIFunction[]) {
+  return functions.reduce((map, func) => {
+    map.set(func.spec.name, func);
+    return map;
+  }, new Map<string, Prompt.AIFunction>());
+}

src/prompt/functions/extract-json.ts

@@ -8,12 +8,15 @@ export function extractJsonObject(str: string): JsonObject {
   const start = str.indexOf('{');
   const end = str.lastIndexOf('}');
   const objectString = str.slice(start, end ? end + 1 : str.length);
+  let repairedObjectString = objectString;
+
   try {
-    const json: JsonObject = JSON.parse(jsonrepair(objectString));
+    repairedObjectString = jsonrepair(objectString);
+    const json: JsonObject = JSON.parse(repairedObjectString);
     return json;
   } catch (error) {
     // Parse again with parse-json for better error messages.
-    const json: JsonObject = parseJson(objectString);
+    const json: JsonObject = parseJson(repairedObjectString);
     return json;
   }
 }

src/prompt/functions/stringify-for-model.test.ts

@@ -0,0 +1,20 @@
+import { describe, expect, it } from 'vitest';
+import { stringifyForModel } from './stringify-for-model.js';
+
+describe('stringifyForModel', () => {
+  it('handles basic objects', () => {
+    const input = {
+      foo: 'bar',
+      nala: ['is', 'cute'],
+      kittens: null,
+      cats: undefined,
+      paws: 4.3,
+    };
+    const result = stringifyForModel(input);
+    expect(result).toEqual(JSON.stringify(input, null));
+  });
+  it('handles empty input', () => {
+    const result = stringifyForModel();
+    expect(result).toEqual('');
+  });
+});

src/prompt/functions/stringify-for-model.ts

@@ -0,0 +1,18 @@
+import type { Jsonifiable } from 'type-fest';
+
+/**
+ * Stringifies a JSON value in a way that's optimized for use with LLM prompts.
+ *
+ * This is intended to be used with `function` and `tool` arguments and responses.
+ */
+export function stringifyForModel(jsonObject: Jsonifiable | void): string {
+  if (jsonObject === undefined) {
+    return '';
+  }
+
+  if (typeof jsonObject === 'string') {
+    return jsonObject;
+  }
+
+  return JSON.stringify(jsonObject, null, 0);
+}

src/prompt/index.ts

@@ -2,6 +2,14 @@ export type { Prompt } from './types.js';
 
 export { Msg } from './utils/message.js';
 export { createAIFunction } from './functions/ai-function.js';
+export {
+  createAIRunner,
+  handleFunctionCallMessage,
+} from './functions/ai-runner.js';
+export { createAIExtractFunction } from './functions/ai-extract-function.js';
 export { extractJsonObject } from './functions/extract-json.js';
 export { extractZodObject } from './functions/extract-zod-object.js';
 export { zodToJsonSchema } from './functions/zod-to-json.js';
+export { getErrorMsg } from './utils/get-error-message.js';
+export { stringifyForModel } from './functions/stringify-for-model.js';
+export * from './utils/errors.js';

src/prompt/types.ts

@@ -1,24 +1,47 @@
 import type { z } from 'zod';
+import type { SetOptional } from 'type-fest';
+import type { Model } from '../index.js';
 
 export namespace Prompt {
   /**
-   * A prompt chain that coordinates the template, functions, and validator.
+   * A runner that iteratively calls the model and handles function calls.
    */
-  export type Chain<Args extends Record<string, any>, Result extends any> = (
-    args: Args
-  ) => Promise<Result>;
+  export type Runner<Content extends any = string> = (
+    params: string | Runner.Params,
+    context?: Model.Ctx
+  ) => Promise<Runner.Response<Content>>;
+
+  export namespace Runner {
+    /** Parameters to execute a runner */
+    export type Params = SetOptional<
+      Model.Chat.Run & Model.Chat.Config,
+      'model'
+    >;
+
+    /** Response from executing a runner */
+    export type Response<Content extends any = string> =
+      | {
+          status: 'success';
+          messages: Prompt.Msg[];
+          content: Content;
+        }
+      | {
+          status: 'error';
+          messages: Prompt.Msg[];
+          error: Error;
+        };
+
+    /** Controls use of functions or tool_calls from OpenAI API */
+    export type Mode = 'tools' | 'functions';
+  }
 
   /**
-   * Turn structured data into a message.
+   * A function used to extract data using OpenAI function calling.
    */
-  export type Template<T = Record<string, any>> = (
-    params: T
-  ) => Promise<Msg[]> | Msg[];
-
-  /**
-   * Validate the output of an LLM call
-   */
-  export type Validator<T> = (input: Msg) => Promise<T> | T;
+  export type ExtractFunction<Schema extends z.ZodObject<any>> = (
+    params: string | Runner.Params,
+    context?: Model.Ctx
+  ) => Promise<z.infer<Schema>>;
 
   export interface AIFunctionSpec {
     name: string;

src/prompt/utils/errors.ts

@@ -0,0 +1,19 @@
+export class AbortError extends Error {
+  readonly name: 'AbortError';
+  readonly originalError: Error;
+
+  constructor(message: string | Error) {
+    super();
+
+    if (message instanceof Error) {
+      this.originalError = message;
+      ({ message } = message);
+    } else {
+      this.originalError = new Error(message);
+      this.originalError.stack = this.stack;
+    }
+
+    this.name = 'AbortError';
+    this.message = message;
+  }
+}

src/prompt/utils/message.ts

@@ -1,11 +1,15 @@
 import dedent from 'dedent';
+import type { Jsonifiable } from 'type-fest';
 import type { Prompt } from '../types.js';
+import { stringifyForModel } from '../functions/stringify-for-model.js';
 
 /**
  * Clean a string by removing extra newlines and indentation.
  * @see: https://github.com/dmnd/dedent
  */
 export function cleanString(text: string): string {
+  // TODO: Should this trim the output as well? could be useful for multiline
+  // templated strings which begin or end with unnecessary newlines.
   const dedenter = dedent.withOptions({ escapeSpecialCharacters: true });
   return dedenter(text);
 }
@@ -79,22 +83,17 @@ export class Msg {
       name?: string;
     }
   ): Prompt.Msg.FuncCall {
-    const { name: msgName } = opts ?? {};
     return {
+      ...opts,
       role: 'assistant',
       content: null,
       function_call,
-      ...(msgName ? { name: msgName } : {}),
     };
   }
 
   /** Create a function result message. */
-  static funcResult(
-    content: string | object | unknown[],
-    name: string
-  ): Prompt.Msg.FuncResult {
-    const contentString =
-      typeof content === 'string' ? content : JSON.stringify(content);
+  static funcResult(content: Jsonifiable, name: string): Prompt.Msg.FuncResult {
+    const contentString = stringifyForModel(content);
     return { role: 'function', content: contentString, name };
   }
 
@@ -106,23 +105,25 @@ export class Msg {
       name?: string;
     }
   ): Prompt.Msg.ToolCall {
-    const { name: msgName } = opts ?? {};
     return {
+      ...opts,
       role: 'assistant',
       content: null,
       tool_calls,
-      ...(msgName ? { name: msgName } : {}),
     };
   }
 
   /** Create a tool call result message. */
   static toolResult(
-    content: string | object | unknown[],
-    tool_call_id: string
+    content: Jsonifiable,
+    tool_call_id: string,
+    opts?: {
+      /** The name of the tool which was called */
+      name?: string;
+    }
   ): Prompt.Msg.ToolResult {
-    const contentString =
-      typeof content === 'string' ? content : JSON.stringify(content);
-    return { role: 'tool', tool_call_id, content: contentString };
+    const contentString = stringifyForModel(content);
+    return { ...opts, role: 'tool', tool_call_id, content: contentString };
   }
 
   /** Get the narrowed message from an EnrichedResponse. */