partners[patch]: Fix flaky token counting tests (#6000)

* partners[patch]: Fix flaky token counting tests

* chore: lint files

Author: bracesproul

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

@@ -13,6 +13,16 @@ import {
 import { CallbackManager } from "@langchain/core/callbacks/manager";
 import { ChatAnthropic } from "../chat_models.js";
 
+async function sleep(ms = 1000): Promise<void> {
+  return new Promise<void>((resolve) => {
+    setTimeout(resolve, ms);
+  });
+}
+
+afterEach(async () => {
+  await sleep();
+});
+
 test("Test ChatAnthropic", async () => {
   const chat = new ChatAnthropic({
     modelName: "claude-3-sonnet-20240229",
@@ -323,6 +333,7 @@ test("Stream tokens", async () => {
   const model = new ChatAnthropic({
     model: "claude-3-haiku-20240307",
     temperature: 0,
+    maxTokens: 10,
   });
   let res: AIMessageChunk | null = null;
   for await (const chunk of await model.stream(
@@ -339,8 +350,8 @@ test("Stream tokens", async () => {
   if (!res?.usage_metadata) {
     return;
   }
-  expect(res.usage_metadata.input_tokens).toBe(34);
-  expect(res.usage_metadata.output_tokens).toBeGreaterThan(10);
+  expect(res.usage_metadata.input_tokens).toBeGreaterThan(1);
+  expect(res.usage_metadata.output_tokens).toBeGreaterThan(1);
   expect(res.usage_metadata.total_tokens).toBe(
     res.usage_metadata.input_tokens + res.usage_metadata.output_tokens
   );

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

@@ -81,7 +81,7 @@ test("Stream token count usage_metadata", async () => {
   if (!res?.usage_metadata) {
     return;
   }
-  expect(res.usage_metadata.input_tokens).toBe(71);
+  expect(res.usage_metadata.input_tokens).toBeGreaterThan(1);
   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
@@ -134,7 +134,7 @@ test("Invoke token count usage_metadata", async () => {
   if (!res?.usage_metadata) {
     return;
   }
-  expect(res.usage_metadata.input_tokens).toBe(71);
+  expect(res.usage_metadata.input_tokens).toBeGreaterThan(1);
   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

libs/langchain-google-genai/src/tests/chat_models.int.test.ts

@@ -445,6 +445,7 @@ test("ChatGoogleGenerativeAI can call withStructuredOutput genai tools and invok
 test("Stream token count usage_metadata", async () => {
   const model = new ChatGoogleGenerativeAI({
     temperature: 0,
+    maxOutputTokens: 10,
   });
   let res: AIMessageChunk | null = null;
   for await (const chunk of await model.stream(
@@ -461,8 +462,8 @@ test("Stream token count usage_metadata", async () => {
   if (!res?.usage_metadata) {
     return;
   }
-  expect(res.usage_metadata.input_tokens).toBe(10);
-  expect(res.usage_metadata.output_tokens).toBeGreaterThan(10);
+  expect(res.usage_metadata.input_tokens).toBeGreaterThan(1);
+  expect(res.usage_metadata.output_tokens).toBeGreaterThan(1);
   expect(res.usage_metadata.total_tokens).toBe(
     res.usage_metadata.input_tokens + res.usage_metadata.output_tokens
   );
@@ -490,15 +491,16 @@ test("streamUsage excludes token usage", async () => {
 test("Invoke token count usage_metadata", async () => {
   const model = new ChatGoogleGenerativeAI({
     temperature: 0,
+    maxOutputTokens: 10,
   });
   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(10);
-  expect(res.usage_metadata.output_tokens).toBeGreaterThan(10);
+  expect(res.usage_metadata.input_tokens).toBeGreaterThan(1);
+  expect(res.usage_metadata.output_tokens).toBeGreaterThan(1);
   expect(res.usage_metadata.total_tokens).toBe(
     res.usage_metadata.input_tokens + res.usage_metadata.output_tokens
   );

libs/langchain-google-vertexai/src/tests/chat_models.int.test.ts

@@ -237,6 +237,7 @@ describe("GAuth Chat", () => {
 test("Stream token count usage_metadata", async () => {
   const model = new ChatVertexAI({
     temperature: 0,
+    maxOutputTokens: 10,
   });
   let res: AIMessageChunk | null = null;
   for await (const chunk of await model.stream(
@@ -253,8 +254,8 @@ test("Stream token count usage_metadata", async () => {
   if (!res?.usage_metadata) {
     return;
   }
-  expect(res.usage_metadata.input_tokens).toBe(9);
-  expect(res.usage_metadata.output_tokens).toBeGreaterThan(10);
+  expect(res.usage_metadata.input_tokens).toBeGreaterThan(1);
+  expect(res.usage_metadata.output_tokens).toBeGreaterThan(1);
   expect(res.usage_metadata.total_tokens).toBe(
     res.usage_metadata.input_tokens + res.usage_metadata.output_tokens
   );
@@ -282,15 +283,16 @@ test("streamUsage excludes token usage", async () => {
 test("Invoke token count usage_metadata", async () => {
   const model = new ChatVertexAI({
     temperature: 0,
+    maxOutputTokens: 10,
   });
   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(9);
-  expect(res.usage_metadata.output_tokens).toBeGreaterThan(10);
+  expect(res.usage_metadata.input_tokens).toBeGreaterThan(1);
+  expect(res.usage_metadata.output_tokens).toBeGreaterThan(1);
   expect(res.usage_metadata.total_tokens).toBe(
     res.usage_metadata.input_tokens + res.usage_metadata.output_tokens
   );

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

@@ -922,6 +922,7 @@ test("Stream token count usage_metadata", async () => {
   const model = new ChatMistralAI({
     model: "codestral-latest",
     temperature: 0,
+    maxTokens: 10,
   });
   let res: AIMessageChunk | null = null;
   for await (const chunk of await model.stream(
@@ -938,8 +939,8 @@ test("Stream token count usage_metadata", async () => {
   if (!res?.usage_metadata) {
     return;
   }
-  expect(res.usage_metadata.input_tokens).toBe(13);
-  expect(res.usage_metadata.output_tokens).toBeGreaterThan(10);
+  expect(res.usage_metadata.input_tokens).toBeGreaterThan(1);
+  expect(res.usage_metadata.output_tokens).toBeGreaterThan(1);
   expect(res.usage_metadata.total_tokens).toBe(
     res.usage_metadata.input_tokens + res.usage_metadata.output_tokens
   );
@@ -969,15 +970,16 @@ test("Invoke token count usage_metadata", async () => {
   const model = new ChatMistralAI({
     model: "codestral-latest",
     temperature: 0,
+    maxTokens: 10,
   });
   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(13);
-  expect(res.usage_metadata.output_tokens).toBeGreaterThan(10);
+  expect(res.usage_metadata.input_tokens).toBeGreaterThan(1);
+  expect(res.usage_metadata.output_tokens).toBeGreaterThan(1);
   expect(res.usage_metadata.total_tokens).toBe(
     res.usage_metadata.input_tokens + res.usage_metadata.output_tokens
   );