ak-gemini 2.4.1 → 2.6.0

package/json-helpers.js

@@ -267,6 +267,111 @@ function findCompleteJSONStructures(text) {
 	return results;
 }
 
+/**
+ * Validates a parsed value against a (subset of) JSON Schema. Dependency-light —
+ * intended to guard fallback structured-output paths where the model isn't forced
+ * to emit schema-valid JSON.
+ *
+ * Supported keywords: `type` (string or array of strings; 'integer' checks for
+ * whole numbers), `required`, `properties`, `additionalProperties: false`,
+ * `items` (single schema), `enum`. Unknown keywords are ignored (lenient).
+ *
+ * @param {*} data - The parsed value to validate
+ * @param {Object} schema - JSON Schema object
+ * @param {string} [path='$'] - Internal: JSON path prefix for error messages
+ * @returns {string[]} Array of human-readable error strings; empty means valid.
+ */
+export function validateSchema(data, schema, path = '$') {
+	// NOTE: ak-gemini enforces structured output natively (responseSchema), so this
+	// validator has no runtime call site here — it exists for cross-package API
+	// symmetry with ak-claude (whose Vertex fallback path relies on it). Keep this
+	// implementation byte-identical to ak-claude/json-helpers.js; apply fixes to both.
+	const errors = [];
+	if (!schema || typeof schema !== 'object') return errors;
+
+	// ── nullable (OpenAPI-style) ── a null value is valid on a nullable node.
+	if (data === null && schema.nullable === true) return errors;
+
+	// ── enum ── strict equality first, then deep-equal for object/array members.
+	if (Array.isArray(schema.enum)) {
+		const target = JSON.stringify(data);
+		const ok = schema.enum.some(v => v === data || JSON.stringify(v) === target);
+		if (!ok) errors.push(`${path}: value ${JSON.stringify(data)} is not one of allowed enum values ${JSON.stringify(schema.enum)}`);
+	}
+
+	// ── type ──
+	if (schema.type !== undefined) {
+		const types = Array.isArray(schema.type) ? schema.type : [schema.type];
+		if (schema.nullable === true) types.push('null');
+		if (!types.some(t => matchesType(data, t))) {
+			errors.push(`${path}: expected type ${types.join('|')} but got ${describeType(data)}`);
+			// If the type is wrong, deeper checks are meaningless — stop here for this node.
+			return errors;
+		}
+	}
+
+	// ── object ── (Object.hasOwn avoids prototype-chain keys like 'toString')
+	const isObject = data !== null && typeof data === 'object' && !Array.isArray(data);
+	if (isObject && (schema.properties || schema.required || schema.additionalProperties === false)) {
+		const props = schema.properties || {};
+
+		if (Array.isArray(schema.required)) {
+			for (const key of schema.required) {
+				if (!Object.hasOwn(data, key)) errors.push(`${path}: missing required property "${key}"`);
+			}
+		}
+
+		if (schema.additionalProperties === false) {
+			for (const key of Object.keys(data)) {
+				if (!Object.hasOwn(props, key)) errors.push(`${path}: unexpected property "${key}" (additionalProperties is false)`);
+			}
+		}
+
+		for (const [key, subSchema] of Object.entries(props)) {
+			if (Object.hasOwn(data, key)) {
+				errors.push(...validateSchema(data[key], /** @type {any} */ (subSchema), `${path}.${key}`));
+			}
+		}
+	}
+
+	// ── array ──
+	if (Array.isArray(data) && schema.items && typeof schema.items === 'object' && !Array.isArray(schema.items)) {
+		data.forEach((item, i) => {
+			errors.push(...validateSchema(item, schema.items, `${path}[${i}]`));
+		});
+	}
+
+	return errors;
+}
+
+/**
+ * @param {*} value
+ * @param {string} type - JSON Schema type name
+ * @returns {boolean}
+ */
+function matchesType(value, type) {
+	switch (type) {
+		case 'string': return typeof value === 'string';
+		case 'number': return typeof value === 'number' && !Number.isNaN(value);
+		case 'integer': return typeof value === 'number' && Number.isInteger(value);
+		case 'boolean': return typeof value === 'boolean';
+		case 'object': return value !== null && typeof value === 'object' && !Array.isArray(value);
+		case 'array': return Array.isArray(value);
+		case 'null': return value === null;
+		default: return true; // unknown type keyword — be lenient
+	}
+}
+
+/**
+ * @param {*} value
+ * @returns {string}
+ */
+function describeType(value) {
+	if (value === null) return 'null';
+	if (Array.isArray(value)) return 'array';
+	return typeof value;
+}
+
 /**
  * Extracts valid JSON from model response text using multiple strategies.
  * @param {string} text - The model response text

package/message.js

@@ -53,6 +53,13 @@ class Message extends BaseGemini {
 		}
 		if (options.responseMimeType) {
 			this.chatConfig.responseMimeType = options.responseMimeType;
+		} else if (options.responseSchema) {
+			// Google's @google/genai requires response_mime_type: 'application/json'
+			// whenever a responseSchema is set — it does NOT add it automatically.
+			// Passing a schema without the mime type yields a 400 INVALID_ARGUMENT
+			// (Vertex) or prose output that fails JSON extraction. Auto-pair them.
+			this.chatConfig.responseMimeType = 'application/json';
+			log.debug("responseSchema set without responseMimeType — defaulting responseMimeType to 'application/json'.");
 		}
 
 		this._isStructured = !!(options.responseSchema || options.responseMimeType === 'application/json');
@@ -71,12 +78,11 @@ class Message extends BaseGemini {
 
 		log.debug(`Initializing ${this.constructor.name} with model: ${this.modelName}...`);
 
-		try {
-			await this.genAIClient.models.list();
-			log.debug(`${this.constructor.name}: API connection successful.`);
-		} catch (e) {
-			throw new Error(`${this.constructor.name} initialization failed: ${e.message}`);
-		}
+		// Connectivity check is opt-in (healthCheck: true). models.list() is a
+		// different IAM surface than generateContent — a service account allowed to
+		// generate but not list publisher models would fail here misleadingly — and
+		// it adds a round-trip per instance. The first send() is a fine error signal.
+		await this._healthCheckPing();
 
 		this._initialized = true;
 		log.debug(`${this.constructor.name}: Initialized (stateless mode).`);
@@ -111,11 +117,18 @@ class Message extends BaseGemini {
 			}
 		}));
 
+		// Compute per-call usage synchronously from THIS response before any other
+		// concurrent send() can mutate instance state. result.usage is safe under
+		// concurrency; getLastUsage() (instance state, updated below) is not.
+		const usage = this._usageFromResponse(result);
+
 		this._captureMetadata(result);
 
 		this._cumulativeUsage = {
 			promptTokens: this.lastResponseMetadata.promptTokens,
 			responseTokens: this.lastResponseMetadata.responseTokens,
+			thoughtsTokens: this.lastResponseMetadata.thoughtsTokens,
+			cachedTokens: this.lastResponseMetadata.cachedTokens,
 			totalTokens: this.lastResponseMetadata.totalTokens,
 			attempts: 1
 		};
@@ -127,7 +140,7 @@ class Message extends BaseGemini {
 		const text = result.text || '';
 		const response = {
 			text,
-			usage: this.getLastUsage()
+			usage
 		};
 
 		// Parse structured data if configured

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.4.1",
+	"version": "2.6.0",
 	"main": "index.js",
 	"files": [
 		"index.js",
@@ -19,7 +19,9 @@
 		"json-helpers.js",
 		"types.d.ts",
 		"logger.js",
-		"GUIDE.md"
+		"GUIDE.md",
+		"UPGRADING.md",
+    "CHANGELOG.md"
 	],
 	"types": "types.d.ts",
 	"exports": {
@@ -65,7 +67,7 @@
 	],
 	"license": "ISC",
 	"dependencies": {
-		"@google/genai": "^2.10.0",
+		"@google/genai": "^2.12.0",
 		"dotenv": "^17.3.1",
 		"pino": "^10.3.1",
 		"pino-pretty": "^13.1.3"

package/rag-agent.js

@@ -213,6 +213,8 @@ class RagAgent extends BaseGemini {
 		this._cumulativeUsage = {
 			promptTokens: this.lastResponseMetadata.promptTokens,
 			responseTokens: this.lastResponseMetadata.responseTokens,
+			thoughtsTokens: this.lastResponseMetadata.thoughtsTokens,
+			cachedTokens: this.lastResponseMetadata.cachedTokens,
 			totalTokens: this.lastResponseMetadata.totalTokens,
 			attempts: 1
 		};

package/tool-agent.js

@@ -200,6 +200,8 @@ class ToolAgent extends BaseGemini {
 		this._cumulativeUsage = {
 			promptTokens: this.lastResponseMetadata.promptTokens,
 			responseTokens: this.lastResponseMetadata.responseTokens,
+			thoughtsTokens: this.lastResponseMetadata.thoughtsTokens,
+			cachedTokens: this.lastResponseMetadata.cachedTokens,
 			totalTokens: this.lastResponseMetadata.totalTokens,
 			attempts: 1
 		};

package/transformer.js

@@ -201,7 +201,7 @@ class Transformer extends BaseGemini {
 		if (opts.labels) messageOptions.labels = opts.labels;
 
 		// Reset cumulative usage tracking
-		this._cumulativeUsage = { promptTokens: 0, responseTokens: 0, totalTokens: 0, attempts: 0 };
+		this._cumulativeUsage = { promptTokens: 0, responseTokens: 0, thoughtsTokens: 0, cachedTokens: 0, totalTokens: 0, attempts: 0 };
 
 		let lastError = null;
 
@@ -215,6 +215,8 @@ class Transformer extends BaseGemini {
 				if (this.lastResponseMetadata) {
 					this._cumulativeUsage.promptTokens += this.lastResponseMetadata.promptTokens || 0;
 					this._cumulativeUsage.responseTokens += this.lastResponseMetadata.responseTokens || 0;
+					this._cumulativeUsage.thoughtsTokens += this.lastResponseMetadata.thoughtsTokens || 0;
+					this._cumulativeUsage.cachedTokens += this.lastResponseMetadata.cachedTokens || 0;
 					this._cumulativeUsage.totalTokens += this.lastResponseMetadata.totalTokens || 0;
 					this._cumulativeUsage.attempts = attempt + 1;
 				}
@@ -389,6 +391,8 @@ Respond with JSON only – no comments or explanations.
 		this._cumulativeUsage = {
 			promptTokens: this.lastResponseMetadata.promptTokens,
 			responseTokens: this.lastResponseMetadata.responseTokens,
+			thoughtsTokens: this.lastResponseMetadata.thoughtsTokens,
+			cachedTokens: this.lastResponseMetadata.cachedTokens,
 			totalTokens: this.lastResponseMetadata.totalTokens,
 			attempts: 1
 		};
@@ -422,7 +426,7 @@ Respond with JSON only – no comments or explanations.
 		this.chatSession = this._createChatSession(exampleHistory);
 
 		this.lastResponseMetadata = null;
-		this._cumulativeUsage = { promptTokens: 0, responseTokens: 0, totalTokens: 0, attempts: 0 };
+		this._cumulativeUsage = { promptTokens: 0, responseTokens: 0, thoughtsTokens: 0, cachedTokens: 0, totalTokens: 0, attempts: 0 };
 
 		log.debug(`Conversation cleared. Preserved ${exampleHistory.length} example items.`);
 	}

package/types.d.ts

@@ -6,11 +6,38 @@ export { ThinkingLevel, HarmCategory, HarmBlockThreshold };
 
 export interface ThinkingConfig {
   includeThoughts?: boolean;
-  /** Token budget for thinking. 0 = disabled, -1 = automatic. */
+  /** Token budget for thinking. 0 = disabled, -1 = automatic. Mutually exclusive with thinkingLevel. */
   thinkingBudget?: number;
   thinkingLevel?: ThinkingLevel;
 }
 
+/**
+ * Cross-package reasoning-effort levels (same option name in ak-claude).
+ * Gemini's wire field is `thinkingConfig.thinkingLevel`; `xhigh` and `max` are
+ * ak-claude-only levels and clamp to `high` here.
+ */
+export type EffortLevel = 'minimal' | 'low' | 'medium' | 'high' | 'xhigh' | 'max';
+
+/** Per-million-token rates for a single model. */
+export interface ModelPricing {
+  input: number;
+  output: number;
+  /** Rate for tokens served from a context cache. Absent when Google publishes none — cost then falls back to `input` (overestimate). */
+  cachedInput?: number;
+  /** >200k-context tier rates, present only on tiered (Pro) models. */
+  inputAbove200k?: number;
+  outputAbove200k?: number;
+  cachedInputAbove200k?: number;
+}
+
+/** ModelPricing decorated with the tier threshold and the table's verification date. */
+export interface ResolvedPricing extends ModelPricing {
+  /** Prompt-token count above which the *Above200k rates apply. */
+  tierThreshold: number;
+  /** ISO date the pricing table was last verified against Google's pricing page. */
+  asOf: string;
+}
+
 export interface SafetySetting {
   category: HarmCategory;
   threshold: HarmBlockThreshold;
@@ -54,6 +81,9 @@ export interface ResponseMetadata {
   requestedModel: string;
   promptTokens: number;
   responseTokens: number;
+  thoughtsTokens?: number;
+  /** Cached-content tokens. A SUBSET of promptTokens, not an addition to it. */
+  cachedTokens?: number;
   totalTokens: number;
   timestamp: number;
   groundingMetadata?: GroundingMetadata | null;
@@ -64,7 +94,17 @@ export interface UsageData {
   promptTokens: number;
   /** CUMULATIVE output tokens across all retry attempts */
   responseTokens: number;
-  /** CUMULATIVE total tokens across all retry attempts */
+  /** CUMULATIVE thinking ("thoughts") tokens; billed at the output rate. Included in totalTokens and estimatedCost. */
+  thoughtsTokens?: number;
+  /**
+   * CUMULATIVE cached-content tokens (`usageMetadata.cachedContentTokenCount`).
+   * NOTE: Gemini INCLUDES these in promptTokens — they are a subset, not an
+   * addition. `estimatedCost` subtracts them from promptTokens and re-bills them
+   * at the discounted cached rate. (ak-claude is the opposite: Anthropic's cache
+   * tokens are excluded from input_tokens and added on top.)
+   */
+  cachedTokens?: number;
+  /** CUMULATIVE total tokens across all retry attempts (includes thoughtsTokens) */
   totalTokens: number;
   /** Number of attempts (1 = first try success, 2+ = retries needed) */
   attempts: number;
@@ -76,6 +116,12 @@ export interface UsageData {
   groundingMetadata?: GroundingMetadata | null;
   /** Model lifecycle status from Google (e.g., 'DEPRECATED'). Surfaced from @google/genai 1.47+. */
   modelStatus?: string | null;
+  /**
+   * Estimated USD cost from MODEL_PRICING (input + cached input + output + thoughts).
+   * `null` means the model's pricing is UNKNOWN — it does NOT mean the call was free.
+   * Cache STORAGE (per token-hour) and image-output tokens are not modelled.
+   */
+  estimatedCost?: number | null;
 }
 
 export interface TransformationExample {
@@ -145,6 +191,12 @@ export interface BaseGeminiOptions {
   chatConfig?: Partial<ChatConfig>;
   /** Thinking features configuration */
   thinkingConfig?: ThinkingConfig | null;
+  /**
+   * Reasoning effort (cross-package option; same name in ak-claude). Maps to
+   * `thinkingConfig.thinkingLevel` and wins over an explicit thinkingLevel /
+   * thinkingBudget. Throws on an unrecognized level.
+   */
+  effort?: EffortLevel | null;
   /** Maximum output tokens (default: 50000, null removes limit) */
   maxOutputTokens?: number | null;
   /** Sampling temperature (shorthand for chatConfig.temperature; wins over chatConfig) */
@@ -592,6 +644,8 @@ export declare class BaseGemini {
   cachedContent: string | null;
   serviceTier: ServiceTier | null;
   includeServerSideToolInvocations: boolean;
+  /** Normalized `effort` option; null when unset. */
+  effort: EffortLevel | null;
 
   init(force?: boolean): Promise<void>;
   seed(examples?: TransformationExample[], opts?: SeedOptions): Promise<any[]>;
@@ -602,11 +656,14 @@ export declare class BaseGemini {
   estimateCost(nextPayload: Record<string, unknown> | string): Promise<{
     inputTokens: number;
     model: string;
-    pricing: { input: number; output: number };
-    estimatedInputCost: number;
+    pricing: ResolvedPricing | null;
+    estimatedInputCost: number | null;
     note: string;
   }>;
 
+  /** @internal Validates the `effort` option; throws on an unknown level. */
+  protected _normalizeEffort(effort?: string | null): EffortLevel | null;
+
   // Context Caching
   createCache(config?: CacheConfig): Promise<CachedContentInfo>;
   getCache(cacheName: string): Promise<CachedContentInfo>;
@@ -757,6 +814,32 @@ export declare class ImageGenerator extends BaseGemini {
 
 export declare function extractJSON(text: string): any;
 export declare function attemptJSONRecovery(text: string, maxAttempts?: number): any | null;
+/** Validates a parsed value against a subset of JSON Schema. Returns error strings ([] means valid). */
+export declare function validateSchema(data: any, schema: Record<string, any>, path?: string): string[];
+
+/** Per-million-token pricing keyed by model id. */
+export declare const MODEL_PRICING: Record<string, ModelPricing>;
+/** ISO date the pricing table was last verified against Google's published rates. */
+export declare const MODEL_PRICING_AS_OF: string;
+/** Floating `-latest` alias → canonical model id, for pricing resolution. */
+export declare const MODEL_ALIASES: Record<string, string>;
+/** Valid values for the `effort` option. */
+export declare const EFFORT_LEVELS: readonly EffortLevel[];
+/** Resolves pricing for a model id (follows -latest aliases). `null` means UNKNOWN, not free. */
+export declare function resolvePricing(modelId: string | null | undefined): ResolvedPricing | null;
+/**
+ * Estimated USD cost from token counts. `null` means UNKNOWN, not free.
+ * `cachedTokens` is a SUBSET of `promptTokens` (Gemini includes them) and is
+ * subtracted before billing at the cached rate. Thinking tokens bill at the
+ * output rate. Tier is selected from `promptTokens` vs 200k.
+ */
+export declare function computeCost(
+  modelId: string | null | undefined,
+  promptTokens: number,
+  responseTokens: number,
+  thoughtsTokens?: number,
+  cachedTokens?: number
+): number | null;
 
 declare const _default: {
   Transformer: typeof Transformer;