Fix the issue #124 (#290)

google-gemini · Nov 5, 2024 · 6ec2c27 · 6ec2c27
1 parent 927f6f0
commit 6ec2c27
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 2 deletions.
diff --git a/src/methods/chat-session.test.ts b/src/methods/chat-session.test.ts
@@ -22,6 +22,8 @@ import * as chaiAsPromised from "chai-as-promised";
 import * as generateContentMethods from "./generate-content";
 import { GenerateContentStreamResult } from "../../types";
 import { ChatSession } from "./chat-session";
+import { getMockResponse } from "../../test-utils/mock-response";
+import * as request from "../requests/request";
 
 use(sinonChai);
 use(chaiAsPromised);
@@ -45,6 +47,15 @@ describe("ChatSession", () => {
       );
     });
   });
+  describe("sendMessageRecitationErrorNotAddingResponseToHistory()", () => {
+    it("generateContent errors should be catchable", async () => {
+      const mockResponse = getMockResponse("unary-failure-citations.json");
+      stub(request, "makeModelRequest").resolves(mockResponse as Response);
+      const chatSession = new ChatSession("MY_API_KEY", "a-model");
+      await chatSession.sendMessage("hello");
+      expect((await chatSession.getHistory()).length).equals(0);
+    });
+  });
   describe("sendMessageStream()", () => {
     it("generateContentStream errors should be catchable", async () => {
       const clock = useFakeTimers();

diff --git a/src/methods/chat-session.ts b/src/methods/chat-session.ts
@@ -110,7 +110,8 @@ export class ChatSession {
       .then((result) => {
         if (
           result.response.candidates &&
-          result.response.candidates.length > 0
+          result.response.candidates.length > 0 &&
+          result.response.candidates[0]?.content !== undefined
         ) {
           this._history.push(newContent);
           const responseContent: Content = {
@@ -179,7 +180,11 @@ export class ChatSession {
       })
       .then((streamResult) => streamResult.response)
       .then((response) => {
-        if (response.candidates && response.candidates.length > 0) {
+        if (
+          response.candidates &&
+          response.candidates.length > 0 &&
+          response.candidates[0]?.content !== undefined
+        ) {
           this._history.push(newContent);
           const responseContent = { ...response.candidates[0].content };
           // Response seems to come back without a role set.

diff --git a/test-utils/mock-responses/unary-failure-citations.json b/test-utils/mock-responses/unary-failure-citations.json
@@ -0,0 +1,13 @@
+{
+    "candidates": [
+        {
+            "finishReason": "RECITATION",
+            "index": 0
+        }
+    ],
+    "usageMetadata": {
+        "promptTokenCount": 18,
+        "totalTokenCount": 18
+    },
+    "modelVersion": "gemini-1.5-flash-001"
+}