Fixing token counts if output_tokens is 0 + adding embedding cost…

… tables (#292)
Scale3-Labs · Oct 15, 2024 · c93ba60 · c93ba60
1 parent 40fbb39
commit c93ba60
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/lib/constants.ts b/lib/constants.ts
@@ -130,6 +130,20 @@ export const OPENAI_PRICING: Record<string, CostTableEntry> = {
     input: 0.0015,
     output: 0.002,
   },
+  // embedding models
+  "text-embedding-ada-002": {
+    input: 0.0001,
+    output: 0,
+  },
+  "text-embedding-3-small": {
+    input: 0.00002,
+    output: 0,
+  },
+  "text-embedding-3-large": {
+    input: 0.00013,
+    output: 0,
+  },
+
 };
 
 export const ANTHROPIC_PRICING: Record<string, CostTableEntry> = {

diff --git a/lib/trace_util.ts b/lib/trace_util.ts
@@ -122,7 +122,7 @@ export function processTrace(trace: any): Trace {
         cost.input += currentcost.input;
         cost.output += currentcost.output;
       } else if (
-        attributes["gen_ai.usage.input_tokens"] &&
+        attributes["gen_ai.usage.input_tokens"] ||
         attributes["gen_ai.usage.output_tokens"]
       ) {
         tokenCounts = {

diff --git a/lib/utils.ts b/lib/utils.ts
@@ -576,7 +576,7 @@ export function calculatePriceFromUsage(
   } else if (vendor === "openai") {
     // check if model is present as key in OPENAI_PRICING
     let correctModel = model;
-    if (model.includes("gpt") || model.includes("o1")) {
+    if (model.includes("gpt") || model.includes("o1") || model.includes("text-embedding")) {
       if (model.includes("gpt-4o-mini")) {
         correctModel = "gpt-4o-mini";
       } else if (model.includes("gpt-4o")) {