fix: safety return types

PiperOrigin-RevId: 591915953

Author: vertex-sdk-bot

src/index_test.ts

@@ -20,7 +20,19 @@ import 'jasmine';
 
 import {ChatSession, GenerativeModel, StartChatParams, VertexAI} from './index';
 import * as StreamFunctions from './process_stream';
-import {CountTokensRequest, FinishReason, GenerateContentRequest, GenerateContentResponse, GenerateContentResult, HarmBlockThreshold, HarmCategory, StreamGenerateContentResult,} from './types/content';
+import {
+  CountTokensRequest,
+  FinishReason,
+  GenerateContentRequest,
+  GenerateContentResponse,
+  GenerateContentResult,
+  HarmBlockThreshold,
+  HarmCategory,
+  HarmProbability,
+  SafetyRating,
+  SafetySetting,
+  StreamGenerateContentResult,
+} from './types/content';
 import {constants} from './util';
 
 const PROJECT = 'test_project';
@@ -38,8 +50,8 @@ const TEST_USER_CHAT_MESSAGE_WITH_GCS_FILE = [
       {
         file_data: {
           file_uri: 'gs://test_bucket/test_image.jpeg',
-          mime_type: 'image/jpeg'
-        }
+          mime_type: 'image/jpeg',
+        },
       },
     ],
   },
@@ -55,10 +67,17 @@ const TEST_USER_CHAT_MESSAGE_WITH_INVALID_GCS_FILE = [
   },
 ];
 
-const TEST_SAFETY_RATINGS = [
+const TEST_SAFETY_SETTINGS: SafetySetting[] = [
   {
     category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
-    threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+    threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+  },
+];
+
+const TEST_SAFETY_RATINGS: SafetyRating[] = [
+  {
+    category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+    probability: HarmProbability.NEGLIGIBLE,
   },
 ];
 const TEST_GENERATION_CONFIG = {
@@ -76,13 +95,14 @@ const TEST_CANDIDATES = [
     finishMessage: '',
     safetyRatings: TEST_SAFETY_RATINGS,
     citationMetadata: {
-      citationSources: [{
-        startIndex: 367,
-        endIndex: 491,
-        uri:
-            'https://www.numerade.com/ask/question/why-does-the-uncertainty-principle-make-it-impossible-to-predict-a-trajectory-for-the-clectron-95172/'
-      }]
-    }
+      citationSources: [
+        {
+          startIndex: 367,
+          endIndex: 491,
+          uri: 'https://www.numerade.com/ask/question/why-does-the-uncertainty-principle-make-it-impossible-to-predict-a-trajectory-for-the-clectron-95172/',
+        },
+      ],
+    },
   },
 ];
 const TEST_MODEL_RESPONSE = {
@@ -178,8 +198,9 @@ describe('VertexAI', () => {
         response: Promise.resolve(TEST_MODEL_RESPONSE),
         stream: testGenerator(),
       };
-      spyOn(StreamFunctions, 'processStream')
-          .and.returnValue(expectedStreamResult);
+      spyOn(StreamFunctions, 'processStream').and.returnValue(
+        expectedStreamResult
+      );
       const resp = await model.generateContent(req);
       expect(resp).toEqual(expectedResult);
     });
@@ -190,15 +211,17 @@ describe('VertexAI', () => {
       const req: GenerateContentRequest = {
         contents: TEST_USER_CHAT_MESSAGE_WITH_INVALID_GCS_FILE,
       };
-      await expectAsync(model.generateContent(req)).toBeRejectedWithError(URIError);
+      await expectAsync(model.generateContent(req)).toBeRejectedWithError(
+        URIError
+      );
     });
   });
 
   describe('generateContent', () => {
     it('returns a GenerateContentResponse when passed safety_settings and generation_config', async () => {
       const req: GenerateContentRequest = {
         contents: TEST_USER_CHAT_MESSAGE,
-        safety_settings: TEST_SAFETY_RATINGS,
+        safety_settings: TEST_SAFETY_SETTINGS,
         generation_config: TEST_GENERATION_CONFIG,
       };
       const expectedResult: GenerateContentResult = {
@@ -223,7 +246,8 @@ describe('VertexAI', () => {
         location: LOCATION,
         apiEndpoint: TEST_ENDPOINT_BASE_PATH,
       });
-      vertexaiWithBasePath.preview['tokenInternalPromise'] = Promise.resolve(TEST_TOKEN);
+      vertexaiWithBasePath.preview['tokenInternalPromise'] =
+        Promise.resolve(TEST_TOKEN);
       model = vertexaiWithBasePath.preview.getGenerativeModel({
         model: 'gemini-pro',
       });
@@ -255,7 +279,8 @@ describe('VertexAI', () => {
         project: PROJECT,
         location: LOCATION,
       });
-      vertexaiWithoutBasePath.preview['tokenInternalPromise'] = Promise.resolve(TEST_TOKEN);
+      vertexaiWithoutBasePath.preview['tokenInternalPromise'] =
+        Promise.resolve(TEST_TOKEN);
       model = vertexaiWithoutBasePath.preview.getGenerativeModel({
         model: 'gemini-pro',
       });
@@ -275,8 +300,9 @@ describe('VertexAI', () => {
         expectedStreamResult
       );
       await model.generateContent(req);
-      expect(requestSpy.calls.allArgs()[0][0].toString())
-          .toContain(`${LOCATION}-aiplatform.googleapis.com`);
+      expect(requestSpy.calls.allArgs()[0][0].toString()).toContain(
+        `${LOCATION}-aiplatform.googleapis.com`
+      );
     });
   });
 
@@ -295,11 +321,12 @@ describe('VertexAI', () => {
         stream: testGenerator(),
       };
       const requestSpy = spyOn(global, 'fetch');
-      spyOn(StreamFunctions, 'processStream')
-          .and.returnValue(expectedStreamResult);
+      spyOn(StreamFunctions, 'processStream').and.returnValue(
+        expectedStreamResult
+      );
       await model.generateContent(reqWithEmptyConfigs);
       const requestArgs = requestSpy.calls.allArgs()[0][1];
-      if (typeof requestArgs == 'object' && requestArgs) {
+      if (typeof requestArgs === 'object' && requestArgs) {
         expect(JSON.stringify(requestArgs['body'])).not.toContain('top_k');
       }
     });
@@ -320,11 +347,12 @@ describe('VertexAI', () => {
         stream: testGenerator(),
       };
       const requestSpy = spyOn(global, 'fetch');
-      spyOn(StreamFunctions, 'processStream')
-          .and.returnValue(expectedStreamResult);
+      spyOn(StreamFunctions, 'processStream').and.returnValue(
+        expectedStreamResult
+      );
       await model.generateContent(reqWithEmptyConfigs);
       const requestArgs = requestSpy.calls.allArgs()[0][1];
-      if (typeof requestArgs == 'object' && requestArgs) {
+      if (typeof requestArgs === 'object' && requestArgs) {
         expect(JSON.stringify(requestArgs['body'])).toContain('top_k');
       }
     });
@@ -342,14 +370,17 @@ describe('VertexAI', () => {
         response: Promise.resolve(TEST_MODEL_RESPONSE),
         stream: testGenerator(),
       };
-      spyOn(StreamFunctions, 'processStream')
-          .and.returnValue(expectedStreamResult);
+      spyOn(StreamFunctions, 'processStream').and.returnValue(
+        expectedStreamResult
+      );
       const resp = await model.generateContent(req);
       console.log(resp.response.candidates[0].citationMetadata, 'yoyoyo');
       expect(
-          resp.response.candidates[0].citationMetadata?.citationSources.length)
-          .toEqual(TEST_MODEL_RESPONSE.candidates[0]
-                       .citationMetadata.citationSources.length);
+        resp.response.candidates[0].citationMetadata?.citationSources.length
+      ).toEqual(
+        TEST_MODEL_RESPONSE.candidates[0].citationMetadata.citationSources
+          .length
+      );
     });
   });
 

src/types/content.ts

@@ -18,7 +18,7 @@
 /**
  * Params used to initialize the Vertex SDK
  * @param{string} project - the project name of your Google Cloud project. It is not the numeric project ID.
- * @param{string} location - the location of your project. 
+ * @param{string} location - the location of your project.
  * @param{string} apiEndpoint - Optional. If not specified, a default value will be resolved by SDK.
  */
 export declare interface VertexInit {
@@ -51,7 +51,7 @@ export declare interface CountTokensResponse {
 
 /**
  * Configuration for initializing a model, for example via getGenerativeModel
- * @param {string} model - model name. 
+ * @param {string} model - model name.
  * @example "gemini-pro"
  */
 export declare interface ModelParams extends BaseModelParams {
@@ -85,7 +85,9 @@ export declare interface GenerationConfig {
   top_p?: number;
   top_k?: number;
 }
-
+/**
+ * Harm categories that would cause prompts or candidates to be blocked.
+ */
 export enum HarmCategory {
   HARM_CATEGORY_UNSPECIFIED = 'HARM_CATEGORY_UNSPECIFIED',
   HARM_CATEGORY_HATE_SPEECH = 'HARM_CATEGORY_HATE_SPEECH',
@@ -94,6 +96,9 @@ export enum HarmCategory {
   HARM_CATEGORY_SEXUALLY_EXPLICIT = 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
 }
 
+/**
+ * Threshhold above which a prompt or candidate will be blocked.
+ */
 export enum HarmBlockThreshold {
   // Unspecified harm block threshold.
   HARM_BLOCK_THRESHOLD_UNSPECIFIED = 'HARM_BLOCK_THRESHOLD_UNSPECIFIED',
@@ -107,12 +112,28 @@ export enum HarmBlockThreshold {
   BLOCK_NONE = 'BLOCK_NONE',
 }
 
+/**
+ * Probability that a prompt or candidate matches a harm category.
+ */
+export enum HarmProbability {
+  // Probability is unspecified.
+  HARM_PROBABILITY_UNSPECIFIED = 'HARM_PROBABILITY_UNSPECIFIED',
+  // Content has a negligible chance of being unsafe.
+  NEGLIGIBLE = 'NEGLIGIBLE',
+  // Content has a low chance of being unsafe.
+  LOW = 'LOW',
+  // Content has a medium chance of being unsafe.
+  MEDIUM = 'MEDIUM',
+  // Content has a high chance of being unsafe.
+  HIGH = 'HIGH',
+}
+
 /**
  * Safety rating for a piece of content
  */
 export declare interface SafetyRating {
   category: HarmCategory;
-  threshold: HarmBlockThreshold;
+  probability: HarmProbability;
 }
 
 /**