ak-gemini 1.0.10 → 1.0.12

package/README.md

@@ -11,6 +11,8 @@ Use this to power LLM-driven data pipelines, JSON mapping, or any automated AI t
 * **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.
+* **Conversation Management:** Clear conversation history while preserving examples, or send stateless one-off messages that don't affect history
+* **Response Metadata:** Access actual model version and token counts from API responses for billing verification and debugging
 * **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
@@ -106,6 +108,25 @@ console.log(validPayload);
 
 ---
 
+### 5. **Conversation Management**
+
+Manage chat history to control costs and isolate requests:
+
+```js
+// Clear conversation history while preserving seeded examples
+await transformer.clearConversation();
+
+// Send a stateless message that doesn't affect chat history
+const result = await transformer.message({ query: "one-off question" }, { stateless: true });
+
+// Check actual model and token usage from last API call
+console.log(transformer.lastResponseMetadata);
+// → { modelVersion: 'gemini-2.5-flash-001', requestedModel: 'gemini-2.5-flash',
+//    promptTokens: 150, responseTokens: 42, totalTokens: 192, timestamp: 1703... }
+```
+
+---
+
 ## API
 
 ### Constructor
@@ -142,10 +163,14 @@ Initializes Gemini chat session (idempotent).
 Seeds the model with example transformations (uses keys from constructor).
 You can omit `examples` to use the `examplesFile` (if provided).
 
-#### `await transformer.message(sourcePayload)`
+#### `await transformer.message(sourcePayload, options?)`
 
 Transforms input JSON to output JSON using the seeded examples and system instructions. Throws if estimated token window would be exceeded.
 
+**Options:**
+- `stateless: true` — Send a one-off message without affecting chat history (uses `generateContent` instead of chat)
+- `labels: {}` — Per-message billing labels
+
 #### `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).
@@ -168,6 +193,31 @@ Resets the Gemini chat session, clearing all history/examples.
 
 Returns the current chat history (for debugging).
 
+#### `await transformer.clearConversation()`
+
+Clears conversation history while preserving seeded examples. Useful for starting fresh user sessions without re-seeding.
+
+---
+
+### Properties
+
+#### `transformer.lastResponseMetadata`
+
+After each API call, contains metadata from the response:
+
+```js
+{
+  modelVersion: string | null,  // Actual model version that responded (e.g., 'gemini-2.5-flash-001')
+  requestedModel: string,       // Model you requested (e.g., 'gemini-2.5-flash')
+  promptTokens: number,         // Tokens in the prompt
+  responseTokens: number,       // Tokens in the response
+  totalTokens: number,          // Total tokens used
+  timestamp: number             // When response was received
+}
+```
+
+Useful for verifying billing, debugging model behavior, and tracking token usage.
+
 ---
 
 ## Examples

package/index.cjs

@@ -82,7 +82,7 @@ var DEFAULT_THINKING_CONFIG = {
   thinkingBudget: 0,
   thinkingLevel: import_genai.ThinkingLevel.MINIMAL
 };
-var DEFAULT_MAX_OUTPUT_TOKENS = 1e5;
+var DEFAULT_MAX_OUTPUT_TOKENS = 5e4;
 var THINKING_SUPPORTED_MODELS = [
   /^gemini-3-flash(-preview)?$/,
   /^gemini-3-pro(-preview|-image-preview)?$/,
@@ -120,6 +120,8 @@ var AITransformer = class {
     this.onlyJSON = true;
     this.asyncValidator = null;
     this.logLevel = "info";
+    this.lastResponseMetadata = null;
+    this.exampleCount = 0;
     AITransformFactory.call(this, options);
     this.init = initChat.bind(this);
     this.seed = seedWithExamples.bind(this);
@@ -135,6 +137,8 @@ var AITransformer = class {
     this.estimate = estimateTokenUsage.bind(this);
     this.estimateTokenUsage = estimateTokenUsage.bind(this);
     this.updateSystemInstructions = updateSystemInstructions.bind(this);
+    this.estimateCost = estimateCost.bind(this);
+    this.clearConversation = clearConversation.bind(this);
   }
 };
 var index_default = AITransformer;
@@ -164,8 +168,17 @@ function AITransformFactory(options = {}) {
     this.logLevel = "info";
     logger_default.level = "info";
   }
+  this.vertexai = options.vertexai || false;
+  this.project = options.project || process.env.GOOGLE_CLOUD_PROJECT || null;
+  this.location = options.location || process.env.GOOGLE_CLOUD_LOCATION || "us-central1";
+  this.googleAuthOptions = options.googleAuthOptions || null;
   this.apiKey = options.apiKey !== void 0 && 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.");
+  if (!this.vertexai && !this.apiKey) {
+    throw new Error("Missing Gemini API key. Provide via options.apiKey or GEMINI_API_KEY env var. For Vertex AI, set vertexai: true with project and location.");
+  }
+  if (this.vertexai && !this.project) {
+    throw new Error("Vertex AI requires a project ID. Provide via options.project or GOOGLE_CLOUD_PROJECT env var.");
+  }
   this.chatConfig = {
     ...DEFAULT_CHAT_CONFIG,
     ...options.chatConfig,
@@ -226,6 +239,10 @@ function AITransformFactory(options = {}) {
   this.onlyJSON = options.onlyJSON !== void 0 ? options.onlyJSON : true;
   this.enableGrounding = options.enableGrounding || false;
   this.groundingConfig = options.groundingConfig || {};
+  this.labels = options.labels || {};
+  if (Object.keys(this.labels).length > 0 && logger_default.level !== "silent") {
+    logger_default.debug(`Billing labels configured: ${JSON.stringify(this.labels)}`);
+  }
   if (this.promptKey === this.answerKey) {
     throw new Error("Source and target keys cannot be the same. Please provide distinct keys.");
   }
@@ -233,10 +250,27 @@ 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}"`);
     logger_default.debug(`Max output tokens set to: ${this.chatConfig.maxOutputTokens}`);
-    logger_default.debug(`Using API key: ${this.apiKey.substring(0, 10)}...`);
+    if (this.vertexai) {
+      logger_default.debug(`Using Vertex AI - Project: ${this.project}, Location: ${this.location}`);
+      if (this.googleAuthOptions?.keyFilename) {
+        logger_default.debug(`Auth: Service account key file: ${this.googleAuthOptions.keyFilename}`);
+      } else if (this.googleAuthOptions?.credentials) {
+        logger_default.debug(`Auth: Inline credentials provided`);
+      } else {
+        logger_default.debug(`Auth: Application Default Credentials (ADC)`);
+      }
+    } else {
+      logger_default.debug(`Using Gemini API with key: ${this.apiKey.substring(0, 10)}...`);
+    }
     logger_default.debug(`Grounding ${this.enableGrounding ? "ENABLED" : "DISABLED"} (costs $35/1k queries)`);
   }
-  const ai = new import_genai.GoogleGenAI({ apiKey: this.apiKey });
+  const clientOptions = this.vertexai ? {
+    vertexai: true,
+    project: this.project,
+    location: this.location,
+    ...this.googleAuthOptions && { googleAuthOptions: this.googleAuthOptions }
+  } : { apiKey: this.apiKey };
+  const ai = new import_genai.GoogleGenAI(clientOptions);
   this.genAIClient = ai;
   this.chat = null;
 }
@@ -246,7 +280,10 @@ async function initChat(force = false) {
   const chatOptions = {
     model: this.modelName,
     // @ts-ignore
-    config: this.chatConfig,
+    config: {
+      ...this.chatConfig,
+      ...Object.keys(this.labels).length > 0 && { labels: this.labels }
+    },
     history: []
   };
   if (this.enableGrounding) {
@@ -326,20 +363,47 @@ ${contextText}
   this.chat = await this.genAIClient.chats.create({
     model: this.modelName,
     // @ts-ignore
-    config: this.chatConfig,
+    config: {
+      ...this.chatConfig,
+      ...Object.keys(this.labels).length > 0 && { labels: this.labels }
+    },
     history: [...currentHistory, ...historyToAdd]
   });
+  this.exampleCount = currentHistory.length + historyToAdd.length;
   const newHistory = this.chat.getHistory();
   logger_default.debug(`Created new chat session with ${newHistory.length} examples.`);
   return newHistory;
 }
-async function rawMessage(sourcePayload) {
+async function rawMessage(sourcePayload, messageOptions = {}) {
   if (!this.chat) {
     throw new Error("Chat session not initialized.");
   }
   const actualPayload = typeof sourcePayload === "string" ? sourcePayload : JSON.stringify(sourcePayload, null, 2);
+  const mergedLabels = { ...this.labels, ...messageOptions.labels || {} };
+  const hasLabels = Object.keys(mergedLabels).length > 0;
   try {
-    const result = await this.chat.sendMessage({ message: actualPayload });
+    const sendParams = { message: actualPayload };
+    if (hasLabels) {
+      sendParams.config = { labels: mergedLabels };
+    }
+    const result = await this.chat.sendMessage(sendParams);
+    this.lastResponseMetadata = {
+      modelVersion: result.modelVersion || null,
+      requestedModel: this.modelName,
+      promptTokens: result.usageMetadata?.promptTokenCount || 0,
+      responseTokens: result.usageMetadata?.candidatesTokenCount || 0,
+      totalTokens: result.usageMetadata?.totalTokenCount || 0,
+      timestamp: Date.now()
+    };
+    if (result.usageMetadata && logger_default.level !== "silent") {
+      logger_default.debug(`API response metadata:`, {
+        modelVersion: result.modelVersion || "not-provided",
+        requestedModel: this.modelName,
+        promptTokens: result.usageMetadata.promptTokenCount,
+        responseTokens: result.usageMetadata.candidatesTokenCount,
+        totalTokens: result.usageMetadata.totalTokenCount
+      });
+    }
     const modelResponse = result.text;
     const extractedJSON = extractJSON(modelResponse);
     if (extractedJSON?.data) {
@@ -357,6 +421,9 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
   if (!this.chat) {
     throw new Error("Chat session not initialized. Please call init() first.");
   }
+  if (options.stateless) {
+    return await statelessMessage.call(this, sourcePayload, options, validatorFn);
+  }
   const maxRetries = options.maxRetries ?? this.maxRetries;
   const retryDelay = options.retryDelay ?? this.retryDelay;
   const enableGroundingForMessage = options.enableGrounding ?? this.enableGrounding;
@@ -398,9 +465,13 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
   } else {
     throw new Error("Invalid source payload. Must be a JSON object or string.");
   }
+  const messageOptions = {};
+  if (options.labels) {
+    messageOptions.labels = options.labels;
+  }
   for (let attempt = 0; attempt <= maxRetries; attempt++) {
     try {
-      const transformedPayload = attempt === 0 ? await this.rawMessage(lastPayload) : await this.rebuild(lastPayload, lastError.message);
+      const transformedPayload = attempt === 0 ? await this.rawMessage(lastPayload, messageOptions) : await this.rebuild(lastPayload, lastError.message);
       lastPayload = transformedPayload;
       if (validatorFn) {
         await validatorFn(transformedPayload);
@@ -444,6 +515,17 @@ Respond with JSON only \u2013 no comments or explanations.
   let result;
   try {
     result = await this.chat.sendMessage({ message: prompt });
+    this.lastResponseMetadata = {
+      modelVersion: result.modelVersion || null,
+      requestedModel: this.modelName,
+      promptTokens: result.usageMetadata?.promptTokenCount || 0,
+      responseTokens: result.usageMetadata?.candidatesTokenCount || 0,
+      totalTokens: result.usageMetadata?.totalTokenCount || 0,
+      timestamp: Date.now()
+    };
+    if (result.usageMetadata && logger_default.level !== "silent") {
+      logger_default.debug(`Rebuild response metadata - tokens used:`, result.usageMetadata.totalTokenCount);
+    }
   } catch (err) {
     throw new Error(`Gemini call failed while repairing payload: ${err.message}`);
   }
@@ -473,13 +555,36 @@ async function estimateTokenUsage(nextPayload) {
   });
   return resp;
 }
+var MODEL_PRICING = {
+  "gemini-2.5-flash": { input: 0.15, output: 0.6 },
+  "gemini-2.5-flash-lite": { input: 0.02, output: 0.1 },
+  "gemini-2.5-pro": { input: 2.5, output: 10 },
+  "gemini-3-pro": { input: 2, output: 12 },
+  "gemini-3-pro-preview": { input: 2, output: 12 },
+  "gemini-2.0-flash": { input: 0.1, output: 0.4 },
+  "gemini-2.0-flash-lite": { input: 0.02, output: 0.1 }
+};
+async function estimateCost(nextPayload) {
+  const tokenInfo = await this.estimateTokenUsage(nextPayload);
+  const pricing = MODEL_PRICING[this.modelName] || { input: 0, output: 0 };
+  return {
+    totalTokens: tokenInfo.totalTokens,
+    model: this.modelName,
+    pricing,
+    estimatedInputCost: tokenInfo.totalTokens / 1e6 * pricing.input,
+    note: "Cost is for input tokens only; output cost depends on response length"
+  };
+}
 async function resetChat() {
   if (this.chat) {
     logger_default.debug("Resetting Gemini chat session...");
     const chatOptions = {
       model: this.modelName,
       // @ts-ignore
-      config: this.chatConfig,
+      config: {
+        ...this.chatConfig,
+        ...Object.keys(this.labels).length > 0 && { labels: this.labels }
+      },
       history: []
     };
     if (this.enableGrounding) {
@@ -510,6 +615,68 @@ async function updateSystemInstructions(newInstructions) {
   logger_default.debug("Updating system instructions and reinitializing chat...");
   await this.init(true);
 }
+async function clearConversation() {
+  if (!this.chat) {
+    logger_default.warn("Cannot clear conversation: chat not initialized.");
+    return;
+  }
+  const history = this.chat.getHistory();
+  const exampleHistory = history.slice(0, this.exampleCount || 0);
+  this.chat = await this.genAIClient.chats.create({
+    model: this.modelName,
+    // @ts-ignore
+    config: {
+      ...this.chatConfig,
+      ...Object.keys(this.labels).length > 0 && { labels: this.labels }
+    },
+    history: exampleHistory
+  });
+  logger_default.debug(`Conversation cleared. Preserved ${exampleHistory.length} example items.`);
+}
+async function statelessMessage(sourcePayload, options = {}, validatorFn = null) {
+  if (!this.chat) {
+    throw new Error("Chat session not initialized. Please call init() first.");
+  }
+  const payloadStr = typeof sourcePayload === "string" ? sourcePayload : JSON.stringify(sourcePayload, null, 2);
+  const contents = [];
+  if (this.exampleCount > 0) {
+    const history = this.chat.getHistory();
+    const exampleHistory = history.slice(0, this.exampleCount);
+    contents.push(...exampleHistory);
+  }
+  contents.push({ role: "user", parts: [{ text: payloadStr }] });
+  const mergedLabels = { ...this.labels, ...options.labels || {} };
+  const result = await this.genAIClient.models.generateContent({
+    model: this.modelName,
+    contents,
+    config: {
+      ...this.chatConfig,
+      ...Object.keys(mergedLabels).length > 0 && { labels: mergedLabels }
+    }
+  });
+  this.lastResponseMetadata = {
+    modelVersion: result.modelVersion || null,
+    requestedModel: this.modelName,
+    promptTokens: result.usageMetadata?.promptTokenCount || 0,
+    responseTokens: result.usageMetadata?.candidatesTokenCount || 0,
+    totalTokens: result.usageMetadata?.totalTokenCount || 0,
+    timestamp: Date.now()
+  };
+  if (result.usageMetadata && logger_default.level !== "silent") {
+    logger_default.debug(`Stateless message metadata:`, {
+      modelVersion: result.modelVersion || "not-provided",
+      promptTokens: result.usageMetadata.promptTokenCount,
+      responseTokens: result.usageMetadata.candidatesTokenCount
+    });
+  }
+  const modelResponse = result.text;
+  const extractedJSON = extractJSON(modelResponse);
+  let transformedPayload = extractedJSON?.data ? extractedJSON.data : extractedJSON;
+  if (validatorFn) {
+    await validatorFn(transformedPayload);
+  }
+  return transformedPayload;
+}
 function attemptJSONRecovery(text, maxAttempts = 100) {
   if (!text || typeof text !== "string") return null;
   try {

package/index.js

@@ -57,7 +57,7 @@ const DEFAULT_THINKING_CONFIG = {
 	thinkingLevel: ThinkingLevel.MINIMAL
 };
 
-const DEFAULT_MAX_OUTPUT_TOKENS = 100000; // Default ceiling for output tokens
+const DEFAULT_MAX_OUTPUT_TOKENS = 50_000; // Default ceiling for output tokens
 
 // Models that support thinking features (as of Dec 2024)
 // Using regex patterns for more precise matching
@@ -112,6 +112,8 @@ class AITransformer {
 		this.onlyJSON = true; // always return JSON
 		this.asyncValidator = null; // for transformWithValidation
 		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()
 		AITransformFactory.call(this, options);
 
 		//external API
@@ -135,6 +137,8 @@ class AITransformer {
 		this.estimate = estimateTokenUsage.bind(this);
 		this.estimateTokenUsage = estimateTokenUsage.bind(this);
 		this.updateSystemInstructions = updateSystemInstructions.bind(this);
+		this.estimateCost = estimateCost.bind(this);
+		this.clearConversation = clearConversation.bind(this);
 	}
 }
 
@@ -180,8 +184,22 @@ function AITransformFactory(options = {}) {
 		log.level = 'info';
 	}
 
+	// Vertex AI configuration
+	this.vertexai = options.vertexai || false;
+	this.project = options.project || process.env.GOOGLE_CLOUD_PROJECT || null;
+	this.location = options.location || process.env.GOOGLE_CLOUD_LOCATION || 'us-central1';
+	this.googleAuthOptions = options.googleAuthOptions || null;
+
+	// API Key (for Gemini API, not Vertex AI)
 	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.");
+
+	// Validate authentication - need either API key (for Gemini API) or Vertex AI config
+	if (!this.vertexai && !this.apiKey) {
+		throw new Error("Missing Gemini API key. Provide via options.apiKey or GEMINI_API_KEY env var. For Vertex AI, set vertexai: true with project and location.");
+	}
+	if (this.vertexai && !this.project) {
+		throw new Error("Vertex AI requires a project ID. Provide via options.project or GOOGLE_CLOUD_PROJECT env var.");
+	}
 
 	// Build chat config, making sure systemInstruction uses the custom instructions
 	this.chatConfig = {
@@ -270,6 +288,12 @@ function AITransformFactory(options = {}) {
 	this.enableGrounding = options.enableGrounding || false;
 	this.groundingConfig = options.groundingConfig || {};
 
+	// Billing labels for cost segmentation
+	this.labels = options.labels || {};
+	if (Object.keys(this.labels).length > 0 && log.level !== 'silent') {
+		log.debug(`Billing labels configured: ${JSON.stringify(this.labels)}`);
+	}
+
 	if (this.promptKey === this.answerKey) {
 		throw new Error("Source and target keys cannot be the same. Please provide distinct keys.");
 	}
@@ -278,12 +302,33 @@ 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}"`);
 		log.debug(`Max output tokens set to: ${this.chatConfig.maxOutputTokens}`);
-		// Log API key prefix for tracking (first 10 chars only for security)
-		log.debug(`Using API key: ${this.apiKey.substring(0, 10)}...`);
+		// Log authentication method
+		if (this.vertexai) {
+			log.debug(`Using Vertex AI - Project: ${this.project}, Location: ${this.location}`);
+			if (this.googleAuthOptions?.keyFilename) {
+				log.debug(`Auth: Service account key file: ${this.googleAuthOptions.keyFilename}`);
+			} else if (this.googleAuthOptions?.credentials) {
+				log.debug(`Auth: Inline credentials provided`);
+			} else {
+				log.debug(`Auth: Application Default Credentials (ADC)`);
+			}
+		} else {
+			log.debug(`Using Gemini API with key: ${this.apiKey.substring(0, 10)}...`);
+		}
 		log.debug(`Grounding ${this.enableGrounding ? 'ENABLED' : 'DISABLED'} (costs $35/1k queries)`);
 	}
 
-	const ai = new GoogleGenAI({ apiKey: this.apiKey });
+	// Initialize Google GenAI client with appropriate configuration
+	const clientOptions = this.vertexai
+		? {
+			vertexai: true,
+			project: this.project,
+			location: this.location,
+			...(this.googleAuthOptions && { googleAuthOptions: this.googleAuthOptions })
+		}
+		: { apiKey: this.apiKey };
+
+	const ai = new GoogleGenAI(clientOptions);
 	this.genAIClient = ai;
 	this.chat = null;
 }
@@ -303,7 +348,10 @@ async function initChat(force = false) {
 	const chatOptions = {
 		model: this.modelName,
 		// @ts-ignore
-		config: this.chatConfig,
+		config: {
+			...this.chatConfig,
+			...(Object.keys(this.labels).length > 0 && { labels: this.labels })
+		},
 		history: [],
 	};
 
@@ -415,10 +463,15 @@ async function seedWithExamples(examples) {
 	this.chat = await this.genAIClient.chats.create({
 		model: this.modelName,
 		// @ts-ignore
-		config: this.chatConfig,
+		config: {
+			...this.chatConfig,
+			...(Object.keys(this.labels).length > 0 && { labels: this.labels })
+		},
 		history: [...currentHistory, ...historyToAdd],
 	});
 
+	// Track example count for clearConversation() and stateless messages
+	this.exampleCount = currentHistory.length + historyToAdd.length;
 
 	const newHistory = this.chat.getHistory();
 	log.debug(`Created new chat session with ${newHistory.length} examples.`);
@@ -436,9 +489,10 @@ async function seedWithExamples(examples) {
  * No validation or retry logic.
  * @this {ExportedAPI}
  * @param {Object|string} sourcePayload - The source payload.
+ * @param {Object} [messageOptions] - Optional per-message options (e.g., labels).
  * @returns {Promise<Object>} - The transformed payload.
  */
-async function rawMessage(sourcePayload) {
+async function rawMessage(sourcePayload, messageOptions = {}) {
 	if (!this.chat) {
 		throw new Error("Chat session not initialized.");
 	}
@@ -447,8 +501,40 @@ async function rawMessage(sourcePayload) {
 		? sourcePayload
 		: JSON.stringify(sourcePayload, null, 2);
 
+	// Merge instance labels with per-message labels (per-message takes precedence)
+	const mergedLabels = { ...this.labels, ...(messageOptions.labels || {}) };
+	const hasLabels = Object.keys(mergedLabels).length > 0;
+
 	try {
-		const result = await this.chat.sendMessage({ message: actualPayload });
+		const sendParams = { message: actualPayload };
+
+		// Add config with labels if we have any
+		if (hasLabels) {
+			sendParams.config = { labels: mergedLabels };
+		}
+
+		const result = await this.chat.sendMessage(sendParams);
+
+		// Capture and log response metadata for model verification and debugging
+		this.lastResponseMetadata = {
+			modelVersion: result.modelVersion || null,
+			requestedModel: this.modelName,
+			promptTokens: result.usageMetadata?.promptTokenCount || 0,
+			responseTokens: result.usageMetadata?.candidatesTokenCount || 0,
+			totalTokens: result.usageMetadata?.totalTokenCount || 0,
+			timestamp: Date.now()
+		};
+
+		if (result.usageMetadata && log.level !== 'silent') {
+			log.debug(`API response metadata:`, {
+				modelVersion: result.modelVersion || 'not-provided',
+				requestedModel: this.modelName,
+				promptTokens: result.usageMetadata.promptTokenCount,
+				responseTokens: result.usageMetadata.candidatesTokenCount,
+				totalTokens: result.usageMetadata.totalTokenCount
+			});
+		}
+
 		const modelResponse = result.text;
 		const extractedJSON = extractJSON(modelResponse); // Assuming extractJSON is defined
 
@@ -479,6 +565,12 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
 	if (!this.chat) {
 		throw new Error("Chat session not initialized. Please call init() first.");
 	}
+
+	// Handle stateless messages separately - they don't add to chat history
+	if (options.stateless) {
+		return await statelessMessage.call(this, sourcePayload, options, validatorFn);
+	}
+
 	const maxRetries = options.maxRetries ?? this.maxRetries;
 	const retryDelay = options.retryDelay ?? this.retryDelay;
 
@@ -542,11 +634,17 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
 		throw new Error("Invalid source payload. Must be a JSON object or string.");
 	}
 
+	// Extract per-message labels for passing to rawMessage
+	const messageOptions = {};
+	if (options.labels) {
+		messageOptions.labels = options.labels;
+	}
+
 	for (let attempt = 0; attempt <= maxRetries; attempt++) {
 		try {
 			// Step 1: Get the transformed payload
 			const transformedPayload = (attempt === 0)
-				? await this.rawMessage(lastPayload) // Use the new raw method
+				? await this.rawMessage(lastPayload, messageOptions) // Use the new raw method with per-message options
 				: await this.rebuild(lastPayload, lastError.message);
 
 			lastPayload = transformedPayload; // Always update lastPayload *before* validation
@@ -590,6 +688,7 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
 
 /**
  * Rebuilds a payload based on server error feedback
+ * @this {ExportedAPI}
  * @param {Object} lastPayload - The payload that failed validation
  * @param {string} serverError - The error message from the server
  * @returns {Promise<Object>} - A new corrected payload
@@ -615,6 +714,20 @@ Respond with JSON only – no comments or explanations.
 	let result;
 	try {
 		result = await this.chat.sendMessage({ message: prompt });
+
+		// Capture and log response metadata for rebuild calls too
+		this.lastResponseMetadata = {
+			modelVersion: result.modelVersion || null,
+			requestedModel: this.modelName,
+			promptTokens: result.usageMetadata?.promptTokenCount || 0,
+			responseTokens: result.usageMetadata?.candidatesTokenCount || 0,
+			totalTokens: result.usageMetadata?.totalTokenCount || 0,
+			timestamp: Date.now()
+		};
+
+		if (result.usageMetadata && log.level !== 'silent') {
+			log.debug(`Rebuild response metadata - tokens used:`, result.usageMetadata.totalTokenCount);
+		}
 	} catch (err) {
 		throw new Error(`Gemini call failed while repairing payload: ${err.message}`);
 	}
@@ -670,6 +783,37 @@ async function estimateTokenUsage(nextPayload) {
 	return resp; // includes totalTokens, possibly breakdown
 }
 
+// Model pricing per million tokens (as of Dec 2025)
+// https://ai.google.dev/gemini-api/docs/pricing
+const MODEL_PRICING = {
+	'gemini-2.5-flash': { input: 0.15, output: 0.60 },
+	'gemini-2.5-flash-lite': { input: 0.02, output: 0.10 },
+	'gemini-2.5-pro': { input: 2.50, output: 10.00 },
+	'gemini-3-pro': { input: 2.00, output: 12.00 },
+	'gemini-3-pro-preview': { input: 2.00, output: 12.00 },
+	'gemini-2.0-flash': { input: 0.10, output: 0.40 },
+	'gemini-2.0-flash-lite': { input: 0.02, output: 0.10 }
+};
+
+/**
+ * Estimates the cost of sending a payload based on token count and model pricing.
+ * @this {ExportedAPI}
+ * @param {object|string} nextPayload - The next user message to be sent (object or string)
+ * @returns {Promise<Object>} - Cost estimation including tokens, model, pricing, and estimated input cost
+ */
+async function estimateCost(nextPayload) {
+	const tokenInfo = await this.estimateTokenUsage(nextPayload);
+	const pricing = MODEL_PRICING[this.modelName] || { input: 0, output: 0 };
+
+	return {
+		totalTokens: tokenInfo.totalTokens,
+		model: this.modelName,
+		pricing: pricing,
+		estimatedInputCost: (tokenInfo.totalTokens / 1_000_000) * pricing.input,
+		note: 'Cost is for input tokens only; output cost depends on response length'
+	};
+}
+
 
 /**
  * Resets the current chat session, clearing all history and examples
@@ -684,7 +828,10 @@ async function resetChat() {
 		const chatOptions = {
 			model: this.modelName,
 			// @ts-ignore
-			config: this.chatConfig,
+			config: {
+				...this.chatConfig,
+				...(Object.keys(this.labels).length > 0 && { labels: this.labels })
+			},
 			history: [],
 		};
 
@@ -733,6 +880,110 @@ async function updateSystemInstructions(newInstructions) {
 	await this.init(true); // Force reinitialize with new instructions
 }
 
+/**
+ * Clears conversation history while preserving seeded examples.
+ * Useful for starting a fresh conversation within the same session
+ * without losing the few-shot learning examples.
+ * @this {ExportedAPI}
+ * @returns {Promise<void>}
+ */
+async function clearConversation() {
+	if (!this.chat) {
+		log.warn("Cannot clear conversation: chat not initialized.");
+		return;
+	}
+
+	const history = this.chat.getHistory();
+	const exampleHistory = history.slice(0, this.exampleCount || 0);
+
+	this.chat = await this.genAIClient.chats.create({
+		model: this.modelName,
+		// @ts-ignore
+		config: {
+			...this.chatConfig,
+			...(Object.keys(this.labels).length > 0 && { labels: this.labels })
+		},
+		history: exampleHistory,
+	});
+
+	log.debug(`Conversation cleared. Preserved ${exampleHistory.length} example items.`);
+}
+
+/**
+ * Sends a one-off message using generateContent (not chat).
+ * Does NOT affect chat history - useful for isolated requests.
+ * @this {ExportedAPI}
+ * @param {Object|string} sourcePayload - The source payload.
+ * @param {Object} [options] - Options including labels.
+ * @param {AsyncValidatorFunction|null} [validatorFn] - Optional validator.
+ * @returns {Promise<Object>} - The transformed payload.
+ */
+async function statelessMessage(sourcePayload, options = {}, validatorFn = null) {
+	if (!this.chat) {
+		throw new Error("Chat session not initialized. Please call init() first.");
+	}
+
+	const payloadStr = typeof sourcePayload === 'string'
+		? sourcePayload
+		: JSON.stringify(sourcePayload, null, 2);
+
+	// Build contents including examples from current chat history
+	const contents = [];
+
+	// Include seeded examples if we have them
+	if (this.exampleCount > 0) {
+		const history = this.chat.getHistory();
+		const exampleHistory = history.slice(0, this.exampleCount);
+		contents.push(...exampleHistory);
+	}
+
+	// Add the user message
+	contents.push({ role: 'user', parts: [{ text: payloadStr }] });
+
+	// Merge labels
+	const mergedLabels = { ...this.labels, ...(options.labels || {}) };
+
+	// Use generateContent instead of chat.sendMessage
+	const result = await this.genAIClient.models.generateContent({
+		model: this.modelName,
+		contents: contents,
+		config: {
+			...this.chatConfig,
+			...(Object.keys(mergedLabels).length > 0 && { labels: mergedLabels })
+		}
+	});
+
+	// Capture and log response metadata
+	this.lastResponseMetadata = {
+		modelVersion: result.modelVersion || null,
+		requestedModel: this.modelName,
+		promptTokens: result.usageMetadata?.promptTokenCount || 0,
+		responseTokens: result.usageMetadata?.candidatesTokenCount || 0,
+		totalTokens: result.usageMetadata?.totalTokenCount || 0,
+		timestamp: Date.now()
+	};
+
+	if (result.usageMetadata && log.level !== 'silent') {
+		log.debug(`Stateless message metadata:`, {
+			modelVersion: result.modelVersion || 'not-provided',
+			promptTokens: result.usageMetadata.promptTokenCount,
+			responseTokens: result.usageMetadata.candidatesTokenCount
+		});
+	}
+
+	const modelResponse = result.text;
+	const extractedJSON = extractJSON(modelResponse);
+
+	let transformedPayload = extractedJSON?.data ? extractedJSON.data : extractedJSON;
+
+	// Validate if a validator is provided
+	if (validatorFn) {
+		await validatorFn(transformedPayload);
+	}
+
+	return transformedPayload;
+}
+
 
 /*
 ----

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

package/types.d.ts

@@ -26,9 +26,31 @@ export interface ChatConfig {
   safetySettings?: SafetySetting[]; // Safety settings array
   responseSchema?: Object; // Schema for validating model responses
   thinkingConfig?: ThinkingConfig; // Thinking features configuration
+  labels?: Record<string, string>; // Labels for billing segmentation
+  tools?: any[]; // Tools configuration (e.g., grounding)
   [key: string]: any; // Additional properties for flexibility
 }
 
+/** Metadata from the last API response, useful for debugging and cost tracking */
+export interface ResponseMetadata {
+  modelVersion: string | null; // The actual model version that responded
+  requestedModel: string; // The model that was requested
+  promptTokens: number; // Number of tokens in the prompt
+  responseTokens: number; // Number of tokens in the response
+  totalTokens: number; // Total tokens used
+  timestamp: number; // Timestamp of when the response was received
+}
+
+/** Options for per-message configuration */
+export interface MessageOptions {
+  labels?: Record<string, string>; // Per-message billing labels
+  stateless?: boolean; // If true, send message without affecting chat history
+  maxRetries?: number; // Override max retries for this message
+  retryDelay?: number; // Override retry delay for this message
+  enableGrounding?: boolean; // Override grounding setting for this message
+  groundingConfig?: Record<string, any>; // Override grounding config for this message
+}
+
 export interface AITransformerContext {
   modelName?: string;
   systemInstructions?: string;
@@ -45,15 +67,19 @@ export interface AITransformerContext {
   maxRetries?: number;
   retryDelay?: number;
   init?: (force?: boolean) => Promise<void>; // Initialization function
-  seed?: () => Promise<void>; // Function to seed the transformer with examples  
-  message?: (payload: Record<string, unknown>) => Promise<Record<string, unknown>>; // Function to send messages to the model
+  seed?: () => Promise<void>; // Function to seed the transformer with examples
+  message?: (payload: Record<string, unknown>, opts?: MessageOptions, validatorFn?: AsyncValidatorFunction | null) => Promise<Record<string, unknown>>; // Function to send messages to the model
   rebuild?: (lastPayload: Record<string, unknown>, serverError: string) => Promise<Record<string, unknown>>; // Function to rebuild the transformer
-  rawMessage?: (payload: Record<string, unknown> | string) => Promise<Record<string, unknown>>; // Function to send raw messages to the model
+  rawMessage?: (payload: Record<string, unknown> | string, messageOptions?: { labels?: Record<string, string> }) => Promise<Record<string, unknown>>; // Function to send raw messages to the model
   genAIClient?: GoogleGenAI; // Google GenAI client instance
   onlyJSON?: boolean; // If true, only JSON responses are allowed
   enableGrounding?: boolean; // Enable Google Search grounding (default: false, WARNING: costs $35/1k queries)
   groundingConfig?: Record<string, any>; // Additional grounding configuration options
-
+  labels?: Record<string, string>; // Custom labels for billing segmentation (keys: 1-63 chars lowercase, values: max 63 chars)
+  estimateTokenUsage?: (nextPayload: Record<string, unknown> | string) => Promise<{ totalTokens: number; breakdown?: any }>;
+  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
 }
 
 export interface TransformationExample {
@@ -71,13 +97,24 @@ export interface ExampleFileContent {
   examples: TransformationExample[];
 }
 
+// Google Auth options for Vertex AI authentication
+// See: https://github.com/googleapis/google-auth-library-nodejs/blob/main/src/auth/googleauth.ts
+export interface GoogleAuthOptions {
+  keyFilename?: string; // Path to a .json, .pem, or .p12 key file
+  keyFile?: string; // Alias for keyFilename
+  credentials?: { client_email?: string; private_key?: string; [key: string]: any }; // Object containing client_email and private_key
+  scopes?: string | string[]; // Required scopes for the API request
+  projectId?: string; // Your project ID (alias for project)
+  universeDomain?: string; // The default service domain for a Cloud universe
+}
+
 export interface AITransformerOptions {
 	// ? https://ai.google.dev/gemini-api/docs/models
   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)
+  maxOutputTokens?: number; // Maximum number of tokens that can be generated in the response (defaults to 50000)
   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)
@@ -91,12 +128,20 @@ export interface AITransformerOptions {
   retryDelay?: number; // Initial retry delay in milliseconds
   // ? https://ai.google.dev/gemini-api/docs/structured-output
   responseSchema?: Object; // Schema for validating model responses
-  apiKey?: string; // API key for Google GenAI
+  apiKey?: string; // API key for Google GenAI (Gemini API)
   onlyJSON?: boolean; // If true, only JSON responses are allowed
   asyncValidator?: AsyncValidatorFunction; // Optional async validator function for response validation
   logLevel?: 'trace' | 'debug' | 'info' | 'warn' | 'error' | 'fatal' | 'none'; // Log level for the logger (defaults to 'info', 'none' disables logging)
   enableGrounding?: boolean; // Enable Google Search grounding (default: false, WARNING: costs $35/1k queries)
   groundingConfig?: Record<string, any>; // Additional grounding configuration options
+  labels?: Record<string, string>; // Custom labels for billing segmentation
+
+  // Vertex AI Authentication Options
+  // Use these instead of apiKey for Vertex AI with service account authentication
+  vertexai?: boolean; // Set to true to use Vertex AI instead of Gemini API
+  project?: string; // Google Cloud project ID (required for Vertex AI)
+  location?: string; // Google Cloud location/region (e.g., 'us-central1') - required for Vertex AI
+  googleAuthOptions?: GoogleAuthOptions; // Authentication options for Vertex AI (keyFilename, credentials, etc.)
 }
 
 // Async validator function type
@@ -126,20 +171,40 @@ export declare class AITransformer {
   logLevel: 'trace' | 'debug' | 'info' | 'warn' | 'error' | 'fatal' | 'none';
   enableGrounding: boolean;
   groundingConfig: Record<string, any>;
-  
+  labels: Record<string, string>;
+  /** Metadata from the last API response (model version, token counts, etc.) */
+  lastResponseMetadata: ResponseMetadata | null;
+  /** Number of history items that are seeded examples (used by clearConversation) */
+  exampleCount: number;
+
   // Methods
   init(force?: boolean): Promise<void>;
   seed(examples?: TransformationExample[]): Promise<any>;
-  message(payload: Record<string, unknown>, opts?: object, validatorFn?: AsyncValidatorFunction | null): Promise<Record<string, unknown>>;
-  rawMessage(sourcePayload: Record<string, unknown> | string): Promise<Record<string, unknown> | any>;
-  transformWithValidation(sourcePayload: Record<string, unknown>, validatorFn: AsyncValidatorFunction, options?: object): Promise<Record<string, unknown>>;
-  messageAndValidate(sourcePayload: Record<string, unknown>, validatorFn: AsyncValidatorFunction, options?: object): Promise<Record<string, unknown>>;
+  /**
+   * Send a message to the model.
+   * @param payload - The payload to transform
+   * @param opts - Options including { stateless: true } to send without affecting history
+   * @param validatorFn - Optional validator function
+   */
+  message(payload: Record<string, unknown>, opts?: MessageOptions, validatorFn?: AsyncValidatorFunction | null): Promise<Record<string, unknown>>;
+  rawMessage(sourcePayload: Record<string, unknown> | string, messageOptions?: { labels?: Record<string, string> }): Promise<Record<string, unknown> | any>;
+  transformWithValidation(sourcePayload: Record<string, unknown>, validatorFn: AsyncValidatorFunction, options?: MessageOptions): Promise<Record<string, unknown>>;
+  messageAndValidate(sourcePayload: Record<string, unknown>, validatorFn: AsyncValidatorFunction, options?: MessageOptions): Promise<Record<string, unknown>>;
   rebuild(lastPayload: Record<string, unknown>, serverError: string): Promise<Record<string, unknown>>;
   reset(): Promise<void>;
   getHistory(): Array<any>;
   estimateTokenUsage(nextPayload: Record<string, unknown> | string): Promise<{ totalTokens: number; breakdown?: any }>;
   estimate(nextPayload: Record<string, unknown> | string): Promise<{ totalTokens: number; breakdown?: any }>;
   updateSystemInstructions(newInstructions: string): Promise<void>;
+  estimateCost(nextPayload: Record<string, unknown> | string): Promise<{
+    totalTokens: number;
+    model: string;
+    pricing: { input: number; output: number };
+    estimatedInputCost: number;
+    note: string;
+  }>;
+  /** Clears conversation history while preserving seeded examples */
+  clearConversation(): Promise<void>;
 }
 
 // Default export