@aikidosec/firewall
Version:
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.
187 lines (186 loc) • 7.28 kB
JavaScript
;
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",
});
});
}
else {
try {
// The new instrumentation wraps the return statements inside the async function
// Inside an async function you can just return normal values (it's not a promise yet)
this.inspectResponse(agent, returnValue);
}
catch (error) {
agent.onErrorThrownByInterceptor({
error: error instanceof Error ? error : new Error("Unknown error"),
method: "generateContent",
module: "@google/genai",
});
}
}
}
isInteractionResponse(response) {
return (!!response &&
typeof response === "object" &&
!Array.isArray(response) &&
"id" in response &&
"status" in response &&
"created" in response);
}
inspectInteractionResponse(agent, response) {
var _a, _b, _c;
if (!this.isInteractionResponse(response)) {
return;
}
const usage = response.usage;
if (!usage) {
return;
}
const inputTokens = (_a = usage.total_input_tokens) !== null && _a !== void 0 ? _a : 0;
const outputTokens = (_b = usage.total_output_tokens) !== null && _b !== void 0 ? _b : 0;
if (inputTokens === 0 && outputTokens === 0) {
return;
}
const aiStats = agent.getAIStatistics();
aiStats.onAICall({
provider: "google",
model: (_c = response.model) !== null && _c !== void 0 ? _c : "unknown",
inputTokens,
outputTokens,
});
}
inspectInteractionReturnValue(agent, returnValue) {
if (returnValue instanceof Promise) {
returnValue
.then((response) => {
this.inspectInteractionResponse(agent, response);
})
.catch((error) => {
agent.onErrorThrownByInterceptor({
error: error,
method: "interactions.create.<promise>",
module: "@google/genai",
});
});
}
else {
try {
this.inspectInteractionResponse(agent, returnValue);
}
catch (error) {
agent.onErrorThrownByInterceptor({
error: error instanceof Error ? error : new Error("Unknown error"),
method: "interactions.create",
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 || ^2.0.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;
},
});
// Access getter once
const interactions = instance.interactions;
if (interactions && typeof interactions.create === "function") {
(0, wrapExport_1.wrapExport)(interactions, "create", pkgInfo, {
kind: undefined,
modifyReturnValue: (_args, returnValue, agent) => {
this.inspectInteractionReturnValue(agent, returnValue);
return returnValue;
},
});
}
});
return exports;
})
.addMultiFileInstrumentation(["dist/node/index.cjs", "dist/node/index.mjs"], [
{
name: "this.generateContent",
nodeType: "FunctionAssignment",
operationKind: "ai_op",
modifyReturnValue: (_args, returnValue, agent) => {
this.inspectReturnValue(agent, returnValue);
return returnValue;
},
},
{
name: "create",
nodeType: "MethodDefinition",
className: "BaseInteractions",
operationKind: "ai_op",
modifyReturnValue: (_args, returnValue, agent) => {
this.inspectInteractionReturnValue(agent, returnValue);
return returnValue;
},
},
]);
}
}
exports.GoogleGenAi = GoogleGenAi;