ak-gemini 1.0.5 → 1.0.7

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.5",
+	"version": "1.0.7",
 	"main": "index.js",
 	"files": [
 		"index.js",
 		"index.cjs",
-		"types.ts",
+		"types.d.ts",
 		"logger.js"
 	],
-	"types": "types.ts",
+	"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": {
@@ -26,17 +32,17 @@
 	},
 	"homepage": "https://github.com/ak--47/ak-gemini#readme",
 	"scripts": {
-		"prepublishOnly": "npm run build:cjs",
-		"post": "npm publish --access public",
+		"prepublishOnly": "npm run typecheck && npm run build:cjs",
+		"post": "npm publish",
 		"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",
-		"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"
+		"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",
+		"coverage": "node --no-warnings --experimental-vm-modules node_modules/jest/bin/jest.js --coverage",
+		"typecheck": "tsc --noEmit",
+		"update:gemini": "npm install @google/genai@latest"
 	},
 	"type": "module",
 	"keywords": [
@@ -46,17 +52,17 @@
 	],
 	"license": "ISC",
 	"dependencies": {
-		"@google-cloud/functions-framework": "^4.0.0",
-		"@google/genai": "^1.4.0",
-		"ak-tools": "^1.0.64",
+		"@google/genai": "^1.34.0",
+		"ak-tools": "^1.1.12",
 		"dotenv": "^16.5.0",
 		"pino": "^9.7.0",
 		"pino-pretty": "^13.0.0"
 	},
 	"devDependencies": {
 		"@types/jest": "^29.5.14",
-		"esbuild": "^0.25.5",
+		"esbuild": "^0.25.11",
 		"jest": "^29.7.0",
-		"nodemon": "^3.1.10"
+		"nodemon": "^3.1.10",
+		"typescript": "^5.9.2"
 	}
 }

package/types.d.ts

@@ -0,0 +1,138 @@
+import type { GoogleGenAI, ThinkingLevel } from '@google/genai';
+
+export { ThinkingLevel };
+
+export interface ThinkingConfig {
+  /** Indicates whether to include thoughts in the response. If true, thoughts are returned only if the model supports thought and thoughts are available. */
+  includeThoughts?: boolean;
+  /** Indicates the thinking budget in tokens. 0 is DISABLED. -1 is AUTOMATIC. The default values and allowed ranges are model dependent. */
+  thinkingBudget?: number;
+  /** Optional. The number of thoughts tokens that the model should generate. */
+  thinkingLevel?: ThinkingLevel;
+}
+
+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
+  thinkingConfig?: ThinkingConfig; // Thinking features configuration
+  [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?: (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?: (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> | 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 {
+  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
+  thinkingConfig?: ThinkingConfig; // Thinking features configuration (defaults to thinkingBudget: 0, thinkingLevel: "MINIMAL")
+  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
+  logLevel?: 'trace' | 'debug' | 'info' | 'warn' | 'error' | 'fatal' | 'none'; // Log level for the logger (defaults to 'info', 'none' disables logging)
+}
+
+// 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;
+  logLevel: 'trace' | 'debug' | 'info' | 'warn' | 'error' | 'fatal' | 'none';
+  
+  // 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;

package/types.ts

@@ -1,65 +0,0 @@
-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
-}
-
-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;
-  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
-  genAIClient?: GoogleGenAI; // Google GenAI client instance
-  
-}
-
-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 {
-  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
-  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
-}
-
-// Async validator function type
-export type AsyncValidatorFunction = (payload: Record<string, unknown>) => Promise<Record<string, unknown>>;