@kalaa/node 1.0.7 → 1.0.9

package/README.md

@@ -55,9 +55,11 @@ 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 
 
+
+> **Order matters for Gemini:** call `kalaa.observe()` before you `require('@google/genai')` anywhere in your app. (OpenAI and Anthropic aren't order-sensitive — only the current Gemini SDK's instance-bound methods require this.)
 ## Configuration
 
 | Option | Env var | Default |

package/costmeter.js

@@ -475,80 +475,90 @@ function patchAnthropicModule() {
   Anthropic.__cm_patched = true;
   log('anthropic module patched (prototype-level)');
 }
-
-function patchGeminiModule() {
+// New SDK: @google/genai — generateContent/generateContentStream are bound
+// per-instance in the Models constructor, not on Models.prototype. Prototype
+// patching is a no-op here; we wrap the instance's own methods instead, right
+// after GoogleGenAI's constructor creates them.
+function patchGeminiUnified() {
   let Mod;
-  try { Mod = require('@google/generative-ai'); } catch { return; } // not installed, skip
-  const GoogleGenerativeAI = Mod.GoogleGenerativeAI;
-  if (!GoogleGenerativeAI || GoogleGenerativeAI.__cm_patched) return;
+  try { Mod = require('@google/genai'); } catch { return; } // not installed, skip
+  const OrigGoogleGenAI = Mod.GoogleGenAI;
+  if (!OrigGoogleGenAI || OrigGoogleGenAI.__cm_patched) return;
+
+  function wrapModelsInstance(models) {
+    if (!models || models.__cm_wrapped) return;
+
+    if (typeof models.generateContent === 'function' && !models.generateContent.__cm) {
+      const orig = models.generateContent.bind(models);
+      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(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.generateContent = wrapped;
+    }
 
-  // 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);
-    return;
-  }
+    if (typeof models.generateContentStream === 'function' && !models.generateContentStream.__cm) {
+      const orig = models.generateContentStream.bind(models);
+      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(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.generateContentStream = wrapped;
+    }
 
-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.
-      let _user, clean = params;
-      if (params && typeof params === 'object' && !Array.isArray(params) && '_user' in params) {
-        ({ _user, ...clean } = params);
-      }
-      try {
-        const resp = await orig.call(this, clean, ...rest);
-        const usage = resp?.response?.usageMetadata || {};
-        track({ provider: 'gemini', model: modelName, usage, 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;
-    GenerativeModel.prototype.generateContent = wrapped;
+    models.__cm_wrapped = true;
   }
 
-  // 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) {
-    const orig = GenerativeModel.prototype.generateContentStream;
-    const wrapped = async function (params, ...rest) {
-      const t0 = Date.now();
-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);
-      }
-      try {
-        const result = await orig.call(this, clean, ...rest);
-        result.response.then(finalResp => {
-          const usage = finalResp?.usageMetadata || {};
-          track({ provider: 'gemini', model: modelName, usage, latencyMs: Date.now() - t0, status: Object.keys(usage).length ? 'ok' : 'no_usage_streamed', endUser: _user });
-        }).catch(() => {});
-        return result;
-      } 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;
-    GenerativeModel.prototype.generateContentStream = wrapped;
-  }
+  // Proxy the constructor: every `new GoogleGenAI(...)` still runs the real
+  // constructor untouched, we just reach into the instance it produces and
+  // wrap `.models` before handing the instance back.
+  const PatchedGoogleGenAI = new Proxy(OrigGoogleGenAI, {
+    construct(Target, args) {
+      const instance = new Target(...args);
+      wrapModelsInstance(instance.models);
+      return instance;
+    },
+  });
+  PatchedGoogleGenAI.__cm_patched = true;
 
-  GoogleGenerativeAI.__cm_patched = true;
-  log('gemini module patched (prototype-level)');
+  Mod.GoogleGenAI = PatchedGoogleGenAI;
+  log('gemini (unified @google/genai) module patched (instance-level via constructor proxy)');
 }
 
 // ── PUBLIC API ───────────────────────────────────────────────────────

package/package.json

@@ -1,6 +1,6 @@
 {
 "name": "@kalaa/node",
-  "version": "1.0.7",
+  "version": "1.0.9",
   "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 }
   }
 }