feat: 🎸 compute token usage for function / tool streaming

packages/lib/chat-model.ts

@@ -1,13 +1,11 @@
 import OpenAI, { ClientOptions } from 'openai';
 import { RunnableToolFunction } from 'openai/lib/RunnableFunction';
-import {
-  ChatCompletionMessageParam,
-  ChatCompletionTool,
-  CompletionUsage,
-} from 'openai/resources';
+import { CompletionUsage } from 'openai/resources';
 import pRetry from 'p-retry';
 
+import { countTokensEstimation } from './count-tokens';
 import failedAttemptHandler from './lc-failed-attempt-hanlder';
+import { promptTokensEstimate } from './tokens-estimate';
 
 const list = () => ['mistery'];
 
@@ -42,16 +40,6 @@ export default class ChatModel {
     });
   }
 
-  static countTokensMessages(messages: ChatCompletionMessageParam[]) {
-    let counter = 0;
-
-    for (const each of messages) {
-      counter += each?.content?.length || 0;
-    }
-
-    return counter / 4;
-  }
-
   async call({
     handleStream,
     signal,
@@ -65,7 +53,14 @@ export default class ChatModel {
       async () => {
         let usage: CompletionUsage = {
           completion_tokens: 0,
-          prompt_tokens: ChatModel.countTokensMessages(otherProps?.messages),
+          prompt_tokens: promptTokensEstimate({
+            tools,
+            useFastApproximation: true,
+            tool_choice: otherProps?.tool_choice,
+            messages: otherProps?.messages,
+            functions: otherProps?.functions,
+            function_call: otherProps?.function_call,
+          }),
           total_tokens: 0,
         };
 
@@ -108,7 +103,10 @@ export default class ChatModel {
 
               handleStream?.(content);
               buffer += content;
-              usage.completion_tokens += 1;
+
+              usage.completion_tokens += countTokensEstimation({
+                text: content || '',
+              });
             }
 
             usage.total_tokens = usage.prompt_tokens + usage.completion_tokens;
@@ -123,6 +121,10 @@ export default class ChatModel {
               stream: false,
             });
 
+            usage.completion_tokens = countTokensEstimation({
+              text: response?.choices?.[0]?.message?.content || '',
+            });
+
             usage.total_tokens = usage.prompt_tokens + usage.completion_tokens;
 
             return {

packages/lib/get-usage-cost.ts

@@ -0,0 +1,20 @@
+import { CompletionUsage } from 'openai/resources';
+
+import { AgentModelName } from '@chaindesk/prisma';
+
+import { ModelConfig } from './config';
+
+type Props = {};
+
+function getUsageCost(props: {
+  modelName: AgentModelName;
+  usage: CompletionUsage;
+}) {
+  return (
+    (props.usage?.prompt_tokens || 0) *
+      ModelConfig[props.modelName]?.providerPriceByInputToken +
+    (props.usage?.completion_tokens || 0) *
+      ModelConfig[props.modelName]?.providerPricePriceByOutputToken
+  );
+}
+export default getUsageCost;

packages/lib/tokens-estimate/functions.ts

@@ -0,0 +1,126 @@
+import OpenAI from 'openai';
+
+type OpenAIFunction = OpenAI.Chat.ChatCompletionCreateParams.Function;
+
+// Types representing the OpenAI function definitions. While the OpenAI client library
+// does have types for function definitions, the properties are just Record<string, unknown>,
+// which isn't very useful for type checking this formatting code.
+export interface FunctionDef extends Omit<OpenAIFunction, 'parameters'> {
+  name: string;
+  description?: string;
+  parameters: ObjectProp;
+}
+
+interface ObjectProp {
+  type: 'object';
+  properties?: {
+    [key: string]: Prop;
+  };
+  required?: string[];
+}
+
+interface AnyOfProp {
+  anyOf: Prop[];
+}
+
+type Prop = {
+  description?: string;
+} & (
+  | AnyOfProp
+  | ObjectProp
+  | {
+      type: 'string';
+      enum?: string[];
+    }
+  | {
+      type: 'number' | 'integer';
+      minimum?: number;
+      maximum?: number;
+      enum?: number[];
+    }
+  | { type: 'boolean' }
+  | { type: 'null' }
+  | {
+      type: 'array';
+      items?: Prop;
+    }
+);
+
+function isAnyOfProp(prop: Prop): prop is AnyOfProp {
+  return (
+    (prop as AnyOfProp).anyOf !== undefined &&
+    Array.isArray((prop as AnyOfProp).anyOf)
+  );
+}
+
+// When OpenAI use functions in the prompt, they format them as TypeScript definitions rather than OpenAPI JSON schemas.
+// This function converts the JSON schemas into TypeScript definitions.
+export function formatFunctionDefinitions(functions: FunctionDef[]) {
+  const lines = ['namespace functions {', ''];
+  for (const f of functions) {
+    if (f.description) {
+      lines.push(`// ${f.description}`);
+    }
+    if (Object.keys(f.parameters.properties ?? {}).length > 0) {
+      lines.push(`type ${f.name} = (_: {`);
+      lines.push(formatObjectProperties(f.parameters, 0));
+      lines.push('}) => any;');
+    } else {
+      lines.push(`type ${f.name} = () => any;`);
+    }
+    lines.push('');
+  }
+  lines.push('} // namespace functions');
+  return lines.join('\n');
+}
+
+// Format just the properties of an object (not including the surrounding braces)
+function formatObjectProperties(obj: ObjectProp, indent: number): string {
+  const lines = [];
+  for (const [name, param] of Object.entries(obj.properties ?? {})) {
+    if (param.description && indent < 2) {
+      lines.push(`// ${param.description}`);
+    }
+    if (obj.required?.includes(name)) {
+      lines.push(`${name}: ${formatType(param, indent)},`);
+    } else {
+      lines.push(`${name}?: ${formatType(param, indent)},`);
+    }
+  }
+  return lines.map((line) => ' '.repeat(indent) + line).join('\n');
+}
+
+// Format a single property type
+function formatType(param: Prop, indent: number): string {
+  if (isAnyOfProp(param)) {
+    return param.anyOf.map((v) => formatType(v, indent)).join(' | ');
+  }
+  switch (param.type) {
+    case 'string':
+      if (param.enum) {
+        return param.enum.map((v) => `"${v}"`).join(' | ');
+      }
+      return 'string';
+    case 'number':
+      if (param.enum) {
+        return param.enum.map((v) => `${v}`).join(' | ');
+      }
+      return 'number';
+    case 'integer':
+      if (param.enum) {
+        return param.enum.map((v) => `${v}`).join(' | ');
+      }
+      return 'number';
+    case 'boolean':
+      return 'boolean';
+    case 'null':
+      return 'null';
+    case 'object':
+      return ['{', formatObjectProperties(param, indent + 2), '}'].join('\n');
+    case 'array':
+      if (param.items) {
+        return `${formatType(param.items, indent)}[]`;
+      }
+      return 'any[]';
+  }
+}