ak-gemini 1.0.7 → 1.0.8

package/index.cjs

@@ -30,6 +30,7 @@ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: tru
 var index_exports = {};
 __export(index_exports, {
   ThinkingLevel: () => import_genai.ThinkingLevel,
+  attemptJSONRecovery: () => attemptJSONRecovery,
   default: () => index_default,
   log: () => logger_default
 });
@@ -79,6 +80,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 +168,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) {
@@ -203,6 +227,7 @@ function AITransformFactory(options = {}) {
   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}`);
   }
   const ai = new import_genai.GoogleGenAI({ apiKey: this.apiKey });
   this.genAIClient = ai;
@@ -423,6 +448,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 +638,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) {
@@ -604,5 +749,6 @@ if (import_meta.url === new URL(`file://${process.argv[1]}`).href) {
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
   ThinkingLevel,
+  attemptJSONRecovery,
   log
 });

package/index.js

@@ -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.`);
+			}
 		}
 	}
 
@@ -241,6 +272,7 @@ 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}`);
 	}
 
 	const ai = new GoogleGenAI({ apiKey: this.apiKey });
@@ -606,6 +638,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 +902,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.8",
 	"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

@@ -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
@@ -74,6 +75,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)