fix: standardize token count properties and add Number conversion (#367)

- Rename token count properties from prompt_tokens/completion_tokens to input_tokens/output_tokens - Add Number() conversion for token counts in trace utilities - Ensure consistent token property naming across all files - Replicate changes from langtrace-internal PR #267 Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Scale3-Labs · Dec 12, 2024 · b58d90c · b58d90c
1 parent 3caf605
commit b58d90c
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 65 deletions.
diff --git a/components/project/traces/conversation-row.tsx b/components/project/traces/conversation-row.tsx
@@ -86,12 +86,12 @@ export const ConversationRow = ({
         attributes["gen_ai.usage.completion_tokens"]
       ) {
         tokenCounts = {
-          input_tokens: tokenCounts.prompt_tokens
-            ? tokenCounts.prompt_tokens +
+          input_tokens: tokenCounts.input_tokens
+            ? tokenCounts.input_tokens +
               attributes["gen_ai.usage.prompt_tokens"]
             : attributes["gen_ai.usage.prompt_tokens"],
-          output_tokens: tokenCounts.completion_tokens
-            ? tokenCounts.completion_tokens +
+          output_tokens: tokenCounts.output_tokens
+            ? tokenCounts.output_tokens +
               attributes["gen_ai.usage.completion_tokens"]
             : attributes["gen_ai.usage.completion_tokens"],
           total_tokens: tokenCounts.total_tokens
@@ -111,12 +111,12 @@ export const ConversationRow = ({
         attributes["gen_ai.usage.output_tokens"]
       ) {
         tokenCounts = {
-          input_tokens: tokenCounts.prompt_tokens
-            ? tokenCounts.prompt_tokens +
+          input_tokens: tokenCounts.input_tokens
+            ? tokenCounts.input_tokens +
               attributes["gen_ai.usage.input_tokens"]
             : attributes["gen_ai.usage.input_tokens"],
-          output_tokens: tokenCounts.completion_tokens
-            ? tokenCounts.completion_tokens +
+          output_tokens: tokenCounts.output_tokens
+            ? tokenCounts.output_tokens +
               attributes["gen_ai.usage.output_tokens"]
             : attributes["gen_ai.usage.output_tokens"],
           total_tokens: tokenCounts.total_tokens

diff --git a/lib/crewai_trace_util.ts b/lib/crewai_trace_util.ts
@@ -297,20 +297,20 @@ export function processCrewAITrace(trace: any): CrewAITrace {
         attributes["gen_ai.usage.completion_tokens"]
       ) {
         tokenCounts = {
-          input_tokens: tokenCounts.prompt_tokens
-            ? tokenCounts.prompt_tokens +
-              attributes["gen_ai.usage.prompt_tokens"]
-            : attributes["gen_ai.usage.prompt_tokens"],
-          output_tokens: tokenCounts.completion_tokens
-            ? tokenCounts.completion_tokens +
-              attributes["gen_ai.usage.completion_tokens"]
-            : attributes["gen_ai.usage.completion_tokens"],
+          input_tokens: tokenCounts.input_tokens
+            ? tokenCounts.input_tokens +
+              Number(attributes["gen_ai.usage.prompt_tokens"])
+            : Number(attributes["gen_ai.usage.prompt_tokens"]),
+          output_tokens: tokenCounts.output_tokens
+            ? tokenCounts.output_tokens +
+              Number(attributes["gen_ai.usage.completion_tokens"])
+            : Number(attributes["gen_ai.usage.completion_tokens"]),
           total_tokens: tokenCounts.total_tokens
             ? tokenCounts.total_tokens +
-              attributes["gen_ai.usage.prompt_tokens"] +
-              attributes["gen_ai.usage.completion_tokens"]
-            : attributes["gen_ai.usage.prompt_tokens"] +
-              attributes["gen_ai.usage.completion_tokens"],
+              Number(attributes["gen_ai.usage.prompt_tokens"]) +
+              Number(attributes["gen_ai.usage.completion_tokens"])
+            : Number(attributes["gen_ai.usage.prompt_tokens"]) +
+              Number(attributes["gen_ai.usage.completion_tokens"]),
         };
 
         // calculate the cost of the current span
@@ -325,20 +325,20 @@ export function processCrewAITrace(trace: any): CrewAITrace {
         attributes["gen_ai.usage.output_tokens"]
       ) {
         tokenCounts = {
-          input_tokens: tokenCounts.prompt_tokens
-            ? tokenCounts.prompt_tokens +
-              attributes["gen_ai.usage.input_tokens"]
-            : attributes["gen_ai.usage.input_tokens"],
-          output_tokens: tokenCounts.completion_tokens
-            ? tokenCounts.completion_tokens +
-              attributes["gen_ai.usage.output_tokens"]
-            : attributes["gen_ai.usage.output_tokens"],
+          input_tokens: tokenCounts.input_tokens
+            ? tokenCounts.input_tokens +
+              Number(attributes["gen_ai.usage.input_tokens"])
+            : Number(attributes["gen_ai.usage.input_tokens"]),
+          output_tokens: tokenCounts.output_tokens
+            ? tokenCounts.output_tokens +
+              Number(attributes["gen_ai.usage.output_tokens"])
+            : Number(attributes["gen_ai.usage.output_tokens"]),
           total_tokens: tokenCounts.total_tokens
             ? tokenCounts.total_tokens +
-              attributes["gen_ai.usage.input_tokens"] +
-              attributes["gen_ai.usage.output_tokens"]
-            : attributes["gen_ai.usage.input_tokens"] +
-              attributes["gen_ai.usage.output_tokens"],
+              Number(attributes["gen_ai.usage.input_tokens"]) +
+              Number(attributes["gen_ai.usage.output_tokens"])
+            : Number(attributes["gen_ai.usage.input_tokens"]) +
+              Number(attributes["gen_ai.usage.output_tokens"]),
         };
         const currentcost = calculatePriceFromUsage(vendor, model, tokenCounts);
         // add the cost of the current span to the total cost

diff --git a/lib/dspy_trace_util.ts b/lib/dspy_trace_util.ts
@@ -158,20 +158,20 @@ export function processDspyTrace(trace: any): DspyTrace {
         attributes["gen_ai.usage.completion_tokens"]
       ) {
         tokenCounts = {
-          input_tokens: tokenCounts.prompt_tokens
-            ? tokenCounts.prompt_tokens +
-              attributes["gen_ai.usage.prompt_tokens"]
-            : attributes["gen_ai.usage.prompt_tokens"],
-          output_tokens: tokenCounts.completion_tokens
-            ? tokenCounts.completion_tokens +
-              attributes["gen_ai.usage.completion_tokens"]
-            : attributes["gen_ai.usage.completion_tokens"],
+          input_tokens: tokenCounts.input_tokens
+            ? tokenCounts.input_tokens +
+              Number(attributes["gen_ai.usage.prompt_tokens"])
+            : Number(attributes["gen_ai.usage.prompt_tokens"]),
+          output_tokens: tokenCounts.output_tokens
+            ? tokenCounts.output_tokens +
+              Number(attributes["gen_ai.usage.completion_tokens"])
+            : Number(attributes["gen_ai.usage.completion_tokens"]),
           total_tokens: tokenCounts.total_tokens
             ? tokenCounts.total_tokens +
-              attributes["gen_ai.usage.prompt_tokens"] +
-              attributes["gen_ai.usage.completion_tokens"]
-            : attributes["gen_ai.usage.prompt_tokens"] +
-              attributes["gen_ai.usage.completion_tokens"],
+              Number(attributes["gen_ai.usage.prompt_tokens"]) +
+              Number(attributes["gen_ai.usage.completion_tokens"])
+            : Number(attributes["gen_ai.usage.prompt_tokens"]) +
+              Number(attributes["gen_ai.usage.completion_tokens"]),
         };
 
         // calculate the cost of the current span
@@ -186,20 +186,20 @@ export function processDspyTrace(trace: any): DspyTrace {
         attributes["gen_ai.usage.output_tokens"]
       ) {
         tokenCounts = {
-          input_tokens: tokenCounts.prompt_tokens
-            ? tokenCounts.prompt_tokens +
-              attributes["gen_ai.usage.input_tokens"]
-            : attributes["gen_ai.usage.input_tokens"],
-          output_tokens: tokenCounts.completion_tokens
-            ? tokenCounts.completion_tokens +
-              attributes["gen_ai.usage.output_tokens"]
-            : attributes["gen_ai.usage.output_tokens"],
+          input_tokens: tokenCounts.input_tokens
+            ? tokenCounts.input_tokens +
+              Number(attributes["gen_ai.usage.input_tokens"])
+            : Number(attributes["gen_ai.usage.input_tokens"]),
+          output_tokens: tokenCounts.output_tokens
+            ? tokenCounts.output_tokens +
+              Number(attributes["gen_ai.usage.output_tokens"])
+            : Number(attributes["gen_ai.usage.output_tokens"]),
           total_tokens: tokenCounts.total_tokens
             ? tokenCounts.total_tokens +
-              attributes["gen_ai.usage.input_tokens"] +
-              attributes["gen_ai.usage.output_tokens"]
-            : attributes["gen_ai.usage.input_tokens"] +
-              attributes["gen_ai.usage.output_tokens"],
+              Number(attributes["gen_ai.usage.input_tokens"]) +
+              Number(attributes["gen_ai.usage.output_tokens"])
+            : Number(attributes["gen_ai.usage.input_tokens"]) +
+              Number(attributes["gen_ai.usage.output_tokens"]),
         };
         const currentcost = calculatePriceFromUsage(vendor, model, tokenCounts);
         // add the cost of the current span to the total cost

diff --git a/lib/trace_util.ts b/lib/trace_util.ts
@@ -105,14 +105,14 @@ export function processTrace(trace: any): Trace {
         attributes["gen_ai.usage.completion_tokens"]
       ) {
         tokenCounts = {
-          input_tokens: tokenCounts.prompt_tokens
-            ? Number(tokenCounts.prompt_tokens) +
+          input_tokens: tokenCounts.input_tokens
+            ? Number(tokenCounts.input_tokens) +
               Number(attributes["gen_ai.usage.prompt_tokens"])
             : Number(attributes["gen_ai.usage.prompt_tokens"]),
-          output_tokens: tokenCounts.completion_tokens
-            ? Number(tokenCounts.completion_tokens) +
+          output_tokens: tokenCounts.output_tokens
+            ? Number(tokenCounts.output_tokens) +
               Number(attributes["gen_ai.usage.completion_tokens"])
-            : Number(attributes["gen_ai.usage.completion_tokens"]), 
+            : Number(attributes["gen_ai.usage.completion_tokens"]),
           total_tokens: tokenCounts.total_tokens
             ? Number(tokenCounts.total_tokens) +
               Number(attributes["gen_ai.usage.prompt_tokens"]) +
@@ -133,12 +133,12 @@ export function processTrace(trace: any): Trace {
         attributes["gen_ai.usage.output_tokens"]
       ) {
         tokenCounts = {
-          input_tokens: tokenCounts.prompt_tokens
-            ? Number(tokenCounts.prompt_tokens) +
+          input_tokens: tokenCounts.input_tokens
+            ? Number(tokenCounts.input_tokens) +
               Number(attributes["gen_ai.usage.input_tokens"])
             : Number(attributes["gen_ai.usage.input_tokens"]),
-          output_tokens: tokenCounts.completion_tokens
-            ? Number(tokenCounts.completion_tokens) +
+          output_tokens: tokenCounts.output_tokens
+            ? Number(tokenCounts.output_tokens) +
               Number(attributes["gen_ai.usage.output_tokens"])
             : Number(attributes["gen_ai.usage.output_tokens"]),
           total_tokens: tokenCounts.total_tokens