Merge pull request #612 from mkbehr/genai-migrate

Migrate to new Gemini client SDK.

Author: vivtsai

functions/package-lock.json

@@ -17,7 +17,7 @@
   "main": "lib/index.js",
   "dependencies": {
     "@deliberation-lab/utils": "file:../utils",
-    "@google/generative-ai": "^0.17.2",
+    "@google/genai": "^1.11.0",
     "@sinclair/typebox": "^0.32.30",
     "firebase-admin": "^12.1.0",
     "firebase-functions": "^5.0.0",

functions/package.json

@@ -4,6 +4,7 @@ import nock = require('nock');
 import {
   ModelGenerationConfig,
   ModelResponse,
+  ModelResponseStatus,
   StructuredOutputType,
   StructuredOutputDataType,
 } from '@deliberation-lab/utils';
@@ -152,5 +153,30 @@ describe('Gemini API', () => {
     expect(parsedResponse).toMatchObject(expectedResponse);
   });
 
-  // TODO(mkbehr): Add tests for error responses.
+  it('handles a 503 error from the server', async () => {
+    nock.cleanAll();
+    nock('https://generativelanguage.googleapis.com')
+      .post(`/v1beta/models/${MODEL_NAME}:generateContent`)
+      .reply(503, 'Service Unavailable');
+
+    const generationConfig: ModelGenerationConfig = {
+      maxTokens: 300,
+      stopSequences: [],
+      temperature: 0.4,
+      topP: 0.9,
+      frequencyPenalty: 0,
+      presencePenalty: 0,
+      customRequestBodyFields: [{ name: 'seed', value: 123 }],
+    };
+
+    const response: ModelResponse = await getGeminiAPIResponse(
+      'testapikey',
+      MODEL_NAME,
+      'This is a test prompt.',
+      generationConfig,
+    );
+
+    expect(response.status).toBe(ModelResponseStatus.PROVIDER_UNAVAILABLE_ERROR);
+    expect(response.errorMessage).toContain('Service Unavailable');
+  });
 });

functions/src/api/gemini.api.test.ts

@@ -1,9 +1,10 @@
 import {
-  GoogleGenerativeAI,
+  ApiError,
+  GoogleGenAI,
   GenerationConfig,
   HarmCategory,
   HarmBlockThreshold,
-} from '@google/generative-ai';
+} from '@google/genai';
 import {
   ModelGenerationConfig,
   StructuredOutputType,
@@ -117,15 +118,18 @@ export async function callGemini(
   modelName = GEMINI_DEFAULT_MODEL,
   parseResponse = false, // parse if structured output
 ) {
-  const genAI = new GoogleGenerativeAI(apiKey);
-  const model = genAI.getGenerativeModel({
+  const genAI = new GoogleGenAI({apiKey});
+
+  const response = await genAI.models.generateContent({
     model: modelName,
-    generationConfig,
-    safetySettings: SAFETY_SETTINGS,
+    contents: [{role: 'user', parts: [{text: prompt}]}],
+    config: {
+      ...generationConfig,
+      safetySettings: SAFETY_SETTINGS,
+    },
   });
 
-  const result = await model.generateContent(prompt);
-  const response = await result.response;
+  console.log(`DEBUG: ${JSON.stringify(response)}`);
 
   if (response.promptFeedback) {
     return {
@@ -236,20 +240,20 @@ export async function getGeminiAPIResponse(
     );
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
   } catch (error: any) {
-    // The GenerativeAI client doesn't return responses in a parseable format,
-    // so try to parse the output string looking for the HTTP status code.
+    if (!(error instanceof ApiError)) {
+      return {
+        status: ModelResponseStatus.UNKNOWN_ERROR,
+        generationConfig: geminiConfig,
+        errorMessage: JSON.stringify(error),
+      };
+    }
     let returnStatus = ModelResponseStatus.UNKNOWN_ERROR;
-    // Match a status code and message between brackets, e.g. "[403 Forbidden]".
-    const statusMatch = error.message.match(/\[(\d{3})[\s\w]*\]/);
-    if (statusMatch) {
-      const statusCode = parseInt(statusMatch[1]);
-      if (statusCode == AUTHENTICATION_FAILURE_ERROR_CODE) {
-        returnStatus = ModelResponseStatus.AUTHENTICATION_ERROR;
-      } else if (statusCode == QUOTA_ERROR_CODE) {
-        returnStatus = ModelResponseStatus.QUOTA_ERROR;
-      } else if (statusCode >= 500 && statusCode < 600) {
-        returnStatus = ModelResponseStatus.PROVIDER_UNAVAILABLE_ERROR;
-      }
+    if (error.status == AUTHENTICATION_FAILURE_ERROR_CODE) {
+      returnStatus = ModelResponseStatus.AUTHENTICATION_ERROR;
+    } else if (error.status == QUOTA_ERROR_CODE) {
+      returnStatus = ModelResponseStatus.QUOTA_ERROR;
+    } else if (error.status >= 500 && error.status < 600) {
+      returnStatus = ModelResponseStatus.PROVIDER_UNAVAILABLE_ERROR;
     }
     return {
       status: returnStatus,