ak-gemini 1.0.8 → 1.0.10

package/index.cjs

@@ -29,6 +29,8 @@ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: tru
 // index.js
 var index_exports = {};
 __export(index_exports, {
+  HarmBlockThreshold: () => import_genai.HarmBlockThreshold,
+  HarmCategory: () => import_genai.HarmCategory,
   ThinkingLevel: () => import_genai.ThinkingLevel,
   attemptJSONRecovery: () => attemptJSONRecovery,
   default: () => index_default,
@@ -132,6 +134,7 @@ var AITransformer = class {
     this.transformWithValidation = prepareAndValidateMessage.bind(this);
     this.estimate = estimateTokenUsage.bind(this);
     this.estimateTokenUsage = estimateTokenUsage.bind(this);
+    this.updateSystemInstructions = updateSystemInstructions.bind(this);
   }
 };
 var index_default = AITransformer;
@@ -221,6 +224,8 @@ function AITransformFactory(options = {}) {
   this.retryDelay = options.retryDelay || 1e3;
   this.asyncValidator = options.asyncValidator || null;
   this.onlyJSON = options.onlyJSON !== void 0 ? options.onlyJSON : true;
+  this.enableGrounding = options.enableGrounding || false;
+  this.groundingConfig = options.groundingConfig || {};
   if (this.promptKey === this.answerKey) {
     throw new Error("Source and target keys cannot be the same. Please provide distinct keys.");
   }
@@ -228,6 +233,8 @@ function AITransformFactory(options = {}) {
     logger_default.debug(`Creating AI Transformer with model: ${this.modelName}`);
     logger_default.debug(`Using keys - Source: "${this.promptKey}", Target: "${this.answerKey}", Context: "${this.contextKey}"`);
     logger_default.debug(`Max output tokens set to: ${this.chatConfig.maxOutputTokens}`);
+    logger_default.debug(`Using API key: ${this.apiKey.substring(0, 10)}...`);
+    logger_default.debug(`Grounding ${this.enableGrounding ? "ENABLED" : "DISABLED"} (costs $35/1k queries)`);
   }
   const ai = new import_genai.GoogleGenAI({ apiKey: this.apiKey });
   this.genAIClient = ai;
@@ -236,12 +243,19 @@ function AITransformFactory(options = {}) {
 async function initChat(force = false) {
   if (this.chat && !force) return;
   logger_default.debug(`Initializing Gemini chat session with model: ${this.modelName}...`);
-  this.chat = await this.genAIClient.chats.create({
+  const chatOptions = {
     model: this.modelName,
     // @ts-ignore
     config: this.chatConfig,
     history: []
-  });
+  };
+  if (this.enableGrounding) {
+    chatOptions.config.tools = [{
+      googleSearch: this.groundingConfig
+    }];
+    logger_default.debug(`Search grounding ENABLED for this session (WARNING: costs $35/1k queries)`);
+  }
+  this.chat = await this.genAIClient.chats.create(chatOptions);
   try {
     await this.genAIClient.models.list();
     logger_default.debug("Gemini API connection successful.");
@@ -345,6 +359,32 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
   }
   const maxRetries = options.maxRetries ?? this.maxRetries;
   const retryDelay = options.retryDelay ?? this.retryDelay;
+  const enableGroundingForMessage = options.enableGrounding ?? this.enableGrounding;
+  const groundingConfigForMessage = options.groundingConfig ?? this.groundingConfig;
+  if (enableGroundingForMessage !== this.enableGrounding) {
+    const originalGrounding = this.enableGrounding;
+    const originalConfig = this.groundingConfig;
+    try {
+      this.enableGrounding = enableGroundingForMessage;
+      this.groundingConfig = groundingConfigForMessage;
+      await this.init(true);
+      if (enableGroundingForMessage) {
+        logger_default.warn(`Search grounding ENABLED for this message (WARNING: costs $35/1k queries)`);
+      } else {
+        logger_default.debug(`Search grounding DISABLED for this message`);
+      }
+    } catch (error) {
+      this.enableGrounding = originalGrounding;
+      this.groundingConfig = originalConfig;
+      throw error;
+    }
+    const restoreGrounding = async () => {
+      this.enableGrounding = originalGrounding;
+      this.groundingConfig = originalConfig;
+      await this.init(true);
+    };
+    options._restoreGrounding = restoreGrounding;
+  }
   let lastError = null;
   let lastPayload = null;
   if (sourcePayload && isJSON(sourcePayload)) {
@@ -366,12 +406,18 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
         await validatorFn(transformedPayload);
       }
       logger_default.debug(`Transformation succeeded on attempt ${attempt + 1}`);
+      if (options._restoreGrounding) {
+        await options._restoreGrounding();
+      }
       return transformedPayload;
     } catch (error) {
       lastError = error;
       logger_default.warn(`Attempt ${attempt + 1} failed: ${error.message}`);
       if (attempt >= maxRetries) {
         logger_default.error(`All ${maxRetries + 1} attempts failed.`);
+        if (options._restoreGrounding) {
+          await options._restoreGrounding();
+        }
         throw new Error(`Transformation failed after ${maxRetries + 1} attempts. Last error: ${error.message}`);
       }
       const delay = retryDelay * Math.pow(2, attempt);
@@ -430,12 +476,19 @@ async function estimateTokenUsage(nextPayload) {
 async function resetChat() {
   if (this.chat) {
     logger_default.debug("Resetting Gemini chat session...");
-    this.chat = await this.genAIClient.chats.create({
+    const chatOptions = {
       model: this.modelName,
       // @ts-ignore
       config: this.chatConfig,
       history: []
-    });
+    };
+    if (this.enableGrounding) {
+      chatOptions.config.tools = [{
+        googleSearch: this.groundingConfig
+      }];
+      logger_default.debug(`Search grounding preserved during reset (WARNING: costs $35/1k queries)`);
+    }
+    this.chat = await this.genAIClient.chats.create(chatOptions);
     logger_default.debug("Chat session reset.");
   } else {
     logger_default.warn("Cannot reset chat session: chat not yet initialized.");
@@ -448,6 +501,15 @@ function getChatHistory() {
   }
   return this.chat.getHistory();
 }
+async function updateSystemInstructions(newInstructions) {
+  if (!newInstructions || typeof newInstructions !== "string") {
+    throw new Error("System instructions must be a non-empty string");
+  }
+  this.systemInstructions = newInstructions.trim();
+  this.chatConfig.systemInstruction = this.systemInstructions;
+  logger_default.debug("Updating system instructions and reinitializing chat...");
+  await this.init(true);
+}
 function attemptJSONRecovery(text, maxAttempts = 100) {
   if (!text || typeof text !== "string") return null;
   try {
@@ -748,6 +810,8 @@ if (import_meta.url === new URL(`file://${process.argv[1]}`).href) {
 }
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
+  HarmBlockThreshold,
+  HarmCategory,
   ThinkingLevel,
   attemptJSONRecovery,
   log

package/index.js

@@ -28,7 +28,7 @@ import u from 'ak-tools';
 import path from 'path';
 import log from './logger.js';
 export { log };
-export { ThinkingLevel };
+export { ThinkingLevel, HarmCategory, HarmBlockThreshold };
 
 
 
@@ -134,6 +134,7 @@ class AITransformer {
 		this.transformWithValidation = prepareAndValidateMessage.bind(this);
 		this.estimate = estimateTokenUsage.bind(this);
 		this.estimateTokenUsage = estimateTokenUsage.bind(this);
+		this.updateSystemInstructions = updateSystemInstructions.bind(this);
 	}
 }
 
@@ -265,6 +266,10 @@ function AITransformFactory(options = {}) {
 	//are we forcing json responses only?
 	this.onlyJSON = options.onlyJSON !== undefined ? options.onlyJSON : true; // If true, only return JSON responses
 
+	// Grounding configuration (disabled by default to avoid costs)
+	this.enableGrounding = options.enableGrounding || false;
+	this.groundingConfig = options.groundingConfig || {};
+
 	if (this.promptKey === this.answerKey) {
 		throw new Error("Source and target keys cannot be the same. Please provide distinct keys.");
 	}
@@ -273,6 +278,9 @@ function AITransformFactory(options = {}) {
 		log.debug(`Creating AI Transformer with model: ${this.modelName}`);
 		log.debug(`Using keys - Source: "${this.promptKey}", Target: "${this.answerKey}", Context: "${this.contextKey}"`);
 		log.debug(`Max output tokens set to: ${this.chatConfig.maxOutputTokens}`);
+		// Log API key prefix for tracking (first 10 chars only for security)
+		log.debug(`Using API key: ${this.apiKey.substring(0, 10)}...`);
+		log.debug(`Grounding ${this.enableGrounding ? 'ENABLED' : 'DISABLED'} (costs $35/1k queries)`);
 	}
 
 	const ai = new GoogleGenAI({ apiKey: this.apiKey });
@@ -291,12 +299,23 @@ async function initChat(force = false) {
 
 	log.debug(`Initializing Gemini chat session with model: ${this.modelName}...`);
 
-	this.chat = await this.genAIClient.chats.create({
+	// Add grounding tools if enabled
+	const chatOptions = {
 		model: this.modelName,
 		// @ts-ignore
 		config: this.chatConfig,
 		history: [],
-	});
+	};
+
+	// Only add tools if grounding is explicitly enabled
+	if (this.enableGrounding) {
+		chatOptions.config.tools = [{
+			googleSearch: this.groundingConfig
+		}];
+		log.debug(`Search grounding ENABLED for this session (WARNING: costs $35/1k queries)`);
+	}
+
+	this.chat = await this.genAIClient.chats.create(chatOptions);
 
 	try {
 		await this.genAIClient.models.list();
@@ -463,6 +482,47 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
 	const maxRetries = options.maxRetries ?? this.maxRetries;
 	const retryDelay = options.retryDelay ?? this.retryDelay;
 
+	// Check if grounding should be enabled for this specific message
+	const enableGroundingForMessage = options.enableGrounding ?? this.enableGrounding;
+	const groundingConfigForMessage = options.groundingConfig ?? this.groundingConfig;
+
+	// Reinitialize chat if grounding settings changed for this message
+	if (enableGroundingForMessage !== this.enableGrounding) {
+		const originalGrounding = this.enableGrounding;
+		const originalConfig = this.groundingConfig;
+
+		try {
+			// Temporarily change grounding settings
+			this.enableGrounding = enableGroundingForMessage;
+			this.groundingConfig = groundingConfigForMessage;
+
+			// Force reinit with new settings
+			await this.init(true);
+
+			// Log the change
+			if (enableGroundingForMessage) {
+				log.warn(`Search grounding ENABLED for this message (WARNING: costs $35/1k queries)`);
+			} else {
+				log.debug(`Search grounding DISABLED for this message`);
+			}
+		} catch (error) {
+			// Restore original settings on error
+			this.enableGrounding = originalGrounding;
+			this.groundingConfig = originalConfig;
+			throw error;
+		}
+
+		// Schedule restoration after message completes
+		const restoreGrounding = async () => {
+			this.enableGrounding = originalGrounding;
+			this.groundingConfig = originalConfig;
+			await this.init(true);
+		};
+
+		// Store restoration function to call after message completes
+		options._restoreGrounding = restoreGrounding;
+	}
+
 	let lastError = null;
 	let lastPayload = null; // Store the payload that caused the validation error
 
@@ -498,6 +558,12 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
 
 			// Step 3: Success!
 			log.debug(`Transformation succeeded on attempt ${attempt + 1}`);
+
+			// Restore original grounding settings if they were changed
+			if (options._restoreGrounding) {
+				await options._restoreGrounding();
+			}
+
 			return transformedPayload;
 
 		} catch (error) {
@@ -506,6 +572,12 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
 
 			if (attempt >= maxRetries) {
 				log.error(`All ${maxRetries + 1} attempts failed.`);
+
+				// Restore original grounding settings even on failure
+				if (options._restoreGrounding) {
+					await options._restoreGrounding();
+				}
+
 				throw new Error(`Transformation failed after ${maxRetries + 1} attempts. Last error: ${error.message}`);
 			}
 
@@ -607,12 +679,24 @@ async function estimateTokenUsage(nextPayload) {
 async function resetChat() {
 	if (this.chat) {
 		log.debug("Resetting Gemini chat session...");
-		this.chat = await this.genAIClient.chats.create({
+
+		// Prepare chat options with grounding if enabled
+		const chatOptions = {
 			model: this.modelName,
 			// @ts-ignore
 			config: this.chatConfig,
 			history: [],
-		});
+		};
+
+		// Only add tools if grounding is explicitly enabled
+		if (this.enableGrounding) {
+			chatOptions.config.tools = [{
+				googleSearch: this.groundingConfig
+			}];
+			log.debug(`Search grounding preserved during reset (WARNING: costs $35/1k queries)`);
+		}
+
+		this.chat = await this.genAIClient.chats.create(chatOptions);
 		log.debug("Chat session reset.");
 	} else {
 		log.warn("Cannot reset chat session: chat not yet initialized.");
@@ -631,6 +715,24 @@ function getChatHistory() {
 	return this.chat.getHistory();
 }
 
+/**
+ * Updates system instructions and reinitializes the chat session
+ * @this {ExportedAPI}
+ * @param {string} newInstructions - The new system instructions
+ * @returns {Promise<void>}
+ */
+async function updateSystemInstructions(newInstructions) {
+	if (!newInstructions || typeof newInstructions !== 'string') {
+		throw new Error('System instructions must be a non-empty string');
+	}
+
+	this.systemInstructions = newInstructions.trim();
+	this.chatConfig.systemInstruction = this.systemInstructions;
+
+	log.debug('Updating system instructions and reinitializing chat...');
+	await this.init(true); // Force reinitialize with new instructions
+}
+
 
 /*
 ----

package/package.json

@@ -2,7 +2,7 @@
 	"name": "ak-gemini",
 	"author": "ak@mixpanel.com",
 	"description": "AK's Generative AI Helper for doing... transforms",
-	"version": "1.0.8",
+	"version": "1.0.10",
 	"main": "index.js",
 	"files": [
 		"index.js",

package/types.d.ts

@@ -1,6 +1,6 @@
-import type { GoogleGenAI, ThinkingLevel } from '@google/genai';
+import type { GoogleGenAI, ThinkingLevel, HarmCategory, HarmBlockThreshold } from '@google/genai';
 
-export { ThinkingLevel };
+export { ThinkingLevel, HarmCategory, HarmBlockThreshold };
 
 export interface ThinkingConfig {
   /** Indicates whether to include thoughts in the response. If true, thoughts are returned only if the model supports thought and thoughts are available. */
@@ -12,8 +12,8 @@ export interface ThinkingConfig {
 }
 
 export interface SafetySetting {
-  category: string; // The harm category
-  threshold: string; // The blocking threshold
+  category: HarmCategory; // The harm category
+  threshold: HarmBlockThreshold; // The blocking threshold
 }
 
 export interface ChatConfig {
@@ -51,7 +51,9 @@ export interface AITransformerContext {
   rawMessage?: (payload: Record<string, unknown> | string) => Promise<Record<string, unknown>>; // Function to send raw messages to the model
   genAIClient?: GoogleGenAI; // Google GenAI client instance
   onlyJSON?: boolean; // If true, only JSON responses are allowed
-  
+  enableGrounding?: boolean; // Enable Google Search grounding (default: false, WARNING: costs $35/1k queries)
+  groundingConfig?: Record<string, any>; // Additional grounding configuration options
+
 }
 
 export interface TransformationExample {
@@ -93,6 +95,8 @@ export interface AITransformerOptions {
   onlyJSON?: boolean; // If true, only JSON responses are allowed
   asyncValidator?: AsyncValidatorFunction; // Optional async validator function for response validation
   logLevel?: 'trace' | 'debug' | 'info' | 'warn' | 'error' | 'fatal' | 'none'; // Log level for the logger (defaults to 'info', 'none' disables logging)
+  enableGrounding?: boolean; // Enable Google Search grounding (default: false, WARNING: costs $35/1k queries)
+  groundingConfig?: Record<string, any>; // Additional grounding configuration options
 }
 
 // Async validator function type
@@ -120,6 +124,8 @@ export declare class AITransformer {
   genAIClient: any;
   chat: any;
   logLevel: 'trace' | 'debug' | 'info' | 'warn' | 'error' | 'fatal' | 'none';
+  enableGrounding: boolean;
+  groundingConfig: Record<string, any>;
   
   // Methods
   init(force?: boolean): Promise<void>;
@@ -133,8 +139,8 @@ export declare class AITransformer {
   getHistory(): Array<any>;
   estimateTokenUsage(nextPayload: Record<string, unknown> | string): Promise<{ totalTokens: number; breakdown?: any }>;
   estimate(nextPayload: Record<string, unknown> | string): Promise<{ totalTokens: number; breakdown?: any }>;
+  updateSystemInstructions(newInstructions: string): Promise<void>;
 }
 
 // Default export
-declare const _default: typeof AITransformer;
-export default _default;
+export default AITransformer;