ak-gemini 1.0.55 → 1.0.58

package/index.cjs

@@ -119,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;
@@ -138,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";

package/index.js

@@ -120,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);
 	}
 }
 
@@ -155,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

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.55",
+	"version": "1.0.58",
 	"main": "index.js",
 	"files": [
 		"index.js",
@@ -32,22 +32,16 @@
 	},
 	"homepage": "https://github.com/ak--47/ak-gemini#readme",
 	"scripts": {
-		"prepublishOnly": "npm run build:cjs",
+		"prepublishOnly": "npm run typecheck && 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: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"
+		"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",
+		"coverage": "node --no-warnings --experimental-vm-modules node_modules/jest/bin/jest.js --coverage",
+		"typecheck": "tsc --noEmit"
 	},
 	"type": "module",
 	"keywords": [
@@ -57,7 +51,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",
@@ -68,6 +61,7 @@
 		"@types/jest": "^29.5.14",
 		"esbuild": "^0.25.5",
 		"jest": "^29.7.0",
-		"nodemon": "^3.1.10"
+		"nodemon": "^3.1.10",
+		"typescript": "^5.9.2"
 	}
-}
+}

package/types.d.ts

@@ -31,20 +31,25 @@ export interface AITransformerContext {
   systemInstructionsKey?: string;
   maxRetries?: number;
   retryDelay?: number;
-  init?: () => Promise<void>; // Initialization function
+  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
-  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
+  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
   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
+  CONTEXT?: Record<string, unknown> | string; // 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
+  INPUT?: Record<string, unknown>; // alias for PROMPT
+  OUTPUT?: Record<string, unknown>; // alias for ANSWER
+  SYSTEM?: string; // system instructions for this example
+  EXPLANATION?: string; // explanation for this example
+  [key: string]: any; // allow additional properties for flexible key mapping
 }
 
 export interface ExampleFileContent {
@@ -58,8 +63,10 @@ export interface AITransformerOptions {
   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
+  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
@@ -70,14 +77,47 @@ export interface AITransformerOptions {
   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 {}
+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> | 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>>;
+  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;