ak-gemini 1.0.54 → 1.0.57

package/index.cjs

@@ -56,10 +56,14 @@ var logger_default = logger;
 // index.js
 var import_meta = {};
 import_dotenv.default.config();
-var { NODE_ENV = "unknown", GEMINI_API_KEY } = process.env;
+var { NODE_ENV = "unknown", GEMINI_API_KEY, LOG_LEVEL = "" } = process.env;
 if (NODE_ENV === "dev") logger_default.level = "debug";
 if (NODE_ENV === "test") logger_default.level = "warn";
 if (NODE_ENV.startsWith("prod")) logger_default.level = "error";
+if (LOG_LEVEL) {
+  logger_default.level = LOG_LEVEL;
+  logger_default.debug(`Setting log level to ${LOG_LEVEL}`);
+}
 var DEFAULT_SAFETY_SETTINGS = [
   { category: import_genai.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold: import_genai.HarmBlockThreshold.BLOCK_NONE },
   { category: import_genai.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold: import_genai.HarmBlockThreshold.BLOCK_NONE }
@@ -115,7 +119,9 @@ var AITransformer = class {
     this.reset = resetChat.bind(this);
     this.getHistory = getChatHistory.bind(this);
     this.messageAndValidate = prepareAndValidateMessage.bind(this);
+    this.transformWithValidation = prepareAndValidateMessage.bind(this);
     this.estimate = estimateTokenUsage.bind(this);
+    this.estimateTokenUsage = estimateTokenUsage.bind(this);
   }
 };
 var index_default = AITransformer;
@@ -134,8 +140,8 @@ function AITransformFactory(options = {}) {
   }
   this.examplesFile = options.examplesFile || null;
   this.exampleData = options.exampleData || null;
-  this.promptKey = options.promptKey || "PROMPT";
-  this.answerKey = options.answerKey || "ANSWER";
+  this.promptKey = options.promptKey || options.sourceKey || "PROMPT";
+  this.answerKey = options.answerKey || options.targetKey || "ANSWER";
   this.contextKey = options.contextKey || "CONTEXT";
   this.explanationKey = options.explanationKey || "EXPLANATION";
   this.systemInstructionsKey = options.systemInstructionsKey || "SYSTEM";
@@ -161,6 +167,12 @@ async function initChat(force = false) {
     config: this.chatConfig,
     history: []
   });
+  try {
+    await this.genAIClient.models.list();
+    logger_default.debug("Gemini API connection successful.");
+  } catch (e) {
+    throw new Error(`Gemini chat initialization failed: ${e.message}`);
+  }
   logger_default.debug("Gemini chat session initialized.");
 }
 async function seedWithExamples(examples) {
@@ -173,6 +185,13 @@ async function seedWithExamples(examples) {
       } catch (err) {
         throw new Error(`Could not load examples from file: ${this.examplesFile}. Please check the file path and format.`);
       }
+    } else if (this.exampleData) {
+      logger_default.debug(`Using example data provided in options.`);
+      if (Array.isArray(this.exampleData)) {
+        examples = this.exampleData;
+      } else {
+        throw new Error(`Invalid example data provided. Expected an array of examples.`);
+      }
     } else {
       logger_default.debug("No examples provided and no examples file specified. Skipping seeding.");
       return;
@@ -214,14 +233,16 @@ ${contextText}
     }
   }
   const currentHistory = this?.chat?.getHistory() || [];
+  logger_default.debug(`Adding ${historyToAdd.length} examples to chat history (${currentHistory.length} current examples)...`);
   this.chat = await this.genAIClient.chats.create({
     model: this.modelName,
     // @ts-ignore
     config: this.chatConfig,
     history: [...currentHistory, ...historyToAdd]
   });
-  logger_default.debug("Transformation examples seeded successfully.");
-  return this.chat.getHistory();
+  const newHistory = this.chat.getHistory();
+  logger_default.debug(`Created new chat session with ${newHistory.length} examples.`);
+  return newHistory;
 }
 async function rawMessage(sourcePayload) {
   if (!this.chat) {
@@ -255,6 +276,10 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
     lastPayload = JSON.stringify(sourcePayload, null, 2);
   } else if (typeof sourcePayload === "string") {
     lastPayload = sourcePayload;
+  } else if (typeof sourcePayload === "boolean" || typeof sourcePayload === "number") {
+    lastPayload = sourcePayload.toString();
+  } else if (sourcePayload === null || sourcePayload === void 0) {
+    lastPayload = JSON.stringify({});
   } else {
     throw new Error("Invalid source payload. Must be a JSON object or string.");
   }

package/index.js

@@ -18,7 +18,7 @@
 //env
 import dotenv from 'dotenv';
 dotenv.config();
-const { NODE_ENV = "unknown", GEMINI_API_KEY } = process.env;
+const { NODE_ENV = "unknown", GEMINI_API_KEY, LOG_LEVEL = "" } = process.env;
 
 
 
@@ -33,6 +33,11 @@ if (NODE_ENV === 'dev') log.level = 'debug';
 if (NODE_ENV === 'test') log.level = 'warn';
 if (NODE_ENV.startsWith('prod')) log.level = 'error';
 
+if (LOG_LEVEL) {
+	log.level = LOG_LEVEL;
+	log.debug(`Setting log level to ${LOG_LEVEL}`);
+}
+
 
 
 // defaults
@@ -115,7 +120,9 @@ class AITransformer {
 		this.reset = resetChat.bind(this);
 		this.getHistory = getChatHistory.bind(this);
 		this.messageAndValidate = prepareAndValidateMessage.bind(this);
+		this.transformWithValidation = prepareAndValidateMessage.bind(this);
 		this.estimate = estimateTokenUsage.bind(this);
+		this.estimateTokenUsage = estimateTokenUsage.bind(this);
 	}
 }
 
@@ -150,8 +157,8 @@ function AITransformFactory(options = {}) {
 	this.exampleData = options.exampleData || null; // can be used instead of examplesFile
 
 	// Use configurable keys with fallbacks
-	this.promptKey = options.promptKey || 'PROMPT';
-	this.answerKey = options.answerKey || 'ANSWER';
+	this.promptKey = options.promptKey || options.sourceKey || 'PROMPT';
+	this.answerKey = options.answerKey || options.targetKey || 'ANSWER';
 	this.contextKey = options.contextKey || 'CONTEXT'; // Optional key for context
 	this.explanationKey = options.explanationKey || 'EXPLANATION'; // Optional key for explanations
 	this.systemInstructionsKey = options.systemInstructionsKey || 'SYSTEM'; // Optional key for system instructions
@@ -196,6 +203,15 @@ async function initChat(force = false) {
 		history: [],
 	});
 
+	try {
+		await this.genAIClient.models.list();
+		log.debug("Gemini API connection successful.");
+	} catch (e) {
+		throw new Error(`Gemini chat initialization failed: ${e.message}`);
+	}
+
+
+
 	log.debug("Gemini chat session initialized.");
 }
 
@@ -218,7 +234,18 @@ async function seedWithExamples(examples) {
 			catch (err) {
 				throw new Error(`Could not load examples from file: ${this.examplesFile}. Please check the file path and format.`);
 			}
-		} else {
+		}
+
+		else if (this.exampleData) {
+			log.debug(`Using example data provided in options.`);
+			if (Array.isArray(this.exampleData)) {
+				examples = this.exampleData;
+			} else {
+				throw new Error(`Invalid example data provided. Expected an array of examples.`);
+			}
+		}
+
+		else {
 			log.debug("No examples provided and no examples file specified. Skipping seeding.");
 			return;
 		}
@@ -267,8 +294,9 @@ async function seedWithExamples(examples) {
 
 	}
 
-	const currentHistory = this?.chat?.getHistory() || [];
 
+	const currentHistory = this?.chat?.getHistory() || [];
+	log.debug(`Adding ${historyToAdd.length} examples to chat history (${currentHistory.length} current examples)...`);
 	this.chat = await this.genAIClient.chats.create({
 		model: this.modelName,
 		// @ts-ignore
@@ -276,9 +304,10 @@ async function seedWithExamples(examples) {
 		history: [...currentHistory, ...historyToAdd],
 	});
 
-	log.debug("Transformation examples seeded successfully.");
 
-	return this.chat.getHistory(); // Return the updated chat history for reference
+	const newHistory = this.chat.getHistory();
+	log.debug(`Created new chat session with ${newHistory.length} examples.`);
+	return newHistory;
 }
 
 /**
@@ -346,7 +375,14 @@ async function prepareAndValidateMessage(sourcePayload, options = {}, validatorF
 		lastPayload = JSON.stringify(sourcePayload, null, 2);
 	} else if (typeof sourcePayload === 'string') {
 		lastPayload = sourcePayload;
-	} else {
+	}
+	else if (typeof sourcePayload === 'boolean' || typeof sourcePayload === 'number') {
+		lastPayload = sourcePayload.toString();
+	}
+	else if (sourcePayload === null || sourcePayload === undefined) {
+		lastPayload = JSON.stringify({}); // Convert null/undefined to empty object
+	}
+	else {
 		throw new Error("Invalid source payload. Must be a JSON object or string.");
 	}
 

package/package.json

@@ -2,19 +2,25 @@
 	"name": "ak-gemini",
 	"author": "ak@mixpanel.com",
 	"description": "AK's Generative AI Helper for doing... transforms",
-	"version": "1.0.54",
+	"version": "1.0.57",
 	"main": "index.js",
 	"files": [
 		"index.js",
 		"index.cjs",
-		"types.ts",
+		"types.d.ts",
 		"logger.js"
 	],
 	"types": "types.d.ts",
 	"exports": {
 		".": {
-			"import": "./index.js",
-			"require": "./index.cjs"
+			"import": {
+				"types": "./types.d.ts",
+				"default": "./index.js"
+			},
+			"require": {
+				"types": "./types.d.ts",
+				"default": "./index.cjs"
+			}
 		}
 	},
 	"repository": {
@@ -29,13 +35,10 @@
 		"prepublishOnly": "npm run build:cjs",
 		"post": "npm publish --access public",
 		"release": "npm version patch && npm publish --access public",
-		"local": "./scripts/local.sh",
-		"fire": "./scripts/fire.sh",
-		"deploy": "./scripts/deploy.sh",
-		"perms": "chmod +x ./scripts/*.sh",
 		"update-deps": "npx npm-check-updates -u && npm install",
 		"prune": "rm -rf tmp/*",
 		"test": "node --no-warnings --experimental-vm-modules node_modules/jest/bin/jest.js",
+		"test:unit": "npm test -- tests/module.test.js",
 		"build:cjs": "esbuild index.js --bundle --platform=node --format=cjs --outfile=index.cjs --external:@google/genai --external:ak-tools --external:dotenv --external:pino-pretty --external:pino"
 	},
 	"type": "module",
@@ -46,7 +49,6 @@
 	],
 	"license": "ISC",
 	"dependencies": {
-		"@google-cloud/functions-framework": "^4.0.0",
 		"@google/genai": "^1.4.0",
 		"ak-tools": "^1.0.64",
 		"dotenv": "^16.5.0",

package/types.d.ts

@@ -0,0 +1,118 @@
+import type { GoogleGenAI } from '@google/genai';
+
+export interface SafetySetting {
+  category: string; // The harm category
+  threshold: string; // The blocking threshold
+}
+
+export interface ChatConfig {
+  responseMimeType?: string; // MIME type for responses
+  temperature?: number; // Controls randomness (0.0 to 1.0)
+  topP?: number; // Controls diversity via nucleus sampling
+  topK?: number; // Controls diversity by limiting top-k tokens
+  systemInstruction?: string; // System instruction for the model
+  safetySettings?: SafetySetting[]; // Safety settings array
+  responseSchema?: Object; // Schema for validating model responses
+  [key: string]: any; // Additional properties for flexibility
+}
+
+export interface AITransformerContext {
+  modelName?: string;
+  systemInstructions?: string;
+  chatConfig?: ChatConfig;
+  genAI?: any;
+  chat?: any;
+  examplesFile?: string | null;
+  exampleData?: TransformationExample[] | null;
+  promptKey?: string;
+  answerKey?: string;
+  contextKey?: string;
+  explanationKey?: string;
+  systemInstructionsKey?: string;
+  maxRetries?: number;
+  retryDelay?: number;
+  init?: () => 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
+  rebuild?: () => Promise<Record<string, unknown>; // Function to rebuild the transformer
+  rawMessage?: (payload: Record<string, unknown>) => 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
+  
+}
+
+export interface TransformationExample {
+  CONTEXT?: Record<string, unknown>; // optional context for the transformation
+  PROMPT?: Record<string, unknown>; // what the user provides as input
+  ANSWER?: Record<string, unknown>; // what the model should return as output
+}
+
+export interface ExampleFileContent {
+  examples: TransformationExample[];
+}
+
+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
+  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)
+  targetKey?: string; // Key name for target data in examples (alias for answerKey)
+  promptKey?: string; // Key for the prompt in examples
+  answerKey?: string; // Key for the answer in examples
+  contextKey?: string; // Key name for context data in examples
+  explanationKey?: string; // Key name for explanation data in examples
+  systemInstructionsKey?: string; // Key for system instructions in examples
+  maxRetries?: number; // Maximum retry attempts for auto-retry functionality
+  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
+  onlyJSON?: boolean; // If true, only JSON responses are allowed
+  asyncValidator?: AsyncValidatorFunction; // Optional async validator function for response validation
+}
+
+// Async validator function type
+export type AsyncValidatorFunction = (payload: Record<string, unknown>) => Promise<unknown>;
+
+
+export declare class AITransformer {
+  // Constructor
+  constructor(options?: AITransformerOptions);
+  
+  // Properties
+  modelName: string;
+  promptKey: string;
+  answerKey: string;
+  contextKey: string;
+  explanationKey: string;
+  systemInstructionKey: string;
+  maxRetries: number;
+  retryDelay: number;
+  systemInstructions: string;
+  chatConfig: ChatConfig;
+  apiKey: string;
+  onlyJSON: boolean;
+  asyncValidator: AsyncValidatorFunction | null;
+  genAIClient: any;
+  chat: any;
+  
+  // 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>>;
+  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>>;
+  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 }>;
+}
+
+// Default export
+declare const _default: typeof AITransformer;
+export default _default;