ak-gemini 1.0.10 → 1.0.11

package/index.cjs

@@ -135,6 +135,7 @@ var AITransformer = class {
     this.estimate = estimateTokenUsage.bind(this);
     this.estimateTokenUsage = estimateTokenUsage.bind(this);
     this.updateSystemInstructions = updateSystemInstructions.bind(this);
+    this.estimateCost = estimateCost.bind(this);
   }
 };
 var index_default = AITransformer;
@@ -164,8 +165,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 +236,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 +247,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 +277,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 +360,29 @@ ${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]
   });
   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);
     const modelResponse = result.text;
     const extractedJSON = extractJSON(modelResponse);
     if (extractedJSON?.data) {
@@ -398,9 +441,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);
@@ -473,13 +520,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) {

package/index.js

@@ -135,6 +135,7 @@ class AITransformer {
 		this.estimate = estimateTokenUsage.bind(this);
 		this.estimateTokenUsage = estimateTokenUsage.bind(this);
 		this.updateSystemInstructions = updateSystemInstructions.bind(this);
+		this.estimateCost = estimateCost.bind(this);
 	}
 }
 
@@ -180,8 +181,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 +285,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 +299,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 +345,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,7 +460,10 @@ 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],
 	});
 
@@ -436,9 +484,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 +496,19 @@ 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);
 		const modelResponse = result.text;
 		const extractedJSON = extractJSON(modelResponse); // Assuming extractJSON is defined
 
@@ -542,11 +602,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
@@ -670,6 +736,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 +781,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: [],
 		};
 

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.11",
 	"main": "index.js",
 	"files": [
 		"index.js",

package/types.d.ts

@@ -26,6 +26,8 @@ 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
 }
 
@@ -48,11 +50,13 @@ export interface AITransformerContext {
   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
   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 }>;
 
 }
 
@@ -71,6 +75,17 @@ 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
@@ -91,12 +106,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,12 +149,13 @@ export declare class AITransformer {
   logLevel: 'trace' | 'debug' | 'info' | 'warn' | 'error' | 'fatal' | 'none';
   enableGrounding: boolean;
   groundingConfig: Record<string, any>;
+  labels: Record<string, string>;
   
   // 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>;
+  rawMessage(sourcePayload: Record<string, unknown> | string, messageOptions?: { labels?: Record<string, 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>>;
   rebuild(lastPayload: Record<string, unknown>, serverError: string): Promise<Record<string, unknown>>;
@@ -140,6 +164,13 @@ export declare class AITransformer {
   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;
+  }>;
 }
 
 // Default export