backend/node_modules/@google/generative-ai/test-integration/node/generate-content.test.ts

/**
 * @license
 * Copyright 2024 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { expect, use } from "chai";
import * as chaiAsPromised from "chai-as-promised";
import {
  GoogleGenerativeAI,
  HarmBlockThreshold,
  HarmCategory,
  Part,
} from "../..";

use(chaiAsPromised);

/**
 * Integration tests against live backend.
 */

describe("generateContent", function () {
  this.timeout(60e3);
  this.slow(10e3);
  // This test can be flaky
  // eslint-disable-next-line no-restricted-properties
  it.skip("streaming - count numbers", async () => {
    const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
    const model = genAI.getGenerativeModel({
      model: "gemini-1.5-flash-latest",
      generationConfig: {
        temperature: 0,
        candidateCount: 1,
      },
    });
    const result = await model.generateContentStream({
      contents: [
        {
          role: "user",
          parts: [
            {
              text: "Count from 1 to 10, put each number into square brackets and on a separate line",
            },
          ],
        },
      ],
    });
    const finalResponse = await result.response;
    expect(finalResponse.candidates.length).to.be.equal(1);
    const text = finalResponse.text();
    expect(text).to.include("[1]");
    expect(text).to.include("[10]");
  });
  it("stream true, invalid argument", async () => {
    const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
    const model = genAI.getGenerativeModel({
      model: "gemini-1.5-flash-latest",
      safetySettings: [
        {
          category: HarmCategory.HARM_CATEGORY_HARASSMENT,
          threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH,
        },
      ],
    });
    await expect(
      model.generateContentStream({
        contents: [
          {
            role: "user",
            parts: [{ inlineData: "This is not an image" } as unknown as Part],
          },
        ],
      }),
    ).to.be.rejectedWith("Invalid value");
  });
  it("non-streaming, simple interface", async () => {
    const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
    const model = genAI.getGenerativeModel({
      model: "gemini-1.5-flash-latest",
      safetySettings: [
        {
          category: HarmCategory.HARM_CATEGORY_HARASSMENT,
          threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH,
        },
      ],
    });
    const result = await model.generateContent("What do cats eat?");
    const response = result.response;
    expect(response.text()).to.not.be.empty;
  });
  it("non-streaming, simple interface, custom API version", async () => {
    const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
    const model = genAI.getGenerativeModel(
      {
        model: "gemini-1.5-flash-latest",
        safetySettings: [
          {
            category: HarmCategory.HARM_CATEGORY_HARASSMENT,
            threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH,
          },
        ],
      },
      { apiVersion: "v1beta" },
    );
    const result = await model.generateContent("What do cats eat?");
    const response = result.response;
    expect(response.text()).to.not.be.empty;
  });
});