ak-gemini 2.0.2 → 2.0.4

package/index.cjs

@@ -32,6 +32,7 @@ __export(index_exports, {
   BaseGemini: () => base_default,
   Chat: () => chat_default,
   CodeAgent: () => code_agent_default,
+  Embedding: () => Embedding,
   HarmBlockThreshold: () => import_genai2.HarmBlockThreshold,
   HarmCategory: () => import_genai2.HarmCategory,
   Message: () => message_default,
@@ -310,7 +311,7 @@ function extractJSON(text) {
 }
 
 // base.js
-import_dotenv.default.config();
+import_dotenv.default.config({ quiet: true });
 var { NODE_ENV = "unknown", LOG_LEVEL = "" } = process.env;
 var DEFAULT_SAFETY_SETTINGS = [
   { category: import_genai.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold: import_genai.HarmBlockThreshold.BLOCK_NONE },
@@ -335,7 +336,8 @@ var MODEL_PRICING = {
   "gemini-3-pro": { input: 2, output: 12 },
   "gemini-3-pro-preview": { input: 2, output: 12 },
   "gemini-2.0-flash": { input: 0.1, output: 0.4 },
-  "gemini-2.0-flash-lite": { input: 0.02, output: 0.1 }
+  "gemini-2.0-flash-lite": { input: 0.02, output: 0.1 },
+  "gemini-embedding-001": { input: 6e-3, output: 0 }
 };
 var BaseGemini = class {
   /**
@@ -361,6 +363,9 @@ var BaseGemini = class {
     }
     this._configureLogLevel(options.logLevel);
     this.labels = options.labels || {};
+    this.enableGrounding = options.enableGrounding || false;
+    this.groundingConfig = options.groundingConfig || {};
+    this.cachedContent = options.cachedContent || null;
     this.chatConfig = {
       temperature: 0.7,
       topP: 0.95,
@@ -433,14 +438,21 @@ var BaseGemini = class {
    * @protected
    */
   _getChatCreateOptions() {
-    return {
+    const opts = {
       model: this.modelName,
       config: {
         ...this.chatConfig,
-        ...this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels }
+        ...this.vertexai && Object.keys(this.labels).length > 0 && { labels: this.labels },
+        ...this.cachedContent && { cachedContent: this.cachedContent }
       },
       history: []
     };
+    if (this.enableGrounding) {
+      const existingTools = opts.config.tools || [];
+      opts.config.tools = [...existingTools, { googleSearch: this.groundingConfig }];
+      logger_default.debug("Search grounding ENABLED (WARNING: costs $35/1k queries)");
+    }
+    return opts;
   }
   // ── Chat Session Management ──────────────────────────────────────────────
   /**
@@ -562,7 +574,8 @@ ${contextText}
       promptTokens: response.usageMetadata?.promptTokenCount || 0,
       responseTokens: response.usageMetadata?.candidatesTokenCount || 0,
       totalTokens: response.usageMetadata?.totalTokenCount || 0,
-      timestamp: Date.now()
+      timestamp: Date.now(),
+      groundingMetadata: response.candidates?.[0]?.groundingMetadata || null
     };
   }
   /**
@@ -582,7 +595,8 @@ ${contextText}
       attempts: useCumulative ? cumulative.attempts : 1,
       modelVersion: meta.modelVersion,
       requestedModel: meta.requestedModel,
-      timestamp: meta.timestamp
+      timestamp: meta.timestamp,
+      groundingMetadata: meta.groundingMetadata || null
     };
   }
   // ── Token Estimation ─────────────────────────────────────────────────────
@@ -627,6 +641,99 @@ ${contextText}
       note: "Cost is for input tokens only; output cost depends on response length"
     };
   }
+  // ── Context Caching ─────────────────────────────────────────────────────
+  /**
+   * Creates a cached content resource for cost reduction on repeated prompts.
+   * Auto-populates model and systemInstruction from this instance if not provided.
+   * @param {Object} [config={}] - Cache configuration
+   * @param {string} [config.model] - Model (defaults to this.modelName)
+   * @param {string} [config.ttl] - Time-to-live (e.g., '3600s')
+   * @param {string} [config.displayName] - Human-readable name
+   * @param {Array} [config.contents] - Content to cache
+   * @param {string} [config.systemInstruction] - System prompt to cache (defaults to this.systemPrompt)
+   * @param {Array} [config.tools] - Tools to cache
+   * @param {Object} [config.toolConfig] - Tool configuration to cache
+   * @returns {Promise<Object>} The created cache resource
+   */
+  async createCache(config = {}) {
+    const cacheConfig = {};
+    if (config.ttl) cacheConfig.ttl = config.ttl;
+    if (config.displayName) cacheConfig.displayName = config.displayName;
+    if (config.contents) cacheConfig.contents = config.contents;
+    if (config.tools) cacheConfig.tools = config.tools;
+    if (config.toolConfig) cacheConfig.toolConfig = config.toolConfig;
+    const sysInstruction = config.systemInstruction !== void 0 ? config.systemInstruction : this.systemPrompt;
+    if (sysInstruction) cacheConfig.systemInstruction = sysInstruction;
+    const cached = await this.genAIClient.caches.create({
+      model: config.model || this.modelName,
+      config: cacheConfig
+    });
+    logger_default.debug(`Cache created: ${cached.name}`);
+    return cached;
+  }
+  /**
+   * Retrieves a cached content resource by name.
+   * @param {string} cacheName - Server-generated resource name
+   * @returns {Promise<Object>} The cached content resource
+   */
+  async getCache(cacheName) {
+    return await this.genAIClient.caches.get({ name: cacheName });
+  }
+  /**
+   * Lists all cached content resources.
+   * @returns {Promise<Object>} Pager of cached content resources
+   */
+  async listCaches() {
+    const pager = await this.genAIClient.caches.list();
+    const results = [];
+    for await (const cache of pager) {
+      results.push(cache);
+    }
+    return results;
+  }
+  /**
+   * Updates a cached content resource (TTL or expiration).
+   * @param {string} cacheName - Server-generated resource name
+   * @param {Object} [config={}] - Update config
+   * @param {string} [config.ttl] - New TTL (e.g., '7200s')
+   * @param {string} [config.expireTime] - New expiration (RFC 3339)
+   * @returns {Promise<Object>} The updated cache resource
+   */
+  async updateCache(cacheName, config = {}) {
+    return await this.genAIClient.caches.update({
+      name: cacheName,
+      config: {
+        ...config.ttl && { ttl: config.ttl },
+        ...config.expireTime && { expireTime: config.expireTime }
+      }
+    });
+  }
+  /**
+   * Deletes a cached content resource.
+   * Clears this.cachedContent if it matches the deleted cache.
+   * @param {string} cacheName - Server-generated resource name
+   * @returns {Promise<void>}
+   */
+  async deleteCache(cacheName) {
+    await this.genAIClient.caches.delete({ name: cacheName });
+    logger_default.debug(`Cache deleted: ${cacheName}`);
+    if (this.cachedContent === cacheName) {
+      this.cachedContent = null;
+    }
+  }
+  /**
+   * Sets the cached content for this instance and reinitializes the session.
+   * @param {string} cacheName - Server-generated cache resource name
+   * @returns {Promise<void>}
+   */
+  async useCache(cacheName) {
+    this.cachedContent = cacheName;
+    delete this.chatConfig.systemInstruction;
+    if (this.chatSession) {
+      await this.init(true);
+    }
+    logger_default.debug(`Using cache: ${cacheName}`);
+  }
   // ── Private Helpers ──────────────────────────────────────────────────────
   /**
    * Configures the log level based on options, env vars, or NODE_ENV.
@@ -722,20 +829,8 @@ var Transformer = class extends base_default {
     this.asyncValidator = options.asyncValidator || null;
     this.maxRetries = options.maxRetries || 3;
     this.retryDelay = options.retryDelay || 1e3;
-    this.enableGrounding = options.enableGrounding || false;
-    this.groundingConfig = options.groundingConfig || {};
     logger_default.debug(`Transformer keys \u2014 Source: "${this.promptKey}", Target: "${this.answerKey}", Context: "${this.contextKey}"`);
   }
-  // ── Chat Create Options Override ──────────────────────────────────────────
-  /** @protected */
-  _getChatCreateOptions() {
-    const opts = super._getChatCreateOptions();
-    if (this.enableGrounding) {
-      opts.config.tools = [{ googleSearch: this.groundingConfig }];
-      logger_default.debug(`Search grounding ENABLED (WARNING: costs $35/1k queries)`);
-    }
-    return opts;
-  }
   // ── Seeding ──────────────────────────────────────────────────────────────
   /**
    * Seeds the chat with transformation examples using the configured key mapping.
@@ -2221,14 +2316,152 @@ ${serialized}` });
 };
 var rag_agent_default = RagAgent;
 
+// embedding.js
+var Embedding = class extends base_default {
+  /**
+   * @param {import('./types.d.ts').EmbeddingOptions} [options={}]
+   */
+  constructor(options = {}) {
+    if (options.modelName === void 0) {
+      options = { ...options, modelName: "gemini-embedding-001" };
+    }
+    if (options.systemPrompt === void 0) {
+      options = { ...options, systemPrompt: null };
+    }
+    super(options);
+    this.taskType = options.taskType || null;
+    this.title = options.title || null;
+    this.outputDimensionality = options.outputDimensionality || null;
+    this.autoTruncate = options.autoTruncate ?? true;
+    logger_default.debug(`Embedding created with model: ${this.modelName}`);
+  }
+  /**
+   * Initialize the Embedding client.
+   * Override: validates API connection only, NO chat session (stateless).
+   * @param {boolean} [force=false]
+   * @returns {Promise<void>}
+   */
+  async init(force = false) {
+    if (this._initialized && !force) return;
+    logger_default.debug(`Initializing ${this.constructor.name} with model: ${this.modelName}...`);
+    try {
+      await this.genAIClient.models.list();
+      logger_default.debug(`${this.constructor.name}: API connection successful.`);
+    } catch (e) {
+      throw new Error(`${this.constructor.name} initialization failed: ${e.message}`);
+    }
+    this._initialized = true;
+    logger_default.debug(`${this.constructor.name}: Initialized (stateless mode).`);
+  }
+  /**
+   * Builds the config object for embedContent calls.
+   * @param {Object} [overrides={}] - Per-call config overrides
+   * @returns {Object} The config object
+   * @private
+   */
+  _buildConfig(overrides = {}) {
+    const config = {};
+    const taskType = overrides.taskType || this.taskType;
+    const title = overrides.title || this.title;
+    const dims = overrides.outputDimensionality || this.outputDimensionality;
+    if (taskType) config.taskType = taskType;
+    if (title) config.title = title;
+    if (dims) config.outputDimensionality = dims;
+    return config;
+  }
+  /**
+  	 * Embed a single text string.
+  	 * @param {string} text - The text to embed
+  	 * @param {Object} [config={}] - Per-call config overrides
+  	 * @param {string} [config.taskType] - Override task type
+  	 * @param {string} [config.title] - Override title
+  	 * @param {number} [config.outputDimensionality] - Override dimensions
+  
+  	 * @returns {Promise<import('./types.d.ts').EmbeddingResult>} The embedding result
+  	 */
+  async embed(text, config = {}) {
+    if (!this._initialized) await this.init();
+    const result = await this.genAIClient.models.embedContent({
+      model: this.modelName,
+      contents: text,
+      config: this._buildConfig(config)
+    });
+    return result.embeddings[0];
+  }
+  /**
+  	 * Embed multiple text strings in a single API call.
+  	 * @param {string[]} texts - Array of texts to embed
+  	 * @param {Object} [config={}] - Per-call config overrides
+  	 * @param {string} [config.taskType] - Override task type
+  	 * @param {string} [config.title] - Override title
+  	 * @param {number} [config.outputDimensionality] - Override dimensions
+  
+  	 * @returns {Promise<import('./types.d.ts').EmbeddingResult[]>} Array of embedding results
+  	 */
+  async embedBatch(texts, config = {}) {
+    if (!this._initialized) await this.init();
+    const result = await this.genAIClient.models.embedContent({
+      model: this.modelName,
+      contents: texts,
+      config: this._buildConfig(config)
+    });
+    return result.embeddings;
+  }
+  /**
+   * Compute cosine similarity between two embedding vectors.
+   * Pure math — no API call.
+   * @param {number[]} a - First embedding vector
+   * @param {number[]} b - Second embedding vector
+   * @returns {number} Cosine similarity between -1 and 1
+   */
+  similarity(a, b) {
+    if (!a || !b || a.length !== b.length) {
+      throw new Error("Vectors must be non-null and have the same length");
+    }
+    let dot = 0;
+    let magA = 0;
+    let magB = 0;
+    for (let i = 0; i < a.length; i++) {
+      dot += a[i] * b[i];
+      magA += a[i] * a[i];
+      magB += b[i] * b[i];
+    }
+    const magnitude = Math.sqrt(magA) * Math.sqrt(magB);
+    if (magnitude === 0) return 0;
+    return dot / magnitude;
+  }
+  // ── No-ops (embeddings don't use chat sessions) ──
+  /** @returns {any[]} Always returns empty array */
+  getHistory() {
+    return [];
+  }
+  /** No-op for Embedding */
+  async clearHistory() {
+  }
+  /** No-op for Embedding */
+  async seed() {
+    logger_default.warn("Embedding.seed() is a no-op \u2014 embeddings do not support few-shot examples.");
+    return [];
+  }
+  /**
+   * @param {any} _nextPayload
+   * @throws {Error} Embedding does not support token estimation
+   * @returns {Promise<{inputTokens: number}>}
+   */
+  async estimate(_nextPayload) {
+    throw new Error("Embedding does not support token estimation. Use embed() directly.");
+  }
+};
+
 // index.js
 var import_genai2 = require("@google/genai");
-var index_default = { Transformer: transformer_default, Chat: chat_default, Message: message_default, ToolAgent: tool_agent_default, CodeAgent: code_agent_default, RagAgent: rag_agent_default };
+var index_default = { Transformer: transformer_default, Chat: chat_default, Message: message_default, ToolAgent: tool_agent_default, CodeAgent: code_agent_default, RagAgent: rag_agent_default, Embedding };
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
   BaseGemini,
   Chat,
   CodeAgent,
+  Embedding,
   HarmBlockThreshold,
   HarmCategory,
   Message,

package/index.js

@@ -26,6 +26,7 @@ export { default as Message } from './message.js';
 export { default as ToolAgent } from './tool-agent.js';
 export { default as CodeAgent } from './code-agent.js';
 export { default as RagAgent } from './rag-agent.js';
+export { default as Embedding } from './embedding.js';
 export { default as BaseGemini } from './base.js';
 export { default as log } from './logger.js';
 export { ThinkingLevel, HarmCategory, HarmBlockThreshold } from '@google/genai';
@@ -39,5 +40,6 @@ import Message from './message.js';
 import ToolAgent from './tool-agent.js';
 import CodeAgent from './code-agent.js';
 import RagAgent from './rag-agent.js';
+import Embedding from './embedding.js';
 
-export default { Transformer, Chat, Message, ToolAgent, CodeAgent, RagAgent };
+export default { Transformer, Chat, Message, ToolAgent, CodeAgent, RagAgent, Embedding };

package/package.json

@@ -2,7 +2,7 @@
 	"name": "ak-gemini",
 	"author": "ak@mixpanel.com",
 	"description": "AK's Generative AI Helper for doing... everything",
-	"version": "2.0.2",
+	"version": "2.0.4",
 	"main": "index.js",
 	"files": [
 		"index.js",
@@ -14,6 +14,7 @@
 		"tool-agent.js",
 		"code-agent.js",
 		"rag-agent.js",
+		"embedding.js",
 		"json-helpers.js",
 		"types.d.ts",
 		"logger.js",

package/transformer.js

@@ -96,27 +96,9 @@ class Transformer extends BaseGemini {
 		this.maxRetries = options.maxRetries || 3;
 		this.retryDelay = options.retryDelay || 1000;
 
-		// ── Grounding ──
-		this.enableGrounding = options.enableGrounding || false;
-		this.groundingConfig = options.groundingConfig || {};
-
 		log.debug(`Transformer keys — Source: "${this.promptKey}", Target: "${this.answerKey}", Context: "${this.contextKey}"`);
 	}
 
-	// ── Chat Create Options Override ──────────────────────────────────────────
-
-	/** @protected */
-	_getChatCreateOptions() {
-		const opts = super._getChatCreateOptions();
-
-		if (this.enableGrounding) {
-			opts.config.tools = [{ googleSearch: this.groundingConfig }];
-			log.debug(`Search grounding ENABLED (WARNING: costs $35/1k queries)`);
-		}
-
-		return opts;
-	}
-
 	// ── Seeding ──────────────────────────────────────────────────────────────
 
 	/**

package/types.d.ts

@@ -32,6 +32,23 @@ export interface ChatConfig {
   [key: string]: any;
 }
 
+export interface GroundingChunk {
+  web?: { uri?: string; title?: string; domain?: string };
+}
+
+export interface GroundingSupport {
+  segment?: any;
+  groundingChunkIndices?: number[];
+  confidenceScores?: number[];
+}
+
+export interface GroundingMetadata {
+  groundingChunks?: GroundingChunk[];
+  groundingSupports?: GroundingSupport[];
+  webSearchQueries?: string[];
+  searchEntryPoint?: { renderedContent?: string };
+}
+
 export interface ResponseMetadata {
   modelVersion: string | null;
   requestedModel: string;
@@ -39,6 +56,7 @@ export interface ResponseMetadata {
   responseTokens: number;
   totalTokens: number;
   timestamp: number;
+  groundingMetadata?: GroundingMetadata | null;
 }
 
 export interface UsageData {
@@ -55,6 +73,7 @@ export interface UsageData {
   /** Model you requested (e.g., 'gemini-2.5-flash') */
   requestedModel: string;
   timestamp: number;
+  groundingMetadata?: GroundingMetadata | null;
 }
 
 export interface TransformationExample {
@@ -77,6 +96,38 @@ export interface GoogleAuthOptions {
   universeDomain?: string;
 }
 
+export interface CacheConfig {
+  /** Model to cache for (defaults to instance modelName) */
+  model?: string;
+  /** Time-to-live duration (e.g., '3600s') */
+  ttl?: string;
+  /** Human-readable display name */
+  displayName?: string;
+  /** Content to cache */
+  contents?: any[];
+  /** System instruction to cache (defaults to instance systemPrompt) */
+  systemInstruction?: string;
+  /** Tools to cache */
+  tools?: any[];
+  /** Tool configuration to cache */
+  toolConfig?: any;
+}
+
+export interface CachedContentInfo {
+  /** Server-generated resource name */
+  name: string;
+  /** User-provided display name */
+  displayName?: string;
+  /** Model this cache is for */
+  model: string;
+  /** Creation timestamp */
+  createTime: string;
+  /** Expiration timestamp */
+  expireTime: string;
+  /** Cache usage metadata */
+  usageMetadata?: { totalTokenCount?: number };
+}
+
 export type AsyncValidatorFunction = (payload: Record<string, unknown>) => Promise<unknown>;
 export type LogLevel = 'trace' | 'debug' | 'info' | 'warn' | 'error' | 'fatal' | 'none';
 
@@ -110,6 +161,14 @@ export interface BaseGeminiOptions {
 
   /** Billing labels for cost segmentation (Vertex AI only) */
   labels?: Record<string, string>;
+
+  /** Enable Google Search grounding (WARNING: costs $35/1k queries) */
+  enableGrounding?: boolean;
+  /** Google Search grounding configuration (searchTypes, excludeDomains, timeRangeFilter) */
+  groundingConfig?: Record<string, any>;
+
+  /** Cached content resource name to use for this session */
+  cachedContent?: string;
 }
 
 export interface TransformerOptions extends BaseGeminiOptions {
@@ -158,6 +217,35 @@ export interface MessageOptions extends BaseGeminiOptions {
   responseMimeType?: string;
 }
 
+export type EmbeddingTaskType = 'RETRIEVAL_DOCUMENT' | 'RETRIEVAL_QUERY' | 'SEMANTIC_SIMILARITY' | 'CLUSTERING' | 'CLASSIFICATION' | 'QUESTION_ANSWERING' | 'FACT_VERIFICATION';
+
+export interface EmbeddingOptions extends BaseGeminiOptions {
+  /** Embedding task type (affects how embeddings are optimized) */
+  taskType?: EmbeddingTaskType;
+  /** Title for the document being embedded (only with RETRIEVAL_DOCUMENT) */
+  title?: string;
+  /** Output dimensionality for the embedding vector */
+  outputDimensionality?: number;
+  /** Whether to auto-truncate long inputs (default: true) */
+  autoTruncate?: boolean;
+}
+
+export interface EmbedConfig {
+  /** Override task type for this call */
+  taskType?: EmbeddingTaskType;
+  /** Override title for this call */
+  title?: string;
+  /** Override output dimensionality for this call */
+  outputDimensionality?: number;
+}
+
+export interface EmbeddingResult {
+  /** The embedding vector */
+  values?: number[];
+  /** Embedding statistics (Vertex AI) */
+  statistics?: { tokenCount?: number; truncated?: boolean };
+}
+
 /** Tool declaration in @google/genai FunctionDeclaration format */
 export interface ToolDeclaration {
   name: string;
@@ -374,6 +462,9 @@ export declare class BaseGemini {
   exampleCount: number;
   labels: Record<string, string>;
   vertexai: boolean;
+  enableGrounding: boolean;
+  groundingConfig: Record<string, any>;
+  cachedContent: string | null;
 
   init(force?: boolean): Promise<void>;
   seed(examples?: TransformationExample[], opts?: SeedOptions): Promise<any[]>;
@@ -388,6 +479,14 @@ export declare class BaseGemini {
     estimatedInputCost: number;
     note: string;
   }>;
+
+  // Context Caching
+  createCache(config?: CacheConfig): Promise<CachedContentInfo>;
+  getCache(cacheName: string): Promise<CachedContentInfo>;
+  listCaches(): Promise<CachedContentInfo[]>;
+  updateCache(cacheName: string, config?: { ttl?: string; expireTime?: string }): Promise<CachedContentInfo>;
+  deleteCache(cacheName: string): Promise<void>;
+  useCache(cacheName: string): Promise<void>;
 }
 
 export declare class Transformer extends BaseGemini {
@@ -401,8 +500,6 @@ export declare class Transformer extends BaseGemini {
   asyncValidator: AsyncValidatorFunction | null;
   maxRetries: number;
   retryDelay: number;
-  enableGrounding: boolean;
-
   seed(examples?: TransformationExample[]): Promise<any[]>;
   send(payload: Record<string, unknown> | string, opts?: SendOptions, validatorFn?: AsyncValidatorFunction | null): Promise<Record<string, unknown>>;
   rawSend(payload: Record<string, unknown> | string, messageOptions?: { labels?: Record<string, string> }): Promise<Record<string, unknown>>;
@@ -496,6 +593,23 @@ export declare class CodeAgent extends BaseGemini {
   stop(): void;
 }
 
+export declare class Embedding extends BaseGemini {
+  constructor(options?: EmbeddingOptions);
+
+  taskType: EmbeddingTaskType | null;
+  title: string | null;
+  outputDimensionality: number | null;
+  autoTruncate: boolean;
+
+  init(force?: boolean): Promise<void>;
+  /** Embed a single text string */
+  embed(text: string, config?: EmbedConfig): Promise<EmbeddingResult>;
+  /** Embed multiple text strings in a single API call */
+  embedBatch(texts: string[], config?: EmbedConfig): Promise<EmbeddingResult[]>;
+  /** Compute cosine similarity between two embedding vectors (-1 to 1) */
+  similarity(a: number[], b: number[]): number;
+}
+
 // ── Module Exports ───────────────────────────────────────────────────────────
 
 export declare function extractJSON(text: string): any;
@@ -508,6 +622,7 @@ declare const _default: {
   ToolAgent: typeof ToolAgent;
   CodeAgent: typeof CodeAgent;
   RagAgent: typeof RagAgent;
+  Embedding: typeof Embedding;
 };
 
 export default _default;