ak-gemini 1.0.7 → 1.0.9

package/index.cjs

@@ -29,7 +29,10 @@ 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,
   log: () => logger_default
 });
@@ -79,6 +82,7 @@ var DEFAULT_THINKING_CONFIG = {
   thinkingBudget: 0,
   thinkingLevel: import_genai.ThinkingLevel.MINIMAL
 };
+var DEFAULT_MAX_OUTPUT_TOKENS = 1e5;
 var THINKING_SUPPORTED_MODELS = [
   /^gemini-3-flash(-preview)?$/,
   /^gemini-3-pro(-preview|-image-preview)?$/,
@@ -166,21 +170,43 @@ function AITransformFactory(options = {}) {
     ...options.chatConfig,
     systemInstruction: this.systemInstructions
   };
+  if (options.maxOutputTokens !== void 0) {
+    if (options.maxOutputTokens === null) {
+      delete this.chatConfig.maxOutputTokens;
+    } else {
+      this.chatConfig.maxOutputTokens = options.maxOutputTokens;
+    }
+  } else if (options.chatConfig?.maxOutputTokens !== void 0) {
+    if (options.chatConfig.maxOutputTokens === null) {
+      delete this.chatConfig.maxOutputTokens;
+    } else {
+      this.chatConfig.maxOutputTokens = options.chatConfig.maxOutputTokens;
+    }
+  } else {
+    this.chatConfig.maxOutputTokens = DEFAULT_MAX_OUTPUT_TOKENS;
+  }
   const modelSupportsThinking = THINKING_SUPPORTED_MODELS.some(
     (pattern) => pattern.test(this.modelName)
   );
-  if (modelSupportsThinking && options.thinkingConfig) {
-    const thinkingConfig = {
-      ...DEFAULT_THINKING_CONFIG,
-      ...options.thinkingConfig
-    };
-    this.chatConfig.thinkingConfig = thinkingConfig;
-    if (logger_default.level !== "silent") {
-      logger_default.debug(`Model ${this.modelName} supports thinking. Applied thinkingConfig:`, thinkingConfig);
-    }
-  } else if (options.thinkingConfig && !modelSupportsThinking) {
-    if (logger_default.level !== "silent") {
-      logger_default.warn(`Model ${this.modelName} does not support thinking features. Ignoring thinkingConfig.`);
+  if (options.thinkingConfig !== void 0) {
+    if (options.thinkingConfig === null) {
+      delete this.chatConfig.thinkingConfig;
+      if (logger_default.level !== "silent") {
+        logger_default.debug(`thinkingConfig set to null - removed from configuration`);
+      }
+    } else if (modelSupportsThinking) {
+      const thinkingConfig = {
+        ...DEFAULT_THINKING_CONFIG,
+        ...options.thinkingConfig
+      };
+      this.chatConfig.thinkingConfig = thinkingConfig;
+      if (logger_default.level !== "silent") {
+        logger_default.debug(`Model ${this.modelName} supports thinking. Applied thinkingConfig:`, thinkingConfig);
+      }
+    } else {
+      if (logger_default.level !== "silent") {
+        logger_default.warn(`Model ${this.modelName} does not support thinking features. Ignoring thinkingConfig.`);
+      }
     }
   }
   if (options.responseSchema) {
@@ -197,12 +223,17 @@ 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.");
   }
   if (logger_default.level !== "silent") {
     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;
@@ -211,12 +242,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.");
@@ -320,6 +358,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)) {
@@ -341,12 +405,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);
@@ -405,12 +475,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.");
@@ -423,6 +500,122 @@ function getChatHistory() {
   }
   return this.chat.getHistory();
 }
+function attemptJSONRecovery(text, maxAttempts = 100) {
+  if (!text || typeof text !== "string") return null;
+  try {
+    return JSON.parse(text);
+  } catch (e) {
+  }
+  let workingText = text.trim();
+  let braces = 0;
+  let brackets = 0;
+  let inString = false;
+  let escapeNext = false;
+  for (let j = 0; j < workingText.length; j++) {
+    const char = workingText[j];
+    if (escapeNext) {
+      escapeNext = false;
+      continue;
+    }
+    if (char === "\\") {
+      escapeNext = true;
+      continue;
+    }
+    if (char === '"') {
+      inString = !inString;
+      continue;
+    }
+    if (!inString) {
+      if (char === "{") braces++;
+      else if (char === "}") braces--;
+      else if (char === "[") brackets++;
+      else if (char === "]") brackets--;
+    }
+  }
+  if ((braces > 0 || brackets > 0 || inString) && workingText.length > 2) {
+    let fixedText = workingText;
+    if (inString) {
+      fixedText += '"';
+    }
+    while (braces > 0) {
+      fixedText += "}";
+      braces--;
+    }
+    while (brackets > 0) {
+      fixedText += "]";
+      brackets--;
+    }
+    try {
+      const result = JSON.parse(fixedText);
+      if (logger_default.level !== "silent") {
+        logger_default.warn(`JSON response appears truncated (possibly hit maxOutputTokens limit). Recovered by adding closing characters.`);
+      }
+      return result;
+    } catch (e) {
+    }
+  }
+  for (let i = 0; i < maxAttempts && workingText.length > 2; i++) {
+    workingText = workingText.slice(0, -1);
+    let braces2 = 0;
+    let brackets2 = 0;
+    let inString2 = false;
+    let escapeNext2 = false;
+    for (let j = 0; j < workingText.length; j++) {
+      const char = workingText[j];
+      if (escapeNext2) {
+        escapeNext2 = false;
+        continue;
+      }
+      if (char === "\\") {
+        escapeNext2 = true;
+        continue;
+      }
+      if (char === '"') {
+        inString2 = !inString2;
+        continue;
+      }
+      if (!inString2) {
+        if (char === "{") braces2++;
+        else if (char === "}") braces2--;
+        else if (char === "[") brackets2++;
+        else if (char === "]") brackets2--;
+      }
+    }
+    if (braces2 === 0 && brackets2 === 0 && !inString2) {
+      try {
+        const result = JSON.parse(workingText);
+        if (logger_default.level !== "silent") {
+          logger_default.warn(`JSON response appears truncated (possibly hit maxOutputTokens limit). Recovered by removing ${i + 1} characters from the end.`);
+        }
+        return result;
+      } catch (e) {
+      }
+    }
+    if (i > 5) {
+      let fixedText = workingText;
+      if (inString2) {
+        fixedText += '"';
+      }
+      while (braces2 > 0) {
+        fixedText += "}";
+        braces2--;
+      }
+      while (brackets2 > 0) {
+        fixedText += "]";
+        brackets2--;
+      }
+      try {
+        const result = JSON.parse(fixedText);
+        if (logger_default.level !== "silent") {
+          logger_default.warn(`JSON response appears truncated (possibly hit maxOutputTokens limit). Recovered by adding closing characters.`);
+        }
+        return result;
+      } catch (e) {
+      }
+    }
+  }
+  return null;
+}
 function isJSON(data) {
   try {
     const attempt = JSON.stringify(data);
@@ -497,6 +690,10 @@ function extractJSON(text) {
   if (isJSONStr(cleanedText)) {
     return JSON.parse(cleanedText);
   }
+  const recoveredJSON = attemptJSONRecovery(text);
+  if (recoveredJSON !== null) {
+    return recoveredJSON;
+  }
   throw new Error(`Could not extract valid JSON from model response. Response preview: ${text.substring(0, 200)}...`);
 }
 function findCompleteJSONStructures(text) {
@@ -603,6 +800,9 @@ 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 };
 
 
 
@@ -57,6 +57,8 @@ const DEFAULT_THINKING_CONFIG = {
 	thinkingLevel: ThinkingLevel.MINIMAL
 };
 
+const DEFAULT_MAX_OUTPUT_TOKENS = 100000; // Default ceiling for output tokens
+
 // Models that support thinking features (as of Dec 2024)
 // Using regex patterns for more precise matching
 const THINKING_SUPPORTED_MODELS = [
@@ -136,6 +138,7 @@ class AITransformer {
 }
 
 export default AITransformer;
+export { attemptJSONRecovery }; // Export for testing
 
 /**
  * factory function to create an AI Transformer instance
@@ -186,25 +189,53 @@ function AITransformFactory(options = {}) {
 		systemInstruction: this.systemInstructions
 	};
 
+	// Handle maxOutputTokens with explicit null check
+	// Priority: options.maxOutputTokens > options.chatConfig.maxOutputTokens > DEFAULT
+	// Setting to null explicitly removes the limit
+	if (options.maxOutputTokens !== undefined) {
+		if (options.maxOutputTokens === null) {
+			delete this.chatConfig.maxOutputTokens;
+		} else {
+			this.chatConfig.maxOutputTokens = options.maxOutputTokens;
+		}
+	} else if (options.chatConfig?.maxOutputTokens !== undefined) {
+		if (options.chatConfig.maxOutputTokens === null) {
+			delete this.chatConfig.maxOutputTokens;
+		} else {
+			this.chatConfig.maxOutputTokens = options.chatConfig.maxOutputTokens;
+		}
+	} else {
+		this.chatConfig.maxOutputTokens = DEFAULT_MAX_OUTPUT_TOKENS;
+	}
+
 	// Only add thinkingConfig if the model supports it
 	const modelSupportsThinking = THINKING_SUPPORTED_MODELS.some(pattern =>
 		pattern.test(this.modelName)
 	);
 
-	if (modelSupportsThinking && options.thinkingConfig) {
-		// Handle thinkingConfig - merge with defaults
-		const thinkingConfig = {
-			...DEFAULT_THINKING_CONFIG,
-			...options.thinkingConfig
-		};
-		this.chatConfig.thinkingConfig = thinkingConfig;
-
-		if (log.level !== 'silent') {
-			log.debug(`Model ${this.modelName} supports thinking. Applied thinkingConfig:`, thinkingConfig);
-		}
-	} else if (options.thinkingConfig && !modelSupportsThinking) {
-		if (log.level !== 'silent') {
-			log.warn(`Model ${this.modelName} does not support thinking features. Ignoring thinkingConfig.`);
+	// Handle thinkingConfig - null explicitly removes it, undefined means not specified
+	if (options.thinkingConfig !== undefined) {
+		if (options.thinkingConfig === null) {
+			// Explicitly remove thinkingConfig if set to null
+			delete this.chatConfig.thinkingConfig;
+			if (log.level !== 'silent') {
+				log.debug(`thinkingConfig set to null - removed from configuration`);
+			}
+		} else if (modelSupportsThinking) {
+			// Handle thinkingConfig - merge with defaults
+			const thinkingConfig = {
+				...DEFAULT_THINKING_CONFIG,
+				...options.thinkingConfig
+			};
+			this.chatConfig.thinkingConfig = thinkingConfig;
+
+			if (log.level !== 'silent') {
+				log.debug(`Model ${this.modelName} supports thinking. Applied thinkingConfig:`, thinkingConfig);
+			}
+		} else {
+			if (log.level !== 'silent') {
+				log.warn(`Model ${this.modelName} does not support thinking features. Ignoring thinkingConfig.`);
+			}
 		}
 	}
 
@@ -234,6 +265,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.");
 	}
@@ -241,6 +276,10 @@ function AITransformFactory(options = {}) {
 	if (log.level !== 'silent') {
 		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 });
@@ -259,12 +298,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();
@@ -431,6 +481,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
 
@@ -466,6 +557,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) {
@@ -474,6 +571,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}`);
 			}
 
@@ -575,12 +678,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.");
@@ -606,6 +721,173 @@ HELPERS
 ----
 */
 
+/**
+ * Attempts to recover truncated JSON by progressively removing characters from the end
+ * until valid JSON is found or recovery fails
+ * @param {string} text - The potentially truncated JSON string
+ * @param {number} maxAttempts - Maximum number of characters to remove
+ * @returns {Object|null} - Parsed JSON object or null if recovery fails
+ */
+function attemptJSONRecovery(text, maxAttempts = 100) {
+	if (!text || typeof text !== 'string') return null;
+
+	// First, try parsing as-is
+	try {
+		return JSON.parse(text);
+	} catch (e) {
+		// Continue with recovery
+	}
+
+	let workingText = text.trim();
+
+	// First attempt: try to close unclosed structures without removing characters
+	// Count open/close braces and brackets in the original text
+	let braces = 0;
+	let brackets = 0;
+	let inString = false;
+	let escapeNext = false;
+
+	for (let j = 0; j < workingText.length; j++) {
+		const char = workingText[j];
+
+		if (escapeNext) {
+			escapeNext = false;
+			continue;
+		}
+
+		if (char === '\\') {
+			escapeNext = true;
+			continue;
+		}
+
+		if (char === '"') {
+			inString = !inString;
+			continue;
+		}
+
+		if (!inString) {
+			if (char === '{') braces++;
+			else if (char === '}') braces--;
+			else if (char === '[') brackets++;
+			else if (char === ']') brackets--;
+		}
+	}
+
+	// Try to fix by just adding closing characters
+	if ((braces > 0 || brackets > 0 || inString) && workingText.length > 2) {
+		let fixedText = workingText;
+
+		// Close any open strings first
+		if (inString) {
+			fixedText += '"';
+		}
+
+		// Add missing closing characters
+		while (braces > 0) {
+			fixedText += '}';
+			braces--;
+		}
+		while (brackets > 0) {
+			fixedText += ']';
+			brackets--;
+		}
+
+		try {
+			const result = JSON.parse(fixedText);
+			if (log.level !== 'silent') {
+				log.warn(`JSON response appears truncated (possibly hit maxOutputTokens limit). Recovered by adding closing characters.`);
+			}
+			return result;
+		} catch (e) {
+			// Simple fix didn't work, continue with more aggressive recovery
+		}
+	}
+
+	// Second attempt: progressively remove characters from the end
+
+	for (let i = 0; i < maxAttempts && workingText.length > 2; i++) {
+		// Remove one character from the end
+		workingText = workingText.slice(0, -1);
+
+		// Count open/close braces and brackets
+		let braces = 0;
+		let brackets = 0;
+		let inString = false;
+		let escapeNext = false;
+
+		for (let j = 0; j < workingText.length; j++) {
+			const char = workingText[j];
+
+			if (escapeNext) {
+				escapeNext = false;
+				continue;
+			}
+
+			if (char === '\\') {
+				escapeNext = true;
+				continue;
+			}
+
+			if (char === '"') {
+				inString = !inString;
+				continue;
+			}
+
+			if (!inString) {
+				if (char === '{') braces++;
+				else if (char === '}') braces--;
+				else if (char === '[') brackets++;
+				else if (char === ']') brackets--;
+			}
+		}
+
+		// If we have balanced braces/brackets, try parsing
+		if (braces === 0 && brackets === 0 && !inString) {
+			try {
+				const result = JSON.parse(workingText);
+				if (log.level !== 'silent') {
+					log.warn(`JSON response appears truncated (possibly hit maxOutputTokens limit). Recovered by removing ${i + 1} characters from the end.`);
+				}
+				return result;
+			} catch (e) {
+				// Continue trying
+			}
+		}
+
+		// After a few attempts, try adding closing characters
+		if (i > 5) {
+			let fixedText = workingText;
+
+			// Close any open strings first
+			if (inString) {
+				fixedText += '"';
+			}
+
+			// Add missing closing characters
+			while (braces > 0) {
+				fixedText += '}';
+				braces--;
+			}
+			while (brackets > 0) {
+				fixedText += ']';
+				brackets--;
+			}
+
+			try {
+				const result = JSON.parse(fixedText);
+				if (log.level !== 'silent') {
+					log.warn(`JSON response appears truncated (possibly hit maxOutputTokens limit). Recovered by adding closing characters.`);
+				}
+				return result;
+			} catch (e) {
+				// Recovery failed, continue trying
+			}
+		}
+	}
+
+	return null;
+}
+
 function isJSON(data) {
 	try {
 		const attempt = JSON.stringify(data);
@@ -703,6 +985,13 @@ function extractJSON(text) {
 		return JSON.parse(cleanedText);
 	}
 
+	// Strategy 6: Last resort - attempt recovery for potentially truncated JSON
+	// This is especially useful when maxOutputTokens might have cut off the response
+	const recoveredJSON = attemptJSONRecovery(text);
+	if (recoveredJSON !== null) {
+		return recoveredJSON;
+	}
+
 	// If all else fails, throw an error with helpful information
 	throw new Error(`Could not extract valid JSON from model response. Response preview: ${text.substring(0, 200)}...`);
 }

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.7",
+	"version": "1.0.9",
 	"main": "index.js",
 	"files": [
 		"index.js",
@@ -39,6 +39,7 @@
 		"prune": "rm -rf tmp/*",
 		"test": "node --no-warnings --experimental-vm-modules node_modules/jest/bin/jest.js",
 		"test:unit": "npm test -- tests/module.test.js",
+		"test:fixed": "npm test -- --testNamePattern=\"should use context in the prompt and transform accordingly|should augment the payload as instructed by system instructions|should succeed on the first try if validation passes|should handle invalid model names|should handle multiple concurrent messages|should use the constructor-provided asyncValidator|should override system instructions from the file\"",
 		"build:cjs": "esbuild index.js --bundle --platform=node --format=cjs --outfile=index.cjs --external:@google/genai --external:ak-tools --external:dotenv --external:pino-pretty --external:pino",
 		"coverage": "node --no-warnings --experimental-vm-modules node_modules/jest/bin/jest.js --coverage",
 		"typecheck": "tsc --noEmit",

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 {
@@ -21,6 +21,7 @@ export interface ChatConfig {
   temperature?: number; // Controls randomness (0.0 to 1.0)
   topP?: number; // Controls diversity via nucleus sampling
   topK?: number; // Controls diversity by limiting top-k tokens
+  maxOutputTokens?: number; // Maximum number of tokens that can be generated in the response
   systemInstruction?: string; // System instruction for the model
   safetySettings?: SafetySetting[]; // Safety settings array
   responseSchema?: Object; // Schema for validating model responses
@@ -50,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 {
@@ -74,6 +77,7 @@ export interface AITransformerOptions {
   systemInstructions?: string; // Custom system instructions for the model
   chatConfig?: ChatConfig; // Configuration object for the chat session
   thinkingConfig?: ThinkingConfig; // Thinking features configuration (defaults to thinkingBudget: 0, thinkingLevel: "MINIMAL")
+  maxOutputTokens?: number; // Maximum number of tokens that can be generated in the response (defaults to 100000)
   examplesFile?: string; // Path to JSON file containing transformation examples
   exampleData?: TransformationExample[]; // Inline examples to seed the transformer
   sourceKey?: string; // Key name for source data in examples (alias for promptKey)
@@ -91,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
@@ -118,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>;