cohere[minor]: Fix token counts, add usage_metadata

libs/langchain-cohere/package.json

@@ -36,7 +36,7 @@
   "license": "MIT",
   "dependencies": {
     "@langchain/core": ">0.1.58 <0.3.0",
-    "cohere-ai": "^7.9.3"
+    "cohere-ai": "^7.10.5"
   },
   "devDependencies": {
     "@jest/globals": "^29.5.0",

libs/langchain-cohere/src/chat_models.ts

@@ -62,7 +62,7 @@ export interface CohereChatCallOptions
 
 function convertMessagesToCohereMessages(
   messages: Array<BaseMessage>
-): Array<Cohere.ChatMessage> {
+): Array<Cohere.Message> {
   const getRole = (role: MessageType) => {
     switch (role) {
       case "system":
@@ -113,7 +113,7 @@ function convertMessagesToCohereMessages(
 export class ChatCohere<
     CallOptions extends CohereChatCallOptions = CohereChatCallOptions
   >
-  extends BaseChatModel<CallOptions>
+  extends BaseChatModel<CallOptions, AIMessageChunk>
   implements ChatCohereInput
 {
   static lc_name() {
@@ -193,8 +193,14 @@ export class ChatCohere<
     const cohereMessages = convertMessagesToCohereMessages(messages);
     // The last message in the array is the most recent, all other messages
     // are apart of the chat history.
-    const { message } = cohereMessages[cohereMessages.length - 1];
-    const chatHistory: Cohere.ChatMessage[] = [];
+    const lastMessage = cohereMessages[cohereMessages.length - 1];
+    if (lastMessage.role === "TOOL") {
+      throw new Error(
+        "Cohere does not support tool messages as the most recent message in chat history."
+      );
+    }
+    const { message } = lastMessage;
+    const chatHistory: Cohere.Message[] = [];
     if (cohereMessages.length > 1) {
       chatHistory.push(...cohereMessages.slice(0, -1));
     }
@@ -241,25 +247,22 @@ export class ChatCohere<
         }
       );
 
-    if ("token_count" in response) {
-      const {
-        response_tokens: completionTokens,
-        prompt_tokens: promptTokens,
-        total_tokens: totalTokens,
-      } = response.token_count as Record<string, number>;
+    if (response.meta?.tokens) {
+      const { inputTokens, outputTokens } = response.meta.tokens;
 
-      if (completionTokens) {
+      if (outputTokens) {
         tokenUsage.completionTokens =
-          (tokenUsage.completionTokens ?? 0) + completionTokens;
+          (tokenUsage.completionTokens ?? 0) + outputTokens;
       }
 
-      if (promptTokens) {
-        tokenUsage.promptTokens = (tokenUsage.promptTokens ?? 0) + promptTokens;
+      if (inputTokens) {
+        tokenUsage.promptTokens = (tokenUsage.promptTokens ?? 0) + inputTokens;
       }
 
-      if (totalTokens) {
-        tokenUsage.totalTokens = (tokenUsage.totalTokens ?? 0) + totalTokens;
-      }
+      tokenUsage.totalTokens =
+        (tokenUsage.totalTokens ?? 0) +
+        (tokenUsage.promptTokens ?? 0) +
+        (tokenUsage.completionTokens ?? 0);
     }
 
     const generationInfo: Record<string, unknown> = { ...response };
@@ -271,6 +274,11 @@ export class ChatCohere<
         message: new AIMessage({
           content: response.text,
           additional_kwargs: generationInfo,
+          usage_metadata: {
+            input_tokens: tokenUsage.promptTokens ?? 0,
+            output_tokens: tokenUsage.completionTokens ?? 0,
+            total_tokens: tokenUsage.totalTokens ?? 0,
+          },
         }),
         generationInfo,
       },
@@ -290,8 +298,14 @@ export class ChatCohere<
     const cohereMessages = convertMessagesToCohereMessages(messages);
     // The last message in the array is the most recent, all other messages
     // are apart of the chat history.
-    const { message } = cohereMessages[cohereMessages.length - 1];
-    const chatHistory: Cohere.ChatMessage[] = [];
+    const lastMessage = cohereMessages[cohereMessages.length - 1];
+    if (lastMessage.role === "TOOL") {
+      throw new Error(
+        "Cohere does not support tool messages as the most recent message in chat history."
+      );
+    }
+    const { message } = lastMessage;
+    const chatHistory: Cohere.Message[] = [];
     if (cohereMessages.length > 1) {
       chatHistory.push(...cohereMessages.slice(0, -1));
     }
@@ -335,6 +349,32 @@ export class ChatCohere<
             ...chunk,
           },
         });
+      } else if (
+        chunk.eventType === "stream-end" &&
+        chunk.response.meta?.tokens &&
+        (chunk.response.meta.tokens.inputTokens ||
+          chunk.response.meta.tokens.outputTokens)
+      ) {
+        // stream-end events contain the final token count
+        const input_tokens = chunk.response.meta.tokens.inputTokens ?? 0;
+        const output_tokens = chunk.response.meta.tokens.outputTokens ?? 0;
+        yield new ChatGenerationChunk({
+          text: "",
+          message: new AIMessageChunk({
+            content: "",
+            additional_kwargs: {
+              eventType: "stream-end",
+            },
+            usage_metadata: {
+              input_tokens,
+              output_tokens,
+              total_tokens: input_tokens + output_tokens,
+            },
+          }),
+          generationInfo: {
+            eventType: "stream-end",
+          },
+        });
       }
     }
   }

libs/langchain-cohere/src/tests/chat_models.int.test.ts

@@ -1,6 +1,6 @@
 /* eslint-disable no-promise-executor-return */
 import { test, expect } from "@jest/globals";
-import { HumanMessage } from "@langchain/core/messages";
+import { AIMessageChunk, HumanMessage } from "@langchain/core/messages";
 import { ChatCohere } from "../chat_models.js";
 
 test("ChatCohere can invoke", async () => {
@@ -58,3 +58,48 @@ test("should abort the request", async () => {
     return ret;
   }).rejects.toThrow("AbortError");
 });
+
+test("Stream token count usage_metadata", async () => {
+  const model = new ChatCohere({
+    model: "command-light",
+    temperature: 0,
+  });
+  let res: AIMessageChunk | null = null;
+  for await (const chunk of await model.stream(
+    "Why is the sky blue? Be concise."
+  )) {
+    if (!res) {
+      res = chunk;
+    } else {
+      res = res.concat(chunk);
+    }
+  }
+  console.log(res);
+  expect(res?.usage_metadata).toBeDefined();
+  if (!res?.usage_metadata) {
+    return;
+  }
+  expect(res.usage_metadata.input_tokens).toBe(71);
+  expect(res.usage_metadata.output_tokens).toBeGreaterThan(10);
+  expect(res.usage_metadata.total_tokens).toBe(
+    res.usage_metadata.input_tokens + res.usage_metadata.output_tokens
+  );
+});
+
+test("Invoke token count usage_metadata", async () => {
+  const model = new ChatCohere({
+    model: "command-light",
+    temperature: 0,
+  });
+  const res = await model.invoke("Why is the sky blue? Be concise.");
+  console.log(res);
+  expect(res?.usage_metadata).toBeDefined();
+  if (!res?.usage_metadata) {
+    return;
+  }
+  expect(res.usage_metadata.input_tokens).toBe(71);
+  expect(res.usage_metadata.output_tokens).toBeGreaterThan(10);
+  expect(res.usage_metadata.total_tokens).toBe(
+    res.usage_metadata.input_tokens + res.usage_metadata.output_tokens
+  );
+});