@kalaa/node 1.0.7 → 1.0.8

package/README.md

@@ -55,7 +55,7 @@ kalaa.withContext('nightly-digest', () => {
 
 - OpenAI: chat completions, including streaming (`stream_options.include_usage`)
 - Anthropic: messages, including streaming
-- Gemini: `generateContent`
+- Gemini: `generateContent` and `generateContentStream` — both `@google/genai` (current) and `@google/generative-ai` (legacy) are supported
 - Whisper/audio transcription duration 
 
 ## Configuration

package/costmeter.js

@@ -475,32 +475,33 @@ function patchAnthropicModule() {
   Anthropic.__cm_patched = true;
   log('anthropic module patched (prototype-level)');
 }
-
 function patchGeminiModule() {
+  patchGeminiLegacy();
+  patchGeminiUnified();
+}
+
+// Old SDK: @google/generative-ai — GoogleGenerativeAI().getGenerativeModel().generateContent()
+function patchGeminiLegacy() {
   let Mod;
   try { Mod = require('@google/generative-ai'); } catch { return; } // not installed, skip
   const GoogleGenerativeAI = Mod.GoogleGenerativeAI;
   if (!GoogleGenerativeAI || GoogleGenerativeAI.__cm_patched) return;
 
-  // Same trick as OpenAI: probe an instance to reach the GenerativeModel
-  // class prototype, patch it once, every model instance inherits it.
   let GenerativeModel;
   try {
     const probe = new GoogleGenerativeAI('cm_probe');
     const model = probe.getGenerativeModel({ model: 'gemini-1.5-flash' });
     GenerativeModel = Object.getPrototypeOf(model).constructor;
   } catch (e) {
-    log('gemini probe failed:', e.message);
+    log('gemini (legacy) probe failed:', e.message);
     return;
   }
 
-if (GenerativeModel?.prototype?.generateContent && !GenerativeModel.prototype.generateContent.__cm) {
+  if (GenerativeModel?.prototype?.generateContent && !GenerativeModel.prototype.generateContent.__cm) {
     const orig = GenerativeModel.prototype.generateContent;
     const wrapped = async function (params, ...rest) {
       const t0 = Date.now();
-const modelName = (this.model || 'gemini').replace(/^models\//, '');
-      // Gemini params can be a plain object, a string, or an array — only
-      // strip _user when it's actually a plain request object carrying one.
+      const modelName = (this.model || 'gemini').replace(/^models\//, '');
       let _user, clean = params;
       if (params && typeof params === 'object' && !Array.isArray(params) && '_user' in params) {
         ({ _user, ...clean } = params);
@@ -519,14 +520,11 @@ const modelName = (this.model || 'gemini').replace(/^models\//, '');
     GenerativeModel.prototype.generateContent = wrapped;
   }
 
-  // generateContentStream: usage arrives only after the stream is fully
-  // consumed, via response.usageMetadata on the final resolved response
-  // object the SDK exposes — same tap pattern as OpenAI/Anthropic streaming.
-if (GenerativeModel?.prototype?.generateContentStream && !GenerativeModel.prototype.generateContentStream.__cm) {
+  if (GenerativeModel?.prototype?.generateContentStream && !GenerativeModel.prototype.generateContentStream.__cm) {
     const orig = GenerativeModel.prototype.generateContentStream;
     const wrapped = async function (params, ...rest) {
       const t0 = Date.now();
-const modelName = (this.model || 'gemini').replace(/^models\//, '');
+      const modelName = (this.model || 'gemini').replace(/^models\//, '');
       let _user, clean = params;
       if (params && typeof params === 'object' && !Array.isArray(params) && '_user' in params) {
         ({ _user, ...clean } = params);
@@ -548,7 +546,81 @@ const modelName = (this.model || 'gemini').replace(/^models\//, '');
   }
 
   GoogleGenerativeAI.__cm_patched = true;
-  log('gemini module patched (prototype-level)');
+  log('gemini (legacy @google/generative-ai) module patched (prototype-level)');
+}
+
+// New SDK: @google/genai — new GoogleGenAI({apiKey}).models.generateContent({model, contents})
+function patchGeminiUnified() {
+  let Mod;
+  try { Mod = require('@google/genai'); } catch { return; } // not installed, skip
+  const GoogleGenAI = Mod.GoogleGenAI;
+  if (!GoogleGenAI || GoogleGenAI.__cm_patched) return;
+
+  let Models;
+  try {
+    const probe = new GoogleGenAI({ apiKey: 'cm_probe' });
+    Models = Object.getPrototypeOf(probe.models).constructor;
+  } catch (e) {
+    log('gemini (unified) probe failed:', e.message);
+    return;
+  }
+
+  if (Models?.prototype?.generateContent && !Models.prototype.generateContent.__cm) {
+    const orig = Models.prototype.generateContent;
+    const wrapped = async function (params, ...rest) {
+      const t0 = Date.now();
+      const { _user, ...clean } = params || {};
+      const modelName = (clean.model || 'gemini').replace(/^models\//, '');
+      try {
+        const resp = await orig.call(this, clean, ...rest);
+        track({ provider: 'gemini', model: modelName, usage: resp?.usageMetadata || {}, latencyMs: Date.now() - t0, endUser: _user });
+        return resp;
+      } catch (err) {
+        track({ provider: 'gemini', model: modelName, usage: {}, latencyMs: Date.now() - t0, status: 'error', meta: { error: String(err.message).slice(0, 200) }, endUser: _user });
+        throw err;
+      }
+    };
+    wrapped.__cm = true;
+    Models.prototype.generateContent = wrapped;
+  }
+
+  if (Models?.prototype?.generateContentStream && !Models.prototype.generateContentStream.__cm) {
+    const orig = Models.prototype.generateContentStream;
+    const wrapped = async function (params, ...rest) {
+      const t0 = Date.now();
+      const { _user, ...clean } = params || {};
+      const modelName = (clean.model || 'gemini').replace(/^models\//, '');
+      try {
+        const stream = await orig.call(this, clean, ...rest);
+        const origIterator = stream[Symbol.asyncIterator].bind(stream);
+        let lastUsage = null;
+        stream[Symbol.asyncIterator] = function () {
+          const it = origIterator();
+          return {
+            async next() {
+              const result = await it.next();
+              if (!result.done && result.value?.usageMetadata) lastUsage = result.value.usageMetadata;
+              if (result.done) {
+                track({ provider: 'gemini', model: modelName, usage: lastUsage || {}, latencyMs: Date.now() - t0, status: lastUsage ? 'ok' : 'no_usage_streamed', endUser: _user });
+              }
+              return result;
+            },
+            return(v) { return it.return ? it.return(v) : Promise.resolve({ value: v, done: true }); },
+            throw(err) { return it.throw ? it.throw(err) : Promise.reject(err); },
+          };
+        };
+        return stream;
+      } catch (err) {
+        track({ provider: 'gemini', model: modelName, usage: {}, latencyMs: Date.now() - t0, status: 'error', meta: { error: String(err.message).slice(0, 200) }, endUser: _user });
+        throw err;
+      }
+    };
+    wrapped.__cm = true;
+    Models.prototype.generateContentStream = wrapped;
+  }
+
+  GoogleGenAI.__cm_patched = true;
+  log('gemini (unified @google/genai) module patched (prototype-level)');
 }
 
 // ── PUBLIC API ───────────────────────────────────────────────────────

package/package.json

@@ -1,6 +1,6 @@
 {
 "name": "@kalaa/node",
-  "version": "1.0.7",
+  "version": "1.0.8",
   "description": "Drop-in AI cost tracking for Node.js and TypeScript. Auto-instruments OpenAI, Anthropic, and Gemini at the class level, with zero code at each call site.",
   "main": "costmeter.js",
   "types": "costmeter.d.ts",
@@ -39,14 +39,16 @@
     "music-metadata": "^7.0.0",
     "node-fetch": "^2.7.0"
   },
-  "peerDependencies": {
+"peerDependencies": {
     "openai": ">=4",
     "@anthropic-ai/sdk": ">=0.20",
-    "@google/generative-ai": ">=0.10"
+    "@google/generative-ai": ">=0.10",
+    "@google/genai": ">=0.1"
   },
   "peerDependenciesMeta": {
     "openai": { "optional": true },
     "@anthropic-ai/sdk": { "optional": true },
-    "@google/generative-ai": { "optional": true }
+    "@google/generative-ai": { "optional": true },
+    "@google/genai": { "optional": true }
   }
 }