feat(community): Introduce callbacks to IBM Watsonx SDK with CallbackHandlers

libs/langchain-community/package.json

@@ -78,7 +78,7 @@
     "@google-cloud/storage": "^7.7.0",
     "@gradientai/nodejs-sdk": "^1.2.0",
     "@huggingface/inference": "^2.6.4",
-    "@ibm-cloud/watsonx-ai": "^1.1.0",
+    "@ibm-cloud/watsonx-ai": "^1.3.0",
     "@jest/globals": "^29.5.0",
     "@lancedb/lancedb": "^0.13.0",
     "@langchain/core": "workspace:*",

libs/langchain-community/src/chat_models/ibm.ts

@@ -18,7 +18,6 @@ import {
   FunctionDefinition,
   StructuredOutputMethodOptions,
 } from "@langchain/core/language_models/base";
-import { CallbackManagerForLLMRun } from "@langchain/core/callbacks/manager";
 import {
   BaseChatModel,
   BaseChatModelCallOptions,
@@ -68,6 +67,7 @@ import { WatsonxAuth, WatsonxParams } from "../types/ibm.js";
 import {
   _convertToolCallIdToMistralCompatible,
   authenticateAndSetInstance,
+  WatsonxCallbackManagerForLLMRun,
   WatsonxToolsOutputParser,
 } from "../utils/ibm.js";
 
@@ -330,7 +330,6 @@ function _convertToolChoiceToWatsonxToolChoice(
       `Unrecognized tool_choice type. Expected string or TextChatParameterTools. Recieved ${toolChoice}`
     );
 }
-
 export class ChatWatsonx<
     CallOptions extends WatsonxCallOptionsChat = WatsonxCallOptionsChat
   >
@@ -541,7 +540,7 @@ export class ChatWatsonx<
   async _generate(
     messages: BaseMessage[],
     options: this["ParsedCallOptions"],
-    runManager?: CallbackManagerForLLMRun
+    runManager?: WatsonxCallbackManagerForLLMRun
   ): Promise<ChatResult> {
     if (this.streaming) {
       const stream = this._streamResponseChunks(messages, options, runManager);
@@ -631,19 +630,25 @@ export class ChatWatsonx<
   async *_streamResponseChunks(
     messages: BaseMessage[],
     options: this["ParsedCallOptions"],
-    _runManager?: CallbackManagerForLLMRun
+    _runManager?: WatsonxCallbackManagerForLLMRun
   ): AsyncGenerator<ChatGenerationChunk> {
     const params = { ...this.invocationParams(options), ...this.scopeId() };
     const watsonxMessages = _convertMessagesToWatsonxMessages(
       messages,
       this.model
     );
+    const watsonxHandlers = _runManager?.handlers.filter(
+      (item) => item.name === "watsonxHandler"
+    )?.[0]?.watsonxHandlers;
     const callback = () =>
-      this.service.textChatStream({
-        ...params,
-        messages: watsonxMessages,
-        returnObject: true,
-      });
+      this.service.textChatStream(
+        {
+          ...params,
+          messages: watsonxMessages,
+          returnObject: true,
+        },
+        watsonxHandlers
+      );
     const stream = await this.completionWithRetry(callback, options);
     let defaultRole;
     let usage: TextChatUsage | undefined;

libs/langchain-community/src/chat_models/tests/ibm.int.test.ts

@@ -13,6 +13,7 @@ import { ChatPromptTemplate } from "@langchain/core/prompts";
 import { tool } from "@langchain/core/tools";
 import { NewTokenIndices } from "@langchain/core/callbacks/base";
 import { ChatWatsonx } from "../ibm.js";
+import { WatsonxCallbackManager } from "../../utils/ibm.js";
 
 describe("Tests for chat", () => {
   describe("Test ChatWatsonx invoke and generate", () => {
@@ -517,7 +518,9 @@ describe("Tests for chat", () => {
         }
       );
       const llmWithTools = service.bindTools([calculatorTool]);
-      const res = await llmWithTools.invoke("What is 3 * 12");
+      const res = await llmWithTools.invoke(
+        "You are bad at calculations and need to use calculator at all times. What is 3 * 12"
+      );
 
       expect(res).toBeInstanceOf(AIMessage);
       expect(res.tool_calls?.[0].name).toBe("calculator");
@@ -831,4 +834,108 @@ describe("Tests for chat", () => {
       expect(typeof result.number2).toBe("number");
     });
   });
+
+  describe("Test watsonx callbacks", () => {
+    test("Single request callback", async () => {
+      let callbackFlag = false;
+      const service = new ChatWatsonx({
+        model: "mistralai/mistral-large",
+        version: "2024-05-31",
+        serviceUrl: process.env.WATSONX_AI_SERVICE_URL ?? "testString",
+        projectId: process.env.WATSONX_AI_PROJECT_ID ?? "testString",
+        maxTokens: 10,
+        callbacks: WatsonxCallbackManager.fromHandlers({
+          requestCallback(req) {
+            callbackFlag = !!req;
+          },
+        }),
+      });
+      const hello = await service.stream("Print hello world");
+      const chunks = [];
+      for await (const chunk of hello) {
+        chunks.push(chunk);
+      }
+      expect(callbackFlag).toBe(true);
+    });
+    test("Single response callback", async () => {
+      let callbackFlag = false;
+      const service = new ChatWatsonx({
+        model: "mistralai/mistral-large",
+        version: "2024-05-31",
+        serviceUrl: process.env.WATSONX_AI_SERVICE_URL ?? "testString",
+        projectId: process.env.WATSONX_AI_PROJECT_ID ?? "testString",
+        maxTokens: 10,
+        callbacks: WatsonxCallbackManager.fromHandlers({
+          responseCallback(res) {
+            callbackFlag = !!res;
+          },
+        }),
+      });
+      const hello = await service.stream("Print hello world");
+      const chunks = [];
+      for await (const chunk of hello) {
+        chunks.push(chunk);
+      }
+      expect(callbackFlag).toBe(true);
+    });
+    test("Both callbacks", async () => {
+      let callbackFlagReq = false;
+      let callbackFlagRes = false;
+      const service = new ChatWatsonx({
+        model: "mistralai/mistral-large",
+        version: "2024-05-31",
+        serviceUrl: process.env.WATSONX_AI_SERVICE_URL ?? "testString",
+        projectId: process.env.WATSONX_AI_PROJECT_ID ?? "testString",
+        maxTokens: 10,
+        callbacks: WatsonxCallbackManager.fromHandlers({
+          requestCallback(req) {
+            callbackFlagReq = !!req;
+          },
+          responseCallback(res) {
+            callbackFlagRes = !!res;
+          },
+        }),
+      });
+      const hello = await service.stream("Print hello world");
+      const chunks = [];
+      for await (const chunk of hello) {
+        chunks.push(chunk);
+      }
+      expect(callbackFlagReq).toBe(true);
+      expect(callbackFlagRes).toBe(true);
+    });
+    test("Multiple callbacks", async () => {
+      let callbackFlagReq = false;
+      let callbackFlagRes = false;
+      let langchainCallback = false;
+
+      const service = new ChatWatsonx({
+        model: "mistralai/mistral-large",
+        version: "2024-05-31",
+        serviceUrl: process.env.WATSONX_AI_SERVICE_URL ?? "testString",
+        projectId: process.env.WATSONX_AI_PROJECT_ID ?? "testString",
+        maxTokens: 10,
+        callbacks: WatsonxCallbackManager.fromHandlers({
+          requestCallback(req) {
+            callbackFlagReq = !!req;
+          },
+          responseCallback(res) {
+            callbackFlagRes = !!res;
+          },
+          async handleLLMEnd(output) {
+            expect(output.generations).toBeDefined();
+            langchainCallback = !!output;
+          },
+        }),
+      });
+      const hello = await service.stream("Print hello world");
+      const chunks = [];
+      for await (const chunk of hello) {
+        chunks.push(chunk);
+      }
+      expect(callbackFlagReq).toBe(true);
+      expect(callbackFlagRes).toBe(true);
+      expect(langchainCallback).toBe(true);
+    });
+  });
 });

libs/langchain-community/src/document_compressors/ibm.ts

@@ -159,7 +159,7 @@ export class WatsonxRerank
         ? {
             index: document.index,
             relevanceScore: document.score,
-            input: document?.input,
+            input: document?.input.text,
           }
         : {
             index: document.index,