Add cost calculation.

README.md

@@ -4,8 +4,7 @@
 
 JS Agent is a composable and extensible framework for creating GPT agents with JavaScript and TypeScript.
 
-While creating an agent prototype is easy, increasing its reliability and robustness is very hard
-and requires considerable experimentation. JS Agent provides building blocks and tooling to help you develop rock-solid agents faster.
+While creating an agent prototype is easy, increasing its reliability and robustness is difficult and requires considerable experimentation. JS Agent provides robust building blocks and tooling to help you develop rock-solid agents faster.
 
 **⚠️ JS Agent is currently in its initial experimental phase. Prior to reaching version 0.1, there may breaking changes in each release.**
 
@@ -20,9 +19,10 @@ See examples below for details on how to implement and run an agent.
 ## Features
 
 - Agent definition and execution
-  - Observable agent runs (to support console output, UIs, server runs, webapps, etc.)
-  - Recording of LLM calls for each agent run
-  - Controller to limit the number of steps
+  - Observe agent runs (to support console output, UIs, server runs, webapps, etc.)
+  - Record all LLM calls of an agent run
+  - Calculate the cost of LLM calls and agent runs
+  - Stop agent runs when certain criteria are met, e.g. to limit the number of steps
 - Supported LLM models
   - OpenAI text completion models (`text-davinci-003` etc.)
   - OpenAI chat completion models (`gpt-4`, `gpt-3.5-turbo`)
@@ -145,9 +145,26 @@ export async function runWikipediaAgent({
       }),
     }),
     controller: $.agent.controller.maxSteps(20),
-    observer: $.agent.showRunInConsole({
-      name: "Wikipedia Agent",
-    }),
+    observer: $.agent.observer.combineObservers(
+      $.agent.observer.showRunInConsole({ name: "Wikipedia Agent" }),
+      {
+        async onRunFinished({ run }) {
+          const runCostInMillicent = await $.agent.calculateRunCostInMillicent({
+            run,
+          });
+
+          console.log(
+            `Run cost: $${(runCostInMillicent / 1000 / 100).toFixed(2)}`
+          );
+
+          console.log(
+            `LLM calls: ${
+              run.recordedCalls.filter((call) => call.success).length
+            }`
+          );
+        },
+      }
+    ),
   });
 }
 ```

examples/wikipedia/README.md

@@ -8,6 +8,7 @@ Answers questions using Wikipedia articles. It searches using a Programmable Sea
 - Custom tool configuration (`readWikipediaArticleAction`, `searchWikipediaAction`)
 - `GenerateNextStepLoop` loop with tools and custom prompt
 - `maxSteps` `RunController` to limit the maximum number of steps
+- Cost calculation and extracting information from LLM calls after the run
 
 ## Usage
 

examples/wikipedia/src/runWikipediaAgent.ts

@@ -80,6 +80,25 @@ export async function runWikipediaAgent({
       }),
     }),
     controller: $.agent.controller.maxSteps(20),
-    observer: $.agent.observer.showRunInConsole({ name: "Wikipedia Agent" }),
+    observer: $.agent.observer.combineObservers(
+      $.agent.observer.showRunInConsole({ name: "Wikipedia Agent" }),
+      {
+        async onRunFinished({ run }) {
+          const runCostInMillicent = await $.agent.calculateRunCostInMillicent({
+            run,
+          });
+
+          console.log(
+            `Run cost: $${(runCostInMillicent / 1000 / 100).toFixed(2)}`
+          );
+
+          console.log(
+            `LLM calls: ${
+              run.recordedCalls.filter((call) => call.success).length
+            }`
+          );
+        },
+      }
+    ),
   });
 }

packages/agent/README.md

@@ -1,9 +1,8 @@
-# JS Agent: Build AI Agents with JS & TS
+# JS Agent: Build GPT Agents with JS & TS
 
-JS Agent is a composable and extensible framework for creating AI agents with JavaScript and TypeScript.
+JS Agent is a composable and extensible framework for creating GPT agents with JavaScript and TypeScript.
 
-While creating an AI agent prototype is easy, increasing its reliability and robustness is very hard
-and requires considerable experimentation. JS Agent provides building blocks and tooling to help you develop rock-solid agents faster.
+While creating an agent prototype is easy, increasing its reliability and robustness is difficult and requires considerable experimentation. JS Agent provides robust building blocks and tooling to help you develop rock-solid agents faster.
 
 **⚠️ JS Agent is currently in its initial experimental phase. Prior to reaching version 0.1, there may breaking changes in each release.**
 
@@ -18,9 +17,10 @@ See examples below for details on how to implement and run an agent.
 ## Features
 
 - Agent definition and execution
-  - Observable agent runs (to support console output, UIs, server runs, webapps, etc.)
-  - Recording of LLM calls for each agent run
-  - Controller to limit the number of steps
+  - Observe agent runs (to support console output, UIs, server runs, webapps, etc.)
+  - Record all LLM calls of an agent run
+  - Calculate the cost of LLM calls and agent runs
+  - Stop agent runs when certain criteria are met, e.g. to limit the number of steps
 - Supported LLM models
   - OpenAI text completion models (`text-davinci-003` etc.)
   - OpenAI chat completion models (`gpt-4`, `gpt-3.5-turbo`)
@@ -143,9 +143,26 @@ export async function runWikipediaAgent({
       }),
     }),
     controller: $.agent.controller.maxSteps(20),
-    observer: $.agent.showRunInConsole({
-      name: "Wikipedia Agent",
-    }),
+    observer: $.agent.observer.combineObservers(
+      $.agent.observer.showRunInConsole({ name: "Wikipedia Agent" }),
+      {
+        async onRunFinished({ run }) {
+          const runCostInMillicent = await $.agent.calculateRunCostInMillicent({
+            run,
+          });
+
+          console.log(
+            `Run cost: $${(runCostInMillicent / 1000 / 100).toFixed(2)}`
+          );
+
+          console.log(
+            `LLM calls: ${
+              run.recordedCalls.filter((call) => call.success).length
+            }`
+          );
+        },
+      }
+    ),
   });
 }
 ```

packages/agent/src/agent/calculateRunCostInMillicent.ts

@@ -0,0 +1,16 @@
+import { calculateOpenAiCallCostInMillicent } from "../provider/openai/calculateOpenAiCallCostInMillicent";
+import { Run } from "./Run";
+
+export const calculateRunCostInMillicent = async ({ run }: { run: Run }) => {
+  const callCostsInMillicent = run.recordedCalls.map((call) => {
+    if (call.success && call.metadata.model.vendor === "openai") {
+      return calculateOpenAiCallCostInMillicent(call);
+    }
+    return undefined;
+  });
+
+  return callCostsInMillicent.reduce(
+    (sum: number, cost) => sum + (cost ?? 0),
+    0
+  );
+};

packages/agent/src/agent/index.ts

@@ -1,6 +1,7 @@
 export * from "./GenerateCall";
 export * from "./Run";
 export * from "./RunContext";
+export * from "./calculateRunCostInMillicent";
 export * as controller from "./controller/index";
 export * as observer from "./observer/index";
 export * from "./runAgent";

packages/agent/src/agent/observer/combineObservers.ts

@@ -0,0 +1,47 @@
+import { RunObserver } from "./RunObserver";
+
+export const combineObservers = (...observers: RunObserver[]): RunObserver => ({
+  onRunStarted: ({ run }) => {
+    observers.forEach((observer) => observer.onRunStarted?.({ run }));
+  },
+
+  onRunFinished: ({ run, result }) => {
+    observers.forEach((observer) => observer.onRunFinished?.({ run, result }));
+  },
+
+  onStepGenerationStarted: ({ run }) => {
+    observers.forEach((observer) =>
+      observer.onStepGenerationStarted?.({ run })
+    );
+  },
+
+  onStepGenerationFinished: ({ run, generatedText, step }) => {
+    observers.forEach((observer) =>
+      observer.onStepGenerationFinished?.({ run, generatedText, step })
+    );
+  },
+
+  onLoopIterationStarted: ({ run, loop }) => {
+    observers.forEach((observer) =>
+      observer.onLoopIterationStarted?.({ run, loop })
+    );
+  },
+
+  onLoopIterationFinished: ({ run, loop }) => {
+    observers.forEach((observer) =>
+      observer.onLoopIterationFinished?.({ run, loop })
+    );
+  },
+
+  onStepExecutionStarted: ({ run, step }) => {
+    observers.forEach((observer) =>
+      observer.onStepExecutionStarted?.({ run, step })
+    );
+  },
+
+  onStepExecutionFinished: ({ run, step, result }) => {
+    observers.forEach((observer) =>
+      observer.onStepExecutionFinished?.({ run, step, result })
+    );
+  },
+});

packages/agent/src/agent/observer/index.ts

@@ -1,2 +1,3 @@
 export * from "./RunObserver";
+export * from "./combineObservers";
 export * from "./showRunInConsole";

packages/agent/src/provider/openai/api/generateChatCompletion.ts

@@ -2,7 +2,7 @@ import axios from "axios";
 import zod from "zod";
 import { OpenAIChatMessage } from "../OpenAIChatMessage";
 
-const OpenAIChatCompletionSchema = zod.object({
+export const OpenAIChatCompletionSchema = zod.object({
   id: zod.string(),
   object: zod.literal("chat.completion"),
   created: zod.number(),

packages/agent/src/provider/openai/api/generateCompletion.ts → packages/agent/src/provider/openai/api/generateTextCompletion.ts

@@ -1,7 +1,7 @@
 import axios from "axios";
 import zod from "zod";
 
-const OpenAICompletionSchema = zod.object({
+export const OpenAITextCompletionSchema = zod.object({
   id: zod.string(),
   object: zod.literal("text_completion"),
   created: zod.number(),
@@ -21,9 +21,9 @@ const OpenAICompletionSchema = zod.object({
   }),
 });
 
-export type OpenAICompletion = zod.infer<typeof OpenAICompletionSchema>;
+export type OpenAITextCompletion = zod.infer<typeof OpenAITextCompletionSchema>;
 
-export type OpenAICompletionModel =
+export type OpenAITextCompletionModel =
   | "text-davinci-003"
   | "text-davinci-002"
   | "code-davinci-002"
@@ -36,7 +36,7 @@ export type OpenAICompletionModel =
   | "babbage"
   | "ada";
 
-export async function generateCompletion({
+export async function generateTextCompletion({
   apiKey,
   prompt,
   model,
@@ -48,13 +48,13 @@ export async function generateCompletion({
 }: {
   apiKey: string;
   prompt: string;
-  model: OpenAICompletionModel;
+  model: OpenAITextCompletionModel;
   n?: number;
   temperature?: number;
   maxTokens?: number;
   presencePenalty?: number;
   frequencyPenalty?: number;
-}): Promise<OpenAICompletion> {
+}): Promise<OpenAITextCompletion> {
   const response = await axios.post(
     "https://api.openai.com/v1/completions",
     JSON.stringify({
@@ -74,5 +74,5 @@ export async function generateCompletion({
     }
   );
 
-  return OpenAICompletionSchema.parse(response.data);
+  return OpenAITextCompletionSchema.parse(response.data);
 }

packages/agent/src/provider/openai/api/index.ts

@@ -1,2 +1,2 @@
 export * from "./generateChatCompletion.js";
-export * from "./generateCompletion.js";
+export * from "./generateTextCompletion.js";

packages/agent/src/provider/openai/calculateChatCompletionCostInMillicent.ts

@@ -0,0 +1,25 @@
+import {
+  OpenAIChatCompletion,
+  OpenAIChatCompletionModel,
+} from "./api/generateChatCompletion";
+
+// see https://openai.com/pricing
+const promptTokenCostInMillicent = {
+  "gpt-4": 3,
+  "gpt-3.5-turbo": 0.2,
+};
+
+const completionTokenCostInMillicent = {
+  "gpt-4": 6,
+  "gpt-3.5-turbo": 0.2,
+};
+
+export const calculateChatCompletionCostInMillicent = ({
+  model,
+  completion,
+}: {
+  model: OpenAIChatCompletionModel;
+  completion: OpenAIChatCompletion;
+}) =>
+  completion.usage.prompt_tokens * promptTokenCostInMillicent[model] +
+  completion.usage.completion_tokens * completionTokenCostInMillicent[model];

packages/agent/src/provider/openai/calculateOpenAiCallCostInMillicent.ts

@@ -0,0 +1,49 @@
+import { GenerateCall } from "../../agent";
+import { OpenAIChatCompletionSchema } from "./api/generateChatCompletion";
+import { OpenAITextCompletionSchema } from "./api/generateTextCompletion";
+import { calculateChatCompletionCostInMillicent } from "./calculateChatCompletionCostInMillicent";
+import { calculateTextCompletionCostInMillicent } from "./calculateTextCompletionCostInMillicent";
+
+export function calculateOpenAiCallCostInMillicent(
+  call: GenerateCall & {
+    success: true;
+  }
+) {
+  if (call.metadata.model.vendor !== "openai") {
+    throw new Error(`Incorrect vendor: ${call.metadata.model.vendor}`);
+  }
+
+  const model = call.metadata.model.name;
+
+  switch (model) {
+    case "gpt-3.5-turbo":
+    case "gpt-4": {
+      return calculateChatCompletionCostInMillicent({
+        model,
+        completion: OpenAIChatCompletionSchema.parse(call.rawOutput),
+      });
+    }
+
+    case "text-davinci-003":
+    case "text-davinci-003":
+    case "text-davinci-002":
+    case "code-davinci-002":
+    case "code-davinci-002":
+    case "text-curie-001":
+    case "text-babbage-001":
+    case "text-ada-001":
+    case "davinci":
+    case "curie":
+    case "babbage":
+    case "ada": {
+      return calculateTextCompletionCostInMillicent({
+        model,
+        completion: OpenAITextCompletionSchema.parse(call.rawOutput),
+      });
+    }
+
+    default: {
+      throw new Error(`Unknown model: ${model}`);
+    }
+  }
+}