ak-gemini 1.0.13 → 1.0.14

package/README.md

@@ -205,14 +205,15 @@ Clears conversation history while preserving seeded examples. Useful for startin
 
 #### `transformer.getLastUsage()`
 
-Returns structured usage data from the last API response for billing verification. Returns `null` if no API call has been made yet.
+Returns structured usage data for billing verification. Token counts are **cumulative across all retry attempts** - if validation failed and a retry was needed, you see the total tokens consumed, not just the final successful call. Returns `null` if no API call has been made yet.
 
 ```js
 const usage = transformer.getLastUsage();
 // {
-//   promptTokens: 150,      // Input tokens (includes system instructions + history + message)
-//   responseTokens: 42,     // Output tokens
-//   totalTokens: 192,       // Total tokens from API
+//   promptTokens: 300,      // CUMULATIVE input tokens across all attempts
+//   responseTokens: 84,     // CUMULATIVE output tokens across all attempts
+//   totalTokens: 384,       // CUMULATIVE total tokens
+//   attempts: 2,            // Number of attempts (1 = first try success, 2+ = retries needed)
 //   modelVersion: 'gemini-2.5-flash-001',  // Actual model that responded
 //   requestedModel: 'gemini-2.5-flash',    // Model you requested
 //   timestamp: 1703...      // When response was received

package/index.cjs

@@ -122,6 +122,12 @@ var AITransformer = class {
     this.logLevel = "info";
     this.lastResponseMetadata = null;
     this.exampleCount = 0;
+    this._cumulativeUsage = {
+      promptTokens: 0,
+      responseTokens: 0,
+      totalTokens: 0,
+      attempts: 0
+    };
     AITransformFactory.call(this, options);
     this.init = initChat.bind(this);
     this.seed = seedWithExamples.bind(this);
@@ -469,9 +475,21 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
   if (options.labels) {
     messageOptions.labels = options.labels;
   }
+  this._cumulativeUsage = {
+    promptTokens: 0,
+    responseTokens: 0,
+    totalTokens: 0,
+    attempts: 0
+  };
   for (let attempt = 0; attempt <= maxRetries; attempt++) {
     try {
       const transformedPayload = attempt === 0 ? await this.rawMessage(lastPayload, messageOptions) : await this.rebuild(lastPayload, lastError.message);
+      if (this.lastResponseMetadata) {
+        this._cumulativeUsage.promptTokens += this.lastResponseMetadata.promptTokens || 0;
+        this._cumulativeUsage.responseTokens += this.lastResponseMetadata.responseTokens || 0;
+        this._cumulativeUsage.totalTokens += this.lastResponseMetadata.totalTokens || 0;
+        this._cumulativeUsage.attempts = attempt + 1;
+      }
       lastPayload = transformedPayload;
       if (validatorFn) {
         await validatorFn(transformedPayload);
@@ -631,6 +649,13 @@ async function clearConversation() {
     },
     history: exampleHistory
   });
+  this.lastResponseMetadata = null;
+  this._cumulativeUsage = {
+    promptTokens: 0,
+    responseTokens: 0,
+    totalTokens: 0,
+    attempts: 0
+  };
   logger_default.debug(`Conversation cleared. Preserved ${exampleHistory.length} example items.`);
 }
 function getLastUsage() {
@@ -638,14 +663,15 @@ function getLastUsage() {
     return null;
   }
   const meta = this.lastResponseMetadata;
+  const cumulative = this._cumulativeUsage || { promptTokens: 0, responseTokens: 0, totalTokens: 0, attempts: 1 };
+  const useCumulative = cumulative.attempts > 0;
   return {
-    // Token breakdown for billing
-    promptTokens: meta.promptTokens,
-    // Input tokens (includes system instructions + history + message)
-    responseTokens: meta.responseTokens,
-    // Output tokens
-    totalTokens: meta.totalTokens,
-    // promptTokens + responseTokens
+    // Token breakdown for billing - CUMULATIVE across all retry attempts
+    promptTokens: useCumulative ? cumulative.promptTokens : meta.promptTokens,
+    responseTokens: useCumulative ? cumulative.responseTokens : meta.responseTokens,
+    totalTokens: useCumulative ? cumulative.totalTokens : meta.totalTokens,
+    // Number of attempts (1 = success on first try, 2+ = retries were needed)
+    attempts: useCumulative ? cumulative.attempts : 1,
     // Model verification for billing cross-check
     modelVersion: meta.modelVersion,
     // Actual model that responded (e.g., 'gemini-2.5-flash-001')
@@ -684,6 +710,12 @@ async function statelessMessage(sourcePayload, options = {}, validatorFn = null)
     totalTokens: result.usageMetadata?.totalTokenCount || 0,
     timestamp: Date.now()
   };
+  this._cumulativeUsage = {
+    promptTokens: this.lastResponseMetadata.promptTokens,
+    responseTokens: this.lastResponseMetadata.responseTokens,
+    totalTokens: this.lastResponseMetadata.totalTokens,
+    attempts: 1
+  };
   if (result.usageMetadata && logger_default.level !== "silent") {
     logger_default.debug(`Stateless message metadata:`, {
       modelVersion: result.modelVersion || "not-provided",

package/index.js

@@ -114,6 +114,13 @@ class AITransformer {
 		this.logLevel = 'info'; // default log level
 		this.lastResponseMetadata = null; // stores metadata from last API response
 		this.exampleCount = 0; // tracks number of example history items from seed()
+		// Cumulative usage tracking across retry attempts
+		this._cumulativeUsage = {
+			promptTokens: 0,
+			responseTokens: 0,
+			totalTokens: 0,
+			attempts: 0
+		};
 		AITransformFactory.call(this, options);
 
 		//external API
@@ -640,6 +647,14 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
 		messageOptions.labels = options.labels;
 	}
 
+	// Reset cumulative usage tracking for this message call
+	this._cumulativeUsage = {
+		promptTokens: 0,
+		responseTokens: 0,
+		totalTokens: 0,
+		attempts: 0
+	};
+
 	for (let attempt = 0; attempt <= maxRetries; attempt++) {
 		try {
 			// Step 1: Get the transformed payload
@@ -647,6 +662,14 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
 				? await this.rawMessage(lastPayload, messageOptions) // Use the new raw method with per-message options
 				: await this.rebuild(lastPayload, lastError.message);
 
+			// Accumulate token usage from this attempt
+			if (this.lastResponseMetadata) {
+				this._cumulativeUsage.promptTokens += this.lastResponseMetadata.promptTokens || 0;
+				this._cumulativeUsage.responseTokens += this.lastResponseMetadata.responseTokens || 0;
+				this._cumulativeUsage.totalTokens += this.lastResponseMetadata.totalTokens || 0;
+				this._cumulativeUsage.attempts = attempt + 1;
+			}
+
 			lastPayload = transformedPayload; // Always update lastPayload *before* validation
 
 			// Step 2: Validate if a validator is provided
@@ -915,15 +938,25 @@ async function clearConversation() {
 		history: exampleHistory,
 	});
 
+	// Reset usage tracking for the new conversation
+	this.lastResponseMetadata = null;
+	this._cumulativeUsage = {
+		promptTokens: 0,
+		responseTokens: 0,
+		totalTokens: 0,
+		attempts: 0
+	};
+
 	log.debug(`Conversation cleared. Preserved ${exampleHistory.length} example items.`);
 }
 
 /**
- * Returns structured usage data from the last API response for billing verification.
+ * Returns structured usage data from the last message call for billing verification.
+ * Includes CUMULATIVE token counts across all retry attempts.
  * Call this after message() or statelessMessage() to get actual token consumption.
  *
  * @this {ExportedAPI}
- * @returns {Object|null} Usage data with promptTokens, responseTokens, totalTokens, modelVersion, etc.
+ * @returns {Object|null} Usage data with promptTokens, responseTokens, totalTokens, attempts, etc.
  *                        Returns null if no API call has been made yet.
  */
 function getLastUsage() {
@@ -932,11 +965,19 @@ function getLastUsage() {
 	}
 
 	const meta = this.lastResponseMetadata;
+	const cumulative = this._cumulativeUsage || { promptTokens: 0, responseTokens: 0, totalTokens: 0, attempts: 1 };
+
+	// Use cumulative tokens if tracking was active (attempts > 0), otherwise fall back to last response
+	const useCumulative = cumulative.attempts > 0;
+
 	return {
-		// Token breakdown for billing
-		promptTokens: meta.promptTokens,      // Input tokens (includes system instructions + history + message)
-		responseTokens: meta.responseTokens,  // Output tokens
-		totalTokens: meta.totalTokens,        // promptTokens + responseTokens
+		// Token breakdown for billing - CUMULATIVE across all retry attempts
+		promptTokens: useCumulative ? cumulative.promptTokens : meta.promptTokens,
+		responseTokens: useCumulative ? cumulative.responseTokens : meta.responseTokens,
+		totalTokens: useCumulative ? cumulative.totalTokens : meta.totalTokens,
+
+		// Number of attempts (1 = success on first try, 2+ = retries were needed)
+		attempts: useCumulative ? cumulative.attempts : 1,
 
 		// Model verification for billing cross-check
 		modelVersion: meta.modelVersion,      // Actual model that responded (e.g., 'gemini-2.5-flash-001')
@@ -1001,6 +1042,14 @@ async function statelessMessage(sourcePayload, options = {}, validatorFn = null)
 		timestamp: Date.now()
 	};
 
+	// Set cumulative usage for stateless message (single attempt, no retries)
+	this._cumulativeUsage = {
+		promptTokens: this.lastResponseMetadata.promptTokens,
+		responseTokens: this.lastResponseMetadata.responseTokens,
+		totalTokens: this.lastResponseMetadata.totalTokens,
+		attempts: 1
+	};
+
 	if (result.usageMetadata && log.level !== 'silent') {
 		log.debug(`Stateless message metadata:`, {
 			modelVersion: result.modelVersion || 'not-provided',

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.13",
+	"version": "1.0.14",
 	"main": "index.js",
 	"files": [
 		"index.js",

package/types.d.ts

@@ -43,9 +43,10 @@ export interface ResponseMetadata {
 
 /** Structured usage data returned by getLastUsage() for billing verification */
 export interface UsageData {
-  promptTokens: number;       // Input tokens (includes system instructions + history + message)
-  responseTokens: number;     // Output tokens
-  totalTokens: number;        // promptTokens + responseTokens
+  promptTokens: number;       // CUMULATIVE input tokens across all retry attempts
+  responseTokens: number;     // CUMULATIVE output tokens across all retry attempts
+  totalTokens: number;        // CUMULATIVE total tokens across all retry attempts
+  attempts: number;           // Number of attempts (1 = first try success, 2+ = retries needed)
   modelVersion: string | null; // Actual model that responded (e.g., 'gemini-2.5-flash-001')
   requestedModel: string;     // Model you requested (e.g., 'gemini-2.5-flash')
   timestamp: number;          // When response was received
@@ -91,6 +92,7 @@ export interface AITransformerContext {
   lastResponseMetadata?: ResponseMetadata | null; // Metadata from the last API response
   exampleCount?: number; // Number of example history items from seed()
   clearConversation?: () => Promise<void>; // Clears conversation history while preserving examples
+  _cumulativeUsage?: { promptTokens: number; responseTokens: number; totalTokens: number; attempts: number }; // Internal cumulative tracking
 }
 
 export interface TransformationExample {