ak-gemini 1.1.0 → 1.1.1

package/index.cjs

@@ -178,7 +178,7 @@ function AITransformFactory(options = {}) {
   }
   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.location = options.location || process.env.GOOGLE_CLOUD_LOCATION || void 0;
   this.googleAuthOptions = options.googleAuthOptions || null;
   this.apiKey = options.apiKey !== void 0 && options.apiKey !== null ? options.apiKey : GEMINI_API_KEY;
   if (!this.vertexai && !this.apiKey) {
@@ -249,7 +249,11 @@ function AITransformFactory(options = {}) {
   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.vertexai) {
+      logger_default.warn(`Billing labels are only supported with Vertex AI. Labels will be ignored.`);
+    } else {
+      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.");
@@ -259,7 +263,7 @@ function AITransformFactory(options = {}) {
     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}`);
     if (this.vertexai) {
-      logger_default.debug(`Using Vertex AI - Project: ${this.project}, Location: ${this.location}`);
+      logger_default.debug(`Using Vertex AI - Project: ${this.project}, Location: ${this.location || "global (default)"}`);
       if (this.googleAuthOptions?.keyFilename) {
         logger_default.debug(`Auth: Service account key file: ${this.googleAuthOptions.keyFilename}`);
       } else if (this.googleAuthOptions?.credentials) {
@@ -275,7 +279,7 @@ function AITransformFactory(options = {}) {
   const clientOptions = this.vertexai ? {
     vertexai: true,
     project: this.project,
-    location: this.location,
+    ...this.location && { location: this.location },
     ...this.googleAuthOptions && { googleAuthOptions: this.googleAuthOptions }
   } : { apiKey: this.apiKey };
   const ai = new import_genai.GoogleGenAI(clientOptions);
@@ -290,7 +294,7 @@ async function initChat(force = false) {
     // @ts-ignore
     config: {
       ...this.chatConfig,
-      ...Object.keys(this.labels).length > 0 && { labels: this.labels }
+      ...this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels }
     },
     history: []
   };
@@ -373,7 +377,7 @@ ${contextText}
     // @ts-ignore
     config: {
       ...this.chatConfig,
-      ...Object.keys(this.labels).length > 0 && { labels: this.labels }
+      ...this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels }
     },
     history: [...currentHistory, ...historyToAdd]
   });
@@ -388,7 +392,7 @@ async function rawMessage(sourcePayload, messageOptions = {}) {
   }
   const actualPayload = typeof sourcePayload === "string" ? sourcePayload : JSON.stringify(sourcePayload, null, 2);
   const mergedLabels = { ...this.labels, ...messageOptions.labels || {} };
-  const hasLabels = Object.keys(mergedLabels).length > 0;
+  const hasLabels = this.vertexai && Object.keys(mergedLabels).length > 0;
   try {
     const sendParams = { message: actualPayload };
     if (hasLabels) {
@@ -603,7 +607,7 @@ async function resetChat() {
       // @ts-ignore
       config: {
         ...this.chatConfig,
-        ...Object.keys(this.labels).length > 0 && { labels: this.labels }
+        ...this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels }
       },
       history: []
     };
@@ -647,7 +651,7 @@ async function clearConversation() {
     // @ts-ignore
     config: {
       ...this.chatConfig,
-      ...Object.keys(this.labels).length > 0 && { labels: this.labels }
+      ...this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels }
     },
     history: exampleHistory
   });
@@ -701,7 +705,7 @@ async function statelessMessage(sourcePayload, options = {}, validatorFn = null)
     contents,
     config: {
       ...this.chatConfig,
-      ...Object.keys(mergedLabels).length > 0 && { labels: mergedLabels }
+      ...this.vertexai && Object.keys(mergedLabels).length > 0 && { labels: mergedLabels }
     }
   });
   this.lastResponseMetadata = {

package/index.js

@@ -194,7 +194,7 @@ function AITransformFactory(options = {}) {
 	// 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.location = options.location || process.env.GOOGLE_CLOUD_LOCATION || undefined;
 	this.googleAuthOptions = options.googleAuthOptions || null;
 
 	// API Key (for Gemini API, not Vertex AI)
@@ -295,10 +295,14 @@ function AITransformFactory(options = {}) {
 	this.enableGrounding = options.enableGrounding || false;
 	this.groundingConfig = options.groundingConfig || {};
 
-	// Billing labels for cost segmentation
+	// Billing labels for cost segmentation (Vertex AI only)
 	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.vertexai) {
+			log.warn(`Billing labels are only supported with Vertex AI. Labels will be ignored.`);
+		} else {
+			log.debug(`Billing labels configured: ${JSON.stringify(this.labels)}`);
+		}
 	}
 
 	if (this.promptKey === this.answerKey) {
@@ -311,7 +315,7 @@ function AITransformFactory(options = {}) {
 		log.debug(`Max output tokens set to: ${this.chatConfig.maxOutputTokens}`);
 		// Log authentication method
 		if (this.vertexai) {
-			log.debug(`Using Vertex AI - Project: ${this.project}, Location: ${this.location}`);
+			log.debug(`Using Vertex AI - Project: ${this.project}, Location: ${this.location || 'global (default)'}`);
 			if (this.googleAuthOptions?.keyFilename) {
 				log.debug(`Auth: Service account key file: ${this.googleAuthOptions.keyFilename}`);
 			} else if (this.googleAuthOptions?.credentials) {
@@ -330,7 +334,7 @@ function AITransformFactory(options = {}) {
 		? {
 			vertexai: true,
 			project: this.project,
-			location: this.location,
+			...(this.location && { location: this.location }),
 			...(this.googleAuthOptions && { googleAuthOptions: this.googleAuthOptions })
 		}
 		: { apiKey: this.apiKey };
@@ -357,7 +361,7 @@ async function initChat(force = false) {
 		// @ts-ignore
 		config: {
 			...this.chatConfig,
-			...(Object.keys(this.labels).length > 0 && { labels: this.labels })
+			...(this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels })
 		},
 		history: [],
 	};
@@ -472,7 +476,7 @@ async function seedWithExamples(examples) {
 		// @ts-ignore
 		config: {
 			...this.chatConfig,
-			...(Object.keys(this.labels).length > 0 && { labels: this.labels })
+			...(this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels })
 		},
 		history: [...currentHistory, ...historyToAdd],
 	});
@@ -509,13 +513,14 @@ async function rawMessage(sourcePayload, messageOptions = {}) {
 		: JSON.stringify(sourcePayload, null, 2);
 
 	// Merge instance labels with per-message labels (per-message takes precedence)
+	// Labels only supported with Vertex AI
 	const mergedLabels = { ...this.labels, ...(messageOptions.labels || {}) };
-	const hasLabels = Object.keys(mergedLabels).length > 0;
+	const hasLabels = this.vertexai && Object.keys(mergedLabels).length > 0;
 
 	try {
 		const sendParams = { message: actualPayload };
 
-		// Add config with labels if we have any
+		// Add config with labels if we have any (Vertex AI only)
 		if (hasLabels) {
 			sendParams.config = { labels: mergedLabels };
 		}
@@ -862,7 +867,7 @@ async function resetChat() {
 			// @ts-ignore
 			config: {
 				...this.chatConfig,
-				...(Object.keys(this.labels).length > 0 && { labels: this.labels })
+				...(this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels })
 			},
 			history: [],
 		};
@@ -933,7 +938,7 @@ async function clearConversation() {
 		// @ts-ignore
 		config: {
 			...this.chatConfig,
-			...(Object.keys(this.labels).length > 0 && { labels: this.labels })
+			...(this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels })
 		},
 		history: exampleHistory,
 	});
@@ -1019,7 +1024,7 @@ async function statelessMessage(sourcePayload, options = {}, validatorFn = null)
 	// Add the user message
 	contents.push({ role: 'user', parts: [{ text: payloadStr }] });
 
-	// Merge labels
+	// Merge labels (Vertex AI only)
 	const mergedLabels = { ...this.labels, ...(options.labels || {}) };
 
 	// Use generateContent instead of chat.sendMessage
@@ -1028,7 +1033,7 @@ async function statelessMessage(sourcePayload, options = {}, validatorFn = null)
 		contents: contents,
 		config: {
 			...this.chatConfig,
-			...(Object.keys(mergedLabels).length > 0 && { labels: mergedLabels })
+			...(this.vertexai && Object.keys(mergedLabels).length > 0 && { labels: mergedLabels })
 		}
 	});
 

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

package/types.d.ts

@@ -87,6 +87,7 @@ export interface AITransformerContext {
   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)
+  vertexai?: boolean; // Whether using Vertex AI (true) or Gemini API (false)
   estimate?: (nextPayload: Record<string, unknown> | string) => Promise<{ inputTokens: number }>;
   getLastUsage?: () => UsageData | null;
   lastResponseMetadata?: ResponseMetadata | null; // Metadata from the last API response
@@ -153,7 +154,7 @@ export interface AITransformerOptions {
   // 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
+  location?: string; // Google Cloud location/region (e.g., 'us-central1'). If not set, defaults to 'global' endpoint
   googleAuthOptions?: GoogleAuthOptions; // Authentication options for Vertex AI (keyFilename, credentials, etc.)
 }
 
@@ -185,6 +186,7 @@ export declare class AITransformer {
   enableGrounding: boolean;
   groundingConfig: Record<string, any>;
   labels: Record<string, string>;
+  vertexai: boolean;
   /** 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) */