bugfixes in google Generative AI chat_model (#3657)

* adding candidateCount + fixing stopSequences not being set in the request * Update chat_models.ts * Small Google fixes * Fix test * Format --------- Co-authored-by: Jacob Lee <[email protected]>
langchain-ai · Dec 14, 2023 · 06fbae0 · 06fbae0 · vercel · Dec 14, 2023
1 parent f601bef
commit 06fbae0
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 1 deletion.
diff --git a/libs/langchain-google-genai/src/chat_models.ts b/libs/langchain-google-genai/src/chat_models.ts
@@ -198,6 +198,8 @@ export class ChatGoogleGenerativeAI
       throw new Error("`topK` must be a positive integer");
     }
 
+    this.stopSequences = fields?.stopSequences ?? this.stopSequences;
+
     this.apiKey = fields?.apiKey ?? getEnvironmentVariable("GOOGLE_API_KEY");
     if (!this.apiKey) {
       throw new Error(

diff --git a/libs/langchain-google-genai/src/tests/chat_models.int.test.ts b/libs/langchain-google-genai/src/tests/chat_models.int.test.ts
@@ -21,6 +21,20 @@ test("Test Google AI generation", async () => {
   expect(res).toBeTruthy();
 });
 
+test("Test Google AI generation with a stop sequence", async () => {
+  const model = new ChatGoogleGenerativeAI({
+    stopSequences: ["two", "2"],
+  });
+  const res = await model.invoke([
+    ["human", `What are the first three positive whole numbers?`],
+  ]);
+  console.log(JSON.stringify(res, null, 2));
+  expect(res).toBeTruthy();
+  expect(res.additional_kwargs.finishReason).toBe("STOP");
+  expect(res.content).not.toContain("2");
+  expect(res.content).not.toContain("two");
+});
+
 test("Test Google AI generation with a system message", async () => {
   const model = new ChatGoogleGenerativeAI({});
   const res = await model.generate([

diff --git a/libs/langchain-google-genai/src/utils.ts b/libs/langchain-google-genai/src/utils.ts
@@ -177,7 +177,7 @@ export function mapGenerateContentResultToChatResult(
     message: new AIMessage({
       content: text,
       name: content === null ? undefined : content.role,
-      additional_kwargs: {},
+      additional_kwargs: generationInfo,
     }),
     generationInfo,
   };
@@ -202,6 +202,8 @@ export function convertResponseContentToChatGenerationChunk(
     message: new AIMessageChunk({
       content: text,
       name: content === null ? undefined : content.role,
+      // Each chunk can have unique "generationInfo", and merging strategy is unclear,
+      // so leave blank for now.
       additional_kwargs: {},
     }),
     generationInfo,