ak-gemini 1.1.13 → 2.0.0

package/json-helpers.js

@@ -0,0 +1,352 @@
+/**
+ * @fileoverview Pure utility functions for JSON extraction and recovery.
+ * Used by Transformer and Message classes to parse AI model responses.
+ */
+
+import log from './logger.js';
+
+/**
+ * Checks if a JavaScript value is a JSON-serializable object or array.
+ * @param {*} data - The value to check
+ * @returns {boolean}
+ */
+export function isJSON(data) {
+	try {
+		const attempt = JSON.stringify(data);
+		if (attempt?.startsWith('{') || attempt?.startsWith('[')) {
+			if (attempt?.endsWith('}') || attempt?.endsWith(']')) {
+				return true;
+			}
+		}
+		return false;
+	} catch (e) {
+		return false;
+	}
+}
+
+/**
+ * Checks if a string is valid JSON that parses to an object or array.
+ * @param {string} string - The string to check
+ * @returns {boolean}
+ */
+export function isJSONStr(string) {
+	if (typeof string !== 'string') return false;
+	try {
+		const result = JSON.parse(string);
+		const type = Object.prototype.toString.call(result);
+		return type === '[object Object]' || type === '[object Array]';
+	} catch (err) {
+		return false;
+	}
+}
+
+/**
+ * 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=100] - Maximum number of characters to remove
+ * @returns {Object|null} - Parsed JSON object or null if recovery fails
+ */
+export 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
+	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;
+
+		if (inString) {
+			fixedText += '"';
+		}
+
+		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++) {
+		workingText = workingText.slice(0, -1);
+
+		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;
+
+			if (inString) {
+				fixedText += '"';
+			}
+
+			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;
+}
+
+/**
+ * Extracts a complete JSON structure from text starting at a given position
+ * using bracket/brace matching.
+ * @param {string} text - The text containing JSON
+ * @param {number} startPos - Position of the opening bracket/brace
+ * @returns {string|null} - The complete JSON structure or null
+ */
+function extractCompleteStructure(text, startPos) {
+	const startChar = text[startPos];
+	const endChar = startChar === '{' ? '}' : ']';
+	let depth = 0;
+	let inString = false;
+	let escaped = false;
+
+	for (let i = startPos; i < text.length; i++) {
+		const char = text[i];
+
+		if (escaped) {
+			escaped = false;
+			continue;
+		}
+
+		if (char === '\\' && inString) {
+			escaped = true;
+			continue;
+		}
+
+		if (char === '"' && !escaped) {
+			inString = !inString;
+			continue;
+		}
+
+		if (!inString) {
+			if (char === startChar) {
+				depth++;
+			} else if (char === endChar) {
+				depth--;
+				if (depth === 0) {
+					return text.substring(startPos, i + 1);
+				}
+			}
+		}
+	}
+
+	return null;
+}
+
+/**
+ * Finds all complete JSON structures (objects and arrays) in text.
+ * @param {string} text - The text to search
+ * @returns {string[]} - Array of JSON structure strings
+ */
+function findCompleteJSONStructures(text) {
+	const results = [];
+	const startChars = ['{', '['];
+
+	for (let i = 0; i < text.length; i++) {
+		if (startChars.includes(text[i])) {
+			const extracted = extractCompleteStructure(text, i);
+			if (extracted) {
+				results.push(extracted);
+			}
+		}
+	}
+
+	return results;
+}
+
+/**
+ * Extracts valid JSON from model response text using multiple strategies.
+ * @param {string} text - The model response text
+ * @returns {Object} - Parsed JSON object
+ * @throws {Error} If no valid JSON can be extracted
+ */
+export function extractJSON(text) {
+	if (!text || typeof text !== 'string') {
+		throw new Error('No text provided for JSON extraction');
+	}
+
+	// Strategy 1: Try parsing the entire response as JSON
+	if (isJSONStr(text.trim())) {
+		return JSON.parse(text.trim());
+	}
+
+	// Strategy 2: Look for JSON code blocks (```json...``` or ```...```)
+	const codeBlockPatterns = [
+		/```json\s*\n?([\s\S]*?)\n?\s*```/gi,
+		/```\s*\n?([\s\S]*?)\n?\s*```/gi
+	];
+
+	for (const pattern of codeBlockPatterns) {
+		const matches = text.match(pattern);
+		if (matches) {
+			for (const match of matches) {
+				const jsonContent = match.replace(/```json\s*\n?/gi, '').replace(/```\s*\n?/gi, '').trim();
+				if (isJSONStr(jsonContent)) {
+					return JSON.parse(jsonContent);
+				}
+			}
+		}
+	}
+
+	// Strategy 3: Look for JSON objects/arrays using bracket matching
+	const jsonPatterns = [
+		/\{[\s\S]*\}/g,
+		/\[[\s\S]*\]/g
+	];
+
+	for (const pattern of jsonPatterns) {
+		const matches = text.match(pattern);
+		if (matches) {
+			for (const match of matches) {
+				const candidate = match.trim();
+				if (isJSONStr(candidate)) {
+					return JSON.parse(candidate);
+				}
+			}
+		}
+	}
+
+	// Strategy 4: Advanced bracket matching for nested structures
+	const advancedExtract = findCompleteJSONStructures(text);
+	if (advancedExtract.length > 0) {
+		for (const candidate of advancedExtract) {
+			if (isJSONStr(candidate)) {
+				return JSON.parse(candidate);
+			}
+		}
+	}
+
+	// Strategy 5: Clean up common formatting issues and retry
+	const cleanedText = text
+		.replace(/^\s*Sure,?\s*here\s+is\s+your?\s+.*?[:\n]/gi, '')
+		.replace(/^\s*Here\s+is\s+the\s+.*?[:\n]/gi, '')
+		.replace(/^\s*The\s+.*?is\s*[:\n]/gi, '')
+		.replace(/\/\*[\s\S]*?\*\//g, '')
+		.replace(/\/\/.*$/gm, '')
+		.trim();
+
+	if (isJSONStr(cleanedText)) {
+		return JSON.parse(cleanedText);
+	}
+
+	// Strategy 6: Last resort - attempt recovery for potentially truncated JSON
+	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)}...`);
+}

package/message.js

@@ -0,0 +1,170 @@
+/**
+ * @fileoverview Message class — stateless one-off messages to AI.
+ * Uses generateContent() instead of chat sessions. No conversation history.
+ */
+
+import BaseGemini from './base.js';
+import { extractJSON } from './json-helpers.js';
+import log from './logger.js';
+
+/**
+ * @typedef {import('./types').MessageOptions} MessageOptions
+ * @typedef {import('./types').MessageResponse} MessageResponse
+ */
+
+/**
+ * Stateless one-off messages to AI.
+ * Each send() call is independent — no conversation history is maintained.
+ * Uses generateContent() directly instead of chat sessions.
+ *
+ * Optionally returns structured data when responseSchema or
+ * responseMimeType: 'application/json' is configured.
+ *
+ * @example
+ * ```javascript
+ * import { Message } from 'ak-gemini';
+ *
+ * // Simple text response
+ * const msg = new Message({
+ *   systemPrompt: 'You are a helpful assistant.'
+ * });
+ * const r = await msg.send('What is the capital of France?');
+ * console.log(r.text); // "The capital of France is Paris."
+ *
+ * // Structured JSON response
+ * const jsonMsg = new Message({
+ *   systemPrompt: 'Extract entities from text.',
+ *   responseMimeType: 'application/json'
+ * });
+ * const r2 = await jsonMsg.send('Alice works at Acme Corp in New York.');
+ * console.log(r2.data); // { entities: [...] }
+ * ```
+ */
+class Message extends BaseGemini {
+	/**
+	 * @param {MessageOptions} [options={}]
+	 */
+	constructor(options = {}) {
+		super(options);
+
+		// ── Structured output config ──
+		if (options.responseSchema) {
+			this.chatConfig.responseSchema = options.responseSchema;
+		}
+		if (options.responseMimeType) {
+			this.chatConfig.responseMimeType = options.responseMimeType;
+		}
+
+		this._isStructured = !!(options.responseSchema || options.responseMimeType === 'application/json');
+
+		log.debug(`Message created (structured=${this._isStructured})`);
+	}
+
+	/**
+	 * Initialize the Message client.
+	 * Override: creates genAIClient only, NO chat session (stateless).
+	 * @param {boolean} [force=false]
+	 * @returns {Promise<void>}
+	 */
+	async init(force = false) {
+		if (this._initialized && !force) return;
+
+		log.debug(`Initializing ${this.constructor.name} with model: ${this.modelName}...`);
+
+		try {
+			await this.genAIClient.models.list();
+			log.debug(`${this.constructor.name}: API connection successful.`);
+		} catch (e) {
+			throw new Error(`${this.constructor.name} initialization failed: ${e.message}`);
+		}
+
+		this._initialized = true;
+		log.debug(`${this.constructor.name}: Initialized (stateless mode).`);
+	}
+
+	/**
+	 * Send a stateless message and get a response.
+	 * Each call is independent — no history is maintained.
+	 *
+	 * @param {Object|string} payload - The message or data to send
+	 * @param {Object} [opts={}] - Per-message options
+	 * @param {Record<string, string>} [opts.labels] - Per-message billing labels
+	 * @returns {Promise<MessageResponse>} Response with text, optional data, and usage
+	 */
+	async send(payload, opts = {}) {
+		if (!this._initialized) await this.init();
+
+		const payloadStr = typeof payload === 'string'
+			? payload
+			: JSON.stringify(payload, null, 2);
+
+		const contents = [{ role: 'user', parts: [{ text: payloadStr }] }];
+
+		const mergedLabels = { ...this.labels, ...(opts.labels || {}) };
+
+		const result = await this.genAIClient.models.generateContent({
+			model: this.modelName,
+			contents: contents,
+			config: {
+				...this.chatConfig,
+				...(this.vertexai && Object.keys(mergedLabels).length > 0 && { labels: mergedLabels })
+			}
+		});
+
+		this._captureMetadata(result);
+
+		this._cumulativeUsage = {
+			promptTokens: this.lastResponseMetadata.promptTokens,
+			responseTokens: this.lastResponseMetadata.responseTokens,
+			totalTokens: this.lastResponseMetadata.totalTokens,
+			attempts: 1
+		};
+
+		if (result.usageMetadata && log.level !== 'silent') {
+			log.debug(`Message response: model=${result.modelVersion || 'unknown'}, tokens=${result.usageMetadata.totalTokenCount}`);
+		}
+
+		const text = result.text || '';
+		const response = {
+			text,
+			usage: this.getLastUsage()
+		};
+
+		// Parse structured data if configured
+		if (this._isStructured) {
+			try {
+				response.data = extractJSON(text);
+			} catch (e) {
+				log.warn(`Could not parse structured response: ${e.message}`);
+				response.data = null;
+			}
+		}
+
+		return response;
+	}
+
+	// ── No-ops for stateless class ──
+
+	/** @returns {Array} Always returns empty array (stateless). */
+	getHistory() { return []; }
+
+	/** No-op (stateless). */
+	async clearHistory() { }
+
+	/** Not supported on Message (stateless). */
+	async seed() {
+		log.warn("Message is stateless — seed() has no effect. Use Transformer or Chat for few-shot learning.");
+		return [];
+	}
+
+	/**
+	 * Not supported on Message (stateless).
+	 * @param {any} [_nextPayload]
+	 * @returns {Promise<{inputTokens: number}>}
+	 */
+	async estimate(_nextPayload) {
+		throw new Error("Message is stateless — use estimate() on Chat or Transformer which have conversation context.");
+	}
+}
+
+export default Message;

package/package.json

@@ -1,12 +1,19 @@
 {
 	"name": "ak-gemini",
 	"author": "ak@mixpanel.com",
-	"description": "AK's Generative AI Helper for doing... transforms",
-	"version": "1.1.13",
+	"description": "AK's Generative AI Helper for doing... everything",
+	"version": "2.0.0",
 	"main": "index.js",
 	"files": [
 		"index.js",
 		"index.cjs",
+		"base.js",
+		"transformer.js",
+		"chat.js",
+		"message.js",
+		"tool-agent.js",
+		"code-agent.js",
+		"json-helpers.js",
 		"types.d.ts",
 		"logger.js"
 	],
@@ -38,7 +45,7 @@
 		"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",
-		"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",
+		"build:cjs": "esbuild index.js --bundle --platform=node --format=cjs --outfile=index.cjs --external:@google/genai --external:dotenv --external:pino-pretty --external:pino",
 		"coverage": "node --no-warnings --experimental-vm-modules node_modules/jest/bin/jest.js --coverage",
 		"typecheck": "tsc --noEmit",
 		"update:gemini": "npm install @google/genai@latest"
@@ -47,21 +54,23 @@
 	"keywords": [
 		"gemini",
 		"ai wrapper",
-		"json transform"
+		"json transform",
+		"chat",
+		"agent",
+		"tool agent"
 	],
 	"license": "ISC",
 	"dependencies": {
-		"@google/genai": "^1.34.0",
-		"ak-tools": "^1.1.12",
-		"dotenv": "^16.5.0",
-		"pino": "^9.7.0",
-		"pino-pretty": "^13.0.0"
+		"@google/genai": "^1.44.0",
+		"dotenv": "^17.3.1",
+		"pino": "^10.3.1",
+		"pino-pretty": "^13.1.3"
 	},
 	"devDependencies": {
-		"@types/jest": "^29.5.14",
-		"esbuild": "^0.25.11",
-		"jest": "^29.7.0",
-		"nodemon": "^3.1.10",
-		"typescript": "^5.9.2"
+		"@types/jest": "^30.0.0",
+		"esbuild": "^0.27.4",
+		"jest": "^30.3.0",
+		"nodemon": "^3.1.14",
+		"typescript": "^5.9.3"
 	}
 }