ak-gemini 1.0.53 → 1.0.55

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 }
@@ -161,6 +165,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,14 +183,22 @@ 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;
     }
   }
-  if (examples?.slice().pop()[this.systemInstructionsKey]) {
+  const instructionExample = examples.find((ex) => ex[this.systemInstructionsKey]);
+  if (instructionExample) {
     logger_default.debug(`Found system instructions in examples; reinitializing chat with new instructions.`);
-    this.systemInstructions = examples.slice().pop()[this.systemInstructionsKey];
+    this.systemInstructions = instructionExample[this.systemInstructionsKey];
     this.chatConfig.systemInstruction = this.systemInstructions;
     await this.init(true);
   }
@@ -213,14 +231,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) {
@@ -254,6 +274,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
@@ -196,6 +201,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,15 +232,27 @@ 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;
 		}
 	}
 
-	if (examples?.slice().pop()[this.systemInstructionsKey]) {
+	const instructionExample = examples.find(ex => ex[this.systemInstructionsKey]);
+	if (instructionExample) {
 		log.debug(`Found system instructions in examples; reinitializing chat with new instructions.`);
-		this.systemInstructions = examples.slice().pop()[this.systemInstructionsKey];
+		this.systemInstructions = instructionExample[this.systemInstructionsKey];
 		this.chatConfig.systemInstruction = this.systemInstructions;
 		await this.init(true); // Reinitialize chat with new system instructions
 	}
@@ -266,8 +292,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
@@ -275,9 +302,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;
 }
 
 /**
@@ -345,7 +373,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.53",
+	"version": "1.0.55",
 	"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": {
@@ -36,6 +42,11 @@
 		"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:all": "./scripts/test-all.sh",
+		"test:integration": "./scripts/test-local.sh",
+		"test:unit": "npm test -- tests/module.test.js",
+		"test:function": "npm test -- tests/function.integration.test.js",
+		"test:http": "npm test -- tests/function.http.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",

package/types.d.ts

@@ -0,0 +1,83 @@
+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
+  targetKey?: string; // Key name for target data 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
+  promptKey?: string; // Key for the prompt in examples
+  answerKey?: string; // Key for the answer in examples
+  contextKey?: string; // Key for the context in examples
+  explanationKey?: string; // Key for the explanation in examples
+}
+
+// Async validator function type
+export type AsyncValidatorFunction = (payload: Record<string, unknown>) => Promise<unknown>;
+
+
+export declare class AITransformer implements AITransformerContext {}