community[minor]: Use usage metadata message field for token reportin…

…g in DatadogLLMObsTracer (#6552) Co-authored-by: Brace Sproul <braceasproul@gmail.com>
langchain-ai · Aug 22, 2024 · b088c33 · b088c33
1 parent 47f73af
commit b088c33
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/libs/langchain-community/src/experimental/callbacks/handlers/datadog.ts b/libs/langchain-community/src/experimental/callbacks/handlers/datadog.ts
@@ -2,7 +2,7 @@ import { BaseCallbackHandlerInput } from "@langchain/core/callbacks/base";
 import { BaseTracer, Run } from "@langchain/core/tracers/base";
 import { getEnvironmentVariable } from "@langchain/core/utils/env";
 import { Document } from "@langchain/core/documents";
-import { BaseMessage } from "@langchain/core/messages";
+import { BaseMessage, isAIMessage } from "@langchain/core/messages";
 import { ChatGeneration } from "@langchain/core/outputs";
 import { KVMap } from "langsmith/schemas";
 
@@ -262,12 +262,13 @@ export class DatadogLLMObsTracer
           messages: outputs?.generations?.flatMap(
             (generations: ChatGeneration[]) =>
               generations.map(({ message, text }) => {
-                const tokenUsage = message?.response_metadata?.tokenUsage;
-                if (tokenUsage) {
+                if (isAIMessage(message) && message?.usage_metadata) {
+                  tokensMetadata.prompt_tokens =
+                    message.usage_metadata.input_tokens;
                   tokensMetadata.completion_tokens =
-                    tokenUsage.completionTokens;
-                  tokensMetadata.prompt_tokens = tokenUsage.promptTokens;
-                  tokensMetadata.total_tokens = tokenUsage.totalTokens;
+                    message.usage_metadata.output_tokens;
+                  tokensMetadata.total_tokens =
+                    message.usage_metadata.total_tokens;
                 }
 
                 return {

diff --git a/libs/langchain-community/src/experimental/callbacks/handlers/tests/datadog.test.ts b/libs/langchain-community/src/experimental/callbacks/handlers/tests/datadog.test.ts
@@ -161,7 +161,7 @@ test("Test llm span with prompt input", async () => {
       generations: [
         [
           {
-            text: "Hi",
+            message: new AIMessage("Hi"),
           },
         ],
       ],
-Original file line number
+Diff line change
@@ Expand Up / @@ -161,7 +161,7 @@ test("Test llm span with prompt input", async () => { @@
           generations: [
             [
               {
-                text: "Hi",
+                message: new AIMessage("Hi"),
               },
             ],
           ],
@@ Expand Down @@