@aikidosec/firewall/sinks/GoogleGenAi.js

Zen by Aikido is an embedded Application Firewall that autonomously protects Node.js apps against common and critical attacks, provides rate limiting, detects malicious traffic (including bots), and more.

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.GoogleGenAi = void 0;
const wrapExport_1 = require("../agent/hooks/wrapExport");
const wrapNewInstance_1 = require("../agent/hooks/wrapNewInstance");
class GoogleGenAi {
    isResponse(response) {
        // Not possible to use isPlainObject because it is a class
        return (!!response &&
            typeof response === "object" &&
            !Array.isArray(response) &&
            "modelVersion" in response &&
            "usageMetadata" in response &&
            typeof response.usageMetadata === "object" &&
            response.usageMetadata !== null);
    }
    getUsageMetadata(usage) {
        var _a, _b, _c, _d, _e, _f;
        const inputTokens = ((_a = usage.promptTokenCount) !== null && _a !== void 0 ? _a : 0) +
            ((_b = usage.cachedContentTokenCount) !== null && _b !== void 0 ? _b : 0) +
            ((_c = usage.toolUsePromptTokenCount) !== null && _c !== void 0 ? _c : 0);
        const outputTokens = ((_d = usage.responseTokenCount) !== null && _d !== void 0 ? _d : 0) +
            ((_e = usage.thoughtsTokenCount) !== null && _e !== void 0 ? _e : 0) +
            ((_f = usage.candidatesTokenCount) !== null && _f !== void 0 ? _f : 0);
        if (inputTokens === 0 && outputTokens === 0) {
            return;
        }
        return {
            inputTokens,
            outputTokens,
        };
    }
    inspectReturnValue(agent, returnValue) {
        if (returnValue instanceof Promise) {
            // Inspect the response after the promise resolves, it won't change the original promise
            returnValue
                .then((response) => {
                this.inspectResponse(agent, response);
            })
                .catch((error) => {
                agent.onErrorThrownByInterceptor({
                    error: error,
                    method: "generateContent.<promise>",
                    module: "@google/genai",
                });
            });
        }
    }
    inspectResponse(agent, response) {
        if (!this.isResponse(response)) {
            return;
        }
        const usage = this.getUsageMetadata(response.usageMetadata);
        if (!usage) {
            return;
        }
        const aiStats = agent.getAIStatistics();
        aiStats.onAICall({
            provider: "google",
            model: response.modelVersion,
            inputTokens: usage.inputTokens,
            outputTokens: usage.outputTokens,
        });
    }
    wrap(hooks) {
        hooks
            .addPackage("@google/genai")
            .withVersion("^1.6.0")
            .onRequire((exports, pkgInfo) => {
            (0, wrapNewInstance_1.wrapNewInstance)(exports, "GoogleGenAI", pkgInfo, (instance) => {
                (0, wrapExport_1.wrapExport)(instance.models, "generateContent", pkgInfo, {
                    kind: undefined,
                    modifyReturnValue: (args, returnValue, agent) => {
                        this.inspectReturnValue(agent, returnValue);
                        return returnValue;
                    },
                });
            });
            return exports;
        });
    }
}
exports.GoogleGenAi = GoogleGenAi;