ak-gemini 1.0.6 → 1.0.8

package/index.cjs

@@ -29,6 +29,8 @@ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: tru
 // index.js
 var index_exports = {};
 __export(index_exports, {
+  ThinkingLevel: () => import_genai.ThinkingLevel,
+  attemptJSONRecovery: () => attemptJSONRecovery,
   default: () => index_default,
   log: () => logger_default
 });
@@ -64,7 +66,7 @@ var DEFAULT_SAFETY_SETTINGS = [
 var DEFAULT_SYSTEM_INSTRUCTIONS = `
 You are an expert JSON transformation engine. Your task is to accurately convert data payloads from one format to another.
 
-You will be provided with example transformations (Source JSON -> Target JSON). 
+You will be provided with example transformations (Source JSON -> Target JSON).
 
 Learn the mapping rules from these examples.
 
@@ -72,8 +74,22 @@ When presented with new Source JSON, apply the learned transformation rules to p
 
 Always respond ONLY with a valid JSON object that strictly adheres to the expected output format.
 
-Do not include any additional text, explanations, or formatting before or after the JSON object. 
+Do not include any additional text, explanations, or formatting before or after the JSON object.
 `;
+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)?$/,
+  /^gemini-2\.5-pro/,
+  /^gemini-2\.5-flash(-preview)?$/,
+  /^gemini-2\.5-flash-lite(-preview)?$/,
+  /^gemini-2\.0-flash$/
+  // Experimental support, exact match only
+];
 var DEFAULT_CHAT_CONFIG = {
   responseMimeType: "application/json",
   temperature: 0.2,
@@ -152,6 +168,45 @@ 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 (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) {
     this.chatConfig.responseSchema = options.responseSchema;
   }
@@ -172,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;
@@ -392,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);
@@ -466,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) {
@@ -572,5 +748,7 @@ 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

@@ -23,11 +23,12 @@ const { NODE_ENV = "unknown", GEMINI_API_KEY, LOG_LEVEL = "" } = process.env;
 
 
 //deps
-import { GoogleGenAI, HarmCategory, HarmBlockThreshold } from '@google/genai';
+import { GoogleGenAI, HarmCategory, HarmBlockThreshold, ThinkingLevel } from '@google/genai';
 import u from 'ak-tools';
 import path from 'path';
 import log from './logger.js';
 export { log };
+export { ThinkingLevel };
 
 
 
@@ -40,7 +41,7 @@ const DEFAULT_SAFETY_SETTINGS = [
 const DEFAULT_SYSTEM_INSTRUCTIONS = `
 You are an expert JSON transformation engine. Your task is to accurately convert data payloads from one format to another.
 
-You will be provided with example transformations (Source JSON -> Target JSON). 
+You will be provided with example transformations (Source JSON -> Target JSON).
 
 Learn the mapping rules from these examples.
 
@@ -48,9 +49,27 @@ When presented with new Source JSON, apply the learned transformation rules to p
 
 Always respond ONLY with a valid JSON object that strictly adheres to the expected output format.
 
-Do not include any additional text, explanations, or formatting before or after the JSON object. 
+Do not include any additional text, explanations, or formatting before or after the JSON object.
 `;
 
+const DEFAULT_THINKING_CONFIG = {
+	thinkingBudget: 0,
+	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 = [
+	/^gemini-3-flash(-preview)?$/,
+	/^gemini-3-pro(-preview|-image-preview)?$/,
+	/^gemini-2\.5-pro/,
+	/^gemini-2\.5-flash(-preview)?$/,
+	/^gemini-2\.5-flash-lite(-preview)?$/,
+	/^gemini-2\.0-flash$/ // Experimental support, exact match only
+];
+
 const DEFAULT_CHAT_CONFIG = {
 	responseMimeType: 'application/json',
 	temperature: 0.2,
@@ -119,6 +138,7 @@ class AITransformer {
 }
 
 export default AITransformer;
+export { attemptJSONRecovery }; // Export for testing
 
 /**
  * factory function to create an AI Transformer instance
@@ -161,6 +181,7 @@ function AITransformFactory(options = {}) {
 
 	this.apiKey = options.apiKey !== undefined && options.apiKey !== null ? options.apiKey : GEMINI_API_KEY;
 	if (!this.apiKey) throw new Error("Missing Gemini API key. Provide via options.apiKey or GEMINI_API_KEY env var.");
+
 	// Build chat config, making sure systemInstruction uses the custom instructions
 	this.chatConfig = {
 		...DEFAULT_CHAT_CONFIG,
@@ -168,6 +189,56 @@ 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)
+	);
+
+	// 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.`);
+			}
+		}
+	}
+
 	// response schema is optional, but if provided, it should be a valid JSON schema
 	if (options.responseSchema) {
 		this.chatConfig.responseSchema = options.responseSchema;
@@ -201,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 });
@@ -566,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);
@@ -663,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.6",
+	"version": "1.0.8",
 	"main": "index.js",
 	"files": [
 		"index.js",
@@ -33,15 +33,17 @@
 	"homepage": "https://github.com/ak--47/ak-gemini#readme",
 	"scripts": {
 		"prepublishOnly": "npm run typecheck && npm run build:cjs",
-		"post": "npm publish --access public",
+		"post": "npm publish",
 		"release": "npm version patch && npm publish --access public",
 		"update-deps": "npx npm-check-updates -u && npm install",
 		"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"
+		"typecheck": "tsc --noEmit",
+		"update:gemini": "npm install @google/genai@latest"
 	},
 	"type": "module",
 	"keywords": [
@@ -51,7 +53,7 @@
 	],
 	"license": "ISC",
 	"dependencies": {
-		"@google/genai": "^1.25.0",
+		"@google/genai": "^1.34.0",
 		"ak-tools": "^1.1.12",
 		"dotenv": "^16.5.0",
 		"pino": "^9.7.0",

package/types.d.ts

@@ -1,4 +1,15 @@
-import type { GoogleGenAI } from '@google/genai';
+import type { GoogleGenAI, ThinkingLevel } from '@google/genai';
+
+export { ThinkingLevel };
+
+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. */
+  includeThoughts?: boolean;
+  /** Indicates the thinking budget in tokens. 0 is DISABLED. -1 is AUTOMATIC. The default values and allowed ranges are model dependent. */
+  thinkingBudget?: number;
+  /** Optional. The number of thoughts tokens that the model should generate. */
+  thinkingLevel?: ThinkingLevel;
+}
 
 export interface SafetySetting {
   category: string; // The harm category
@@ -10,9 +21,11 @@ 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
+  thinkingConfig?: ThinkingConfig; // Thinking features configuration
   [key: string]: any; // Additional properties for flexibility
 }
 
@@ -61,6 +74,8 @@ export interface AITransformerOptions {
   modelName?: string; // The Gemini model to use
   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)