google-genai[patch]: Hookup callbacks to stream & generate methods (#…

…3708)

* google-genai[patch]: Hookup callbacks to stream & generate methods

* chore: lint files

libs/langchain-google-genai/src/chat_models.ts

@@ -247,7 +247,7 @@ export class ChatGoogleGenerativeAI
   async _generate(
     messages: BaseMessage[],
     options: this["ParsedCallOptions"],
-    _runManager?: CallbackManagerForLLMRun
+    runManager?: CallbackManagerForLLMRun
   ): Promise<ChatResult> {
     const prompt = convertBaseMessagesToContent(
       messages,
@@ -272,14 +272,17 @@ export class ChatGoogleGenerativeAI
         return output;
       }
     );
-
-    return mapGenerateContentResultToChatResult(res.response);
+    const generationResult = mapGenerateContentResultToChatResult(res.response);
+    await runManager?.handleLLMNewToken(
+      generationResult.generations[0].text ?? ""
+    );
+    return generationResult;
   }
 
   async *_streamResponseChunks(
     messages: BaseMessage[],
     options: this["ParsedCallOptions"],
-    _runManager?: CallbackManagerForLLMRun
+    runManager?: CallbackManagerForLLMRun
   ): AsyncGenerator<ChatGenerationChunk> {
     const prompt = convertBaseMessagesToContent(
       messages,
@@ -302,6 +305,7 @@ export class ChatGoogleGenerativeAI
       }
 
       yield chunk;
+      await runManager?.handleLLMNewToken(chunk.text ?? "");
     }
   }
 }

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

@@ -73,3 +73,43 @@ test("Test Google AI multimodal generation", async () => {
   console.log(JSON.stringify(res, null, 2));
   expect(res).toBeTruthy();
 });
+
+test("Test Google AI handleLLMNewToken callback", async () => {
+  const model = new ChatGoogleGenerativeAI({});
+  let tokens = "";
+  const res = await model.call(
+    [new HumanMessage("what is 1 + 1?")],
+    undefined,
+    [
+      {
+        handleLLMNewToken(token: string) {
+          tokens += token;
+        },
+      },
+    ]
+  );
+  console.log({ tokens });
+  const responseContent = typeof res.content === "string" ? res.content : "";
+  expect(tokens).toBe(responseContent);
+});
+
+test("Test Google AI handleLLMNewToken callback with streaming", async () => {
+  const model = new ChatGoogleGenerativeAI({});
+  let tokens = "";
+  const res = await model.stream([new HumanMessage("what is 1 + 1?")], {
+    callbacks: [
+      {
+        handleLLMNewToken(token: string) {
+          tokens += token;
+        },
+      },
+    ],
+  });
+  console.log({ tokens });
+  let responseContent = "";
+  for await (const streamItem of res) {
+    responseContent += streamItem.content;
+  }
+  console.log({ tokens });
+  expect(tokens).toBe(responseContent);
+});