ak-gemini 1.0.4 → 1.0.5

package/README.md

@@ -7,12 +7,13 @@ Use this to power LLM-driven data pipelines, JSON mapping, or any automated AI t
 
 ## Features
 
-* **Model-Agnostic**: Configure for any Gemini model (`gemini-2.0-flash` by default)
-* **Declarative Examples**: Seed transformations using example mappings, with support for custom keys (`PROMPT`, `ANSWER`, `CONTEXT`, or your own)
-* **Automatic Validation & Repair**: Validate outputs with your own async function; auto-repair failed payloads with LLM feedback loop (exponential backoff, fully configurable)
-* **Strong TypeScript/JSDoc Typings**: All public APIs fully typed (see `/types`)
-* **Minimal API Surface**: Dead simple, no ceremony—init, seed, transform, validate.
-* **Robust Logging**: Pluggable logger for all steps, easy debugging
+* **Model-Agnostic:** Use any Gemini model (`gemini-2.0-flash` by default)
+* **Declarative Few-shot Examples:** Seed transformations using example mappings, with support for custom keys (`PROMPT`, `ANSWER`, `CONTEXT`, or your own)
+* **Automatic Validation & Repair:** Validate outputs with your own async function; auto-repair failed payloads with LLM feedback loop (exponential backoff, fully configurable)
+* **Token Counting & Safety:** Preview the *exact* Gemini token consumption for any operation—including all examples, instructions, and your input—before sending, so you can avoid window errors and manage costs.
+* **Strong TypeScript/JSDoc Typings:** All public APIs fully typed (see `/types`)
+* **Minimal API Surface:** Dead simple, no ceremony—init, seed, transform, validate.
+* **Robust Logging:** Pluggable logger for all steps, easy debugging
 
 ---
 
@@ -43,7 +44,7 @@ or pass it directly in the constructor options.
 ### 2. **Basic Example**
 
 ```js
-import AITransformer from 'ai-transformer';
+import AITransformer from 'ak-gemini';
 
 const transformer = new AITransformer({
   modelName: 'gemini-2.0-flash',    // or your preferred Gemini model
@@ -72,7 +73,22 @@ console.log(result);
 
 ---
 
-### 3. **Automatic Validation & Self-Healing**
+### 3. **Token Window Safety/Preview**
+
+Before calling `.message()` or `.seed()`, you can preview the exact token usage that will be sent to Gemini—*including* your system instructions, examples, and user input. This is vital for avoiding window errors and managing context size:
+
+```js
+const { totalTokens, breakdown } = await transformer.estimateTokenUsage({ name: "Bob" });
+console.log(`Total tokens: ${totalTokens}`);
+console.log(breakdown); // See per-section token counts
+
+// Optional: abort or trim if over limit
+if (totalTokens > 32000) throw new Error("Request too large for selected Gemini model");
+```
+
+---
+
+### 4. **Automatic Validation & Self-Healing**
 
 You can pass a custom async validator—if it fails, the transformer will attempt to self-correct using LLM feedback, retrying up to `maxRetries` times:
 
@@ -127,7 +143,12 @@ You can omit `examples` to use the `examplesFile` (if provided).
 
 #### `await transformer.message(sourcePayload)`
 
-Transforms input JSON to output JSON using the seeded examples and system instructions.
+Transforms input JSON to output JSON using the seeded examples and system instructions. Throws if estimated token window would be exceeded.
+
+#### `await transformer.estimateTokenUsage(sourcePayload)`
+
+Returns `{ totalTokens, breakdown }` for the *full request* that would be sent to Gemini (system instructions + all examples + your sourcePayload as the new prompt).
+Lets you preview token window safety and abort/trim as needed.
 
 #### `await transformer.transformWithValidation(sourcePayload, validatorFn, options?)`
 
@@ -187,10 +208,19 @@ const result = await transformer.transformWithValidation(
 
 ---
 
+## Token Window Management & Error Handling
+
+* Throws on missing `GEMINI_API_KEY`
+* `.message()` and `.seed()` will *estimate* and prevent calls that would exceed Gemini's model window
+* All API and parsing errors surfaced as `Error` with context
+* Validator and retry failures include the number of attempts and last error
+
+---
+
 ## Testing
 
 * **Jest test suite included**
-* Mocks Google Gemini, logger, ak-tools
+* Real API integration tests as well as local unit tests
 * 100% coverage for all error cases, configuration options, edge cases
 
 Run tests with:
@@ -200,13 +230,3 @@ npm test
 ```
 
 ---
-
-## Error Handling
-
-* Throws on missing `GEMINI_API_KEY`
-* All API and parsing errors surfaced as `Error` with context
-* Validator and retry failures include the number of attempts and last error
-
----
-
-

package/index.cjs

@@ -108,6 +108,7 @@ var AITransformer = class {
     this.reset = resetChat.bind(this);
     this.getHistory = getChatHistory.bind(this);
     this.transformWithValidation = transformWithValidation.bind(this);
+    this.estimate = estimateTokenUsage.bind(this);
   }
 };
 function AITransformFactory(options = {}) {
@@ -135,7 +136,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}"`);
-  this.genAIClient = new import_genai.GoogleGenAI({ apiKey: this.apiKey });
+  const ai = new import_genai.GoogleGenAI({ apiKey: this.apiKey });
+  this.genAIClient = ai;
   this.chat = null;
 }
 async function initChat() {
@@ -186,7 +188,7 @@ async function seedWithExamples(examples) {
       historyToAdd.push({ role: "model", parts: [{ text: answerText }] });
     }
   }
-  const currentHistory = this.chat.getHistory();
+  const currentHistory = this?.chat?.getHistory() || [];
   this.chat = await this.genAIClient.chats.create({
     model: this.modelName,
     // @ts-ignore
@@ -246,6 +248,25 @@ async function transformWithValidation(sourcePayload, validatorFn, options = {})
     }
   }
 }
+async function estimateTokenUsage(nextPayload) {
+  const contents = [];
+  if (this.systemInstructions) {
+    contents.push({ parts: [{ text: this.systemInstructions }] });
+  }
+  if (this.chat && typeof this.chat.getHistory === "function") {
+    const history = this.chat.getHistory();
+    if (Array.isArray(history) && history.length > 0) {
+      contents.push(...history);
+    }
+  }
+  const nextMessage = typeof nextPayload === "string" ? nextPayload : JSON.stringify(nextPayload, null, 2);
+  contents.push({ parts: [{ text: nextMessage }] });
+  const resp = await this.genAIClient.models.countTokens({
+    model: this.modelName,
+    contents
+  });
+  return resp;
+}
 async function rebuildPayload(lastPayload, serverError) {
   await this.init();
   const prompt = `

package/index.js

@@ -96,6 +96,7 @@ export default class AITransformer {
 		this.reset = resetChat.bind(this);
 		this.getHistory = getChatHistory.bind(this);
 		this.transformWithValidation = transformWithValidation.bind(this);
+		this.estimate = estimateTokenUsage.bind(this);
 	}
 }
 
@@ -143,7 +144,8 @@ 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}"`);
 
-	this.genAIClient = new GoogleGenAI({ apiKey: this.apiKey });
+	const ai = new GoogleGenAI({ apiKey: this.apiKey });
+	this.genAIClient = ai;
 	this.chat = null;
 }
 
@@ -221,7 +223,7 @@ async function seedWithExamples(examples) {
 		}
 	}
 
-	const currentHistory = this.chat.getHistory();
+	const currentHistory = this?.chat?.getHistory() || [];
 
 	this.chat = await this.genAIClient.chats.create({
 		model: this.modelName,
@@ -317,6 +319,47 @@ async function transformWithValidation(sourcePayload, validatorFn, options = {})
 	}
 }
 
+
+/**
+ * Estimate total token usage if you were to send a new payload as the next message.
+ * Considers system instructions, current chat history (including examples), and the new message.
+ * @param {object|string} nextPayload - The next user message to be sent (object or string)
+ * @returns {Promise<{ totalTokens: number, ... }>} - The result of Gemini's countTokens API
+ */
+async function estimateTokenUsage(nextPayload) {
+	// Compose the conversation contents, Gemini-style
+	const contents = [];
+
+	// (1) System instructions (if applicable)
+	if (this.systemInstructions) {
+		// Add as a 'system' part; adjust role if Gemini supports
+		contents.push({ parts: [{ text: this.systemInstructions }] });
+	}
+
+	// (2) All current chat history (seeded examples + real user/model turns)
+	if (this.chat && typeof this.chat.getHistory === "function") {
+		const history = this.chat.getHistory();
+		if (Array.isArray(history) && history.length > 0) {
+			contents.push(...history);
+		}
+	}
+
+	// (3) The next user message
+	const nextMessage = typeof nextPayload === "string"
+		? nextPayload
+		: JSON.stringify(nextPayload, null, 2);
+
+	contents.push({ parts: [{ text: nextMessage }] });
+
+	// Call Gemini's token estimator
+	const resp = await this.genAIClient.models.countTokens({
+		model: this.modelName,
+		contents,
+	});
+
+	return resp; // includes totalTokens, possibly breakdown
+}
+
 /**
  * Rebuilds a payload based on server error feedback
  * @param {Object} lastPayload - The payload that failed validation

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.4",
+	"version": "1.0.5",
 	"main": "index.js",
 	"files": [
 		"index.js",
@@ -47,7 +47,7 @@
 	"license": "ISC",
 	"dependencies": {
 		"@google-cloud/functions-framework": "^4.0.0",
-		"@google/genai": "^1.3.0",
+		"@google/genai": "^1.4.0",
 		"ak-tools": "^1.0.64",
 		"dotenv": "^16.5.0",
 		"pino": "^9.7.0",