cloudbroker-sdk 0.5.0 → 0.7.0

package/src/wrapper.js

@@ -4,7 +4,7 @@
  */
 
 import { CloudBrokerClient } from "./client.js";
-import { PRICING } from "./pricing.js";
+import { PRICING, priceGemini } from "./pricing.js";
 import { cbStore } from "./context.js";
 
 export class CloudBrokerAI {
@@ -35,12 +35,18 @@ export class CloudBrokerAI {
     // Adiciona trackDeploy ao cliente wrappado
     client.trackDeploy = cb.trackDeploy.bind(cb);
 
-    // Detecta provider pelo tipo do cliente
+    // Detecta provider pelo tipo do cliente.
+    // O sinal do Gemini é `models.generateContent`, não `models`: o cliente da
+    // OpenAI também tem `.models` (list/retrieve), então checar só o objeto
+    // casaria com o cliente errado.
     const isOpenAI    = client?.constructor?.name === "OpenAI"    || !!client?.chat?.completions;
     const isAnthropic = client?.constructor?.name === "Anthropic" || !!client?.messages?.create;
+    const isGemini    = client?.constructor?.name === "GoogleGenAI"
+                     || typeof client?.models?.generateContent === "function";
 
     if (isOpenAI)    return CloudBrokerAI.#wrapOpenAI(client, cb, { product, team, debug });
     if (isAnthropic) return CloudBrokerAI.#wrapAnthropic(client, cb, { product, team, debug });
+    if (isGemini)    return CloudBrokerAI.#wrapGemini(client, cb, { product, team, debug });
 
     if (debug) console.warn("[CloudBroker] Provider não reconhecido — retornando cliente original.");
     return client;
@@ -254,6 +260,123 @@ export class CloudBrokerAI {
     return client;
   }
 
+  /**
+   * Instrumenta o cliente do @google/genai (`new GoogleGenAI({apiKey})`).
+   *
+   * Diferenças relevantes em relação a OpenAI/Anthropic:
+   *
+   *  - A API é de argumento único (`models.generateContent({model, contents})`),
+   *    sem segundo parâmetro de options. Por isso o `cbLabel` é aceito das duas
+   *    formas: no segundo argumento (convenção dos outros wrappers) e dentro do
+   *    próprio params — neste caso removido antes de repassar, pra nunca vazar
+   *    campo desconhecido para a API do provider.
+   *  - O usage vem em `usageMetadata` (camelCase). Os nomes foram conferidos nos
+   *    tipos do @google/genai@2.15.0, não na doc pública: a doc mostra um
+   *    formato snake_case que pertence a outra superfície da API, e usá-lo
+   *    daria token zero em silêncio.
+   *  - `costUSD` pode ser `null` (modelo fora da tabela de preço). Isso é
+   *    deliberado: registrar tokens sem inventar preço. Quem consome precisa
+   *    olhar `unpriced` — tratar `null` como zero vira "custo real zero".
+   */
+  static #wrapGemini(client, cb, { product, team, debug }) {
+    const models = client.models;
+    if (!models || typeof models.generateContent !== "function") return client;
+
+    // Extrai o cbLabel e devolve os params limpos pra repassar ao provider.
+    const splitLabel = (params, options) => {
+      let label = null, fwd = params;
+      try {
+        label = options?.cbLabel ?? null;
+        if (!label && params && typeof params === "object" && params.cbLabel != null) {
+          label = params.cbLabel;
+        }
+        if (!label) label = cbStore.getStore()?.label ?? null;
+        if (params && typeof params === "object" && "cbLabel" in params) {
+          const { cbLabel: _l, ...rest } = params;
+          fwd = rest;
+        }
+      } catch { /* instrumentação nunca altera a chamada do cliente */ }
+      return { label, fwd };
+    };
+
+    const buildEvent = (model, usage, { label, start, streaming, endpoint }) => {
+      const priced = priceGemini(model, usage || {});
+      return {
+        provider:       "gemini",
+        model,
+        product:        product || null,
+        team:           team    || null,
+        label,
+        inputTokens:    priced.inputTokens,
+        outputTokens:   priced.outputTokens,
+        totalTokens:    priced.totalTokens,
+        costUSD:        priced.costUSD,
+        // Só o Gemini separa esses dois, e são justamente os que mais mexem no
+        // valor: cache barateia ~10x, thinking é cobrado como output.
+        cachedTokens:   priced.cachedTokens,
+        thoughtsTokens: priced.thoughtsTokens,
+        unpriced:       priced.unpriced,
+        pricingNote:    priced.reason,
+        latencyMs:      Date.now() - start,
+        endpoint,
+        streaming:      !!streaming,
+      };
+    };
+
+    // ── generateContent (não-streaming) ───────────────────────────────────────
+    const originalGenerate = models.generateContent.bind(models);
+    models.generateContent = async (params, options) => {
+      const start = Date.now();
+      const { label, fwd } = splitLabel(params, options);
+
+      const res = await originalGenerate(fwd);
+
+      try {
+        const model = res?.modelVersion || params?.model || "unknown";
+        cb.track(buildEvent(model, res?.usageMetadata, {
+          label, start, streaming: false, endpoint: "models.generateContent",
+        }));
+      } catch (err) {
+        if (debug) console.warn("[CloudBroker] track error:", err.message);
+      }
+
+      return res;
+    };
+
+    // ── generateContentStream ────────────────────────────────────────────────
+    if (typeof models.generateContentStream === "function") {
+      const originalStream = models.generateContentStream.bind(models);
+      models.generateContentStream = async (params, options) => {
+        const start = Date.now();
+        const { label, fwd } = splitLabel(params, options);
+
+        const stream = await originalStream(fwd);
+
+        let model = params?.model || "unknown";
+        let usage = null;
+        const onChunk = (chunk) => {
+          if (chunk?.modelVersion) model = chunk.modelVersion;
+          // O usage do Gemini é cumulativo e chega no(s) chunk(s) final(is):
+          // o último que vier é o total da chamada.
+          if (chunk?.usageMetadata) usage = chunk.usageMetadata;
+        };
+        const onDone = () => {
+          try {
+            cb.track(buildEvent(model, usage, {
+              label, start, streaming: true, endpoint: "models.generateContentStream",
+            }));
+          } catch (err) {
+            if (debug) console.warn("[CloudBroker] track error:", err.message);
+          }
+        };
+
+        return CloudBrokerAI.#instrumentStream(stream, onChunk, onDone);
+      };
+    }
+
+    return client;
+  }
+
   /**
    * Registra um deploy para análise de custo
    * Atalho para CloudBrokerClient.trackDeploy() sem precisar envolver um cliente de IA