✨ Added ability to modify Safety Settings

Author: EvanZhouDev

src/index.ts

@@ -1,4 +1,4 @@
-import { Command } from "./types";
+import { Command, HarmCategory, SafetyThreshold } from "./types";
 
 import type {
 	ChatAskOptions,
@@ -15,7 +15,7 @@ import type {
 	QueryResponseMap,
 } from "./types";
 
-import { getFileType, handleReader, pairToMessage } from "./utils";
+import { SafetyError, getFileType, handleReader, pairToMessage } from "./utils";
 
 const uploadFile = async ({
 	file,
@@ -162,6 +162,7 @@ class Gemini {
 
 	static TEXT = "text" as const;
 	static JSON = "json" as const;
+	static SafetyThreshold = SafetyThreshold;
 
 	constructor(key: string, options: Partial<GeminiOptions> = {}) {
 		if (!options.fetch && typeof fetch !== "function") {
@@ -281,7 +282,7 @@ class Gemini {
 		<F extends Format>(format: F = Gemini.TEXT as F) =>
 		(response: GeminiResponse): FormatType<F> => {
 			if (response.candidates[0].finishReason === "SAFETY") {
-				throw new Error(
+				throw new SafetyError(
 					`Your prompt was blocked by Google. Here are the Harm Categories: \n${JSON.stringify(
 						response.candidates[0].safetyRatings,
 						null,
@@ -337,10 +338,35 @@ class Gemini {
 				maxOutputTokens: 2048,
 				data: [],
 				messages: [],
+				safetySettings: {
+					hate: Gemini.SafetyThreshold.BLOCK_SOME,
+					sexual: Gemini.SafetyThreshold.BLOCK_SOME,
+					harassment: Gemini.SafetyThreshold.BLOCK_SOME,
+					dangerous: Gemini.SafetyThreshold.BLOCK_SOME,
+				},
 			},
 			...options,
 		};
 
+		const safety_settings = [
+			{
+				category: HarmCategory.HateSpeech,
+				threshold: parsedOptions.safetySettings.hate,
+			},
+			{
+				category: HarmCategory.SexuallyExplicit,
+				threshold: parsedOptions.safetySettings.sexual,
+			},
+			{
+				category: HarmCategory.Harassment,
+				threshold: parsedOptions.safetySettings.harassment,
+			},
+			{
+				category: HarmCategory.DangerousContent,
+				threshold: parsedOptions.safetySettings.dangerous,
+			},
+		];
+
 		const command = parsedOptions.stream
 			? Command.StreamGenerate
 			: Command.Generate;
@@ -378,6 +404,7 @@ class Gemini {
 				topP: parsedOptions.topP,
 				topK: parsedOptions.topK,
 			},
+			safety_settings,
 		};
 
 		const response: Response = await this.query(
@@ -440,6 +467,12 @@ class Chat {
 			...{
 				data: [],
 				format: Gemini.TEXT as F,
+				safetySettings: {
+					hate: Gemini.SafetyThreshold.BLOCK_SOME,
+					sexual: Gemini.SafetyThreshold.BLOCK_SOME,
+					harassment: Gemini.SafetyThreshold.BLOCK_SOME,
+					dangerous: Gemini.SafetyThreshold.BLOCK_SOME,
+				},
 			},
 			...options,
 		};

src/types.ts

@@ -60,6 +60,13 @@ export enum Command {
 	Count = "countTokens",
 }
 
+export enum HarmCategory {
+	HateSpeech = "HARM_CATEGORY_HATE_SPEECH",
+	SexuallyExplicit = "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+	Harassment = "HARM_CATEGORY_HARASSMENT",
+	DangerousContent = "HARM_CATEGORY_DANGEROUS_CONTENT",
+}
+
 /**
  * The body used for the API call to generateContent or streamGenerateContent
  */
@@ -71,6 +78,7 @@ type GenerateContentBody = {
 		topP: number;
 		topK: number;
 	};
+	safety_settings: { category: HarmCategory; threshold: SafetyThreshold }[];
 };
 
 /**
@@ -143,6 +151,12 @@ export type CommandOptionMap<F extends Format = TextFormat> = {
 		maxOutputTokens: number;
 		model: string;
 		data: Buffer[];
+		safetySettings: {
+			hate: SafetyThreshold;
+			sexual: SafetyThreshold;
+			harassment: SafetyThreshold;
+			dangerous: SafetyThreshold;
+		};
 		messages: ([string, string] | Message)[];
 		stream?(stream: CommandResponseMap<F>[Command.StreamGenerate]): void;
 	};
@@ -154,6 +168,17 @@ export type CommandOptionMap<F extends Format = TextFormat> = {
 	};
 };
 
+export enum SafetyThreshold {
+	// Content with NEGLIGIBLE will be allowed.
+	BLOCK_MOST = "BLOCK_LOW_AND_ABOVE",
+	// Content with NEGLIGIBLE and LOW will be allowed.
+	BLOCK_SOME = "BLOCK_MEDIUM_AND_ABOVE",
+	// Content with NEGLIGIBLE, LOW, and MEDIUM will be allowed.
+	BLOCK_FEW = "BLOCK_ONLY_HIGH",
+	// All content will be allowed.
+	BLOCK_NONE = "BLOCK_NONE",
+}
+
 export type FormatType<T> = T extends JSONFormat ? GeminiResponse : string;
 
 export type ChatOptions = {