@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.
215 lines (214 loc) • 8.53 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.AiSDK = void 0;
const wrapExport_1 = require("../agent/hooks/wrapExport");
const isPlainObject_1 = require("../helpers/isPlainObject");
class AiSDK {
inspectAiCall(agent, args, response) {
if (!this.isResult(response)) {
return;
}
const provider = this.getProviderFromArgs(args);
if (!provider) {
return;
}
const modelName = this.getModelName(response);
const usage = this.getUsage(response.usage);
if (!usage) {
return;
}
const aiStats = agent.getAIStatistics();
aiStats.onAICall({
provider: provider,
model: modelName,
inputTokens: usage.inputTokens,
outputTokens: usage.outputTokens,
});
}
isResult(result) {
if (result &&
typeof result === "object" && // It is not a plain object
!Array.isArray(result) &&
"usage" in result &&
(0, isPlainObject_1.isPlainObject)(result.usage) &&
"response" in result &&
result.response &&
(0, isPlainObject_1.isPlainObject)(result.response) &&
typeof result.response.modelId === "string") {
return true;
}
return false;
}
getProviderFromArgs(args) {
if (!Array.isArray(args) || args.length === 0) {
return undefined;
}
const firstArg = args[0];
if (!(0, isPlainObject_1.isPlainObject)(firstArg)) {
return undefined;
}
if (!firstArg.model || typeof firstArg.model !== "object") {
return undefined;
}
if (!("provider" in firstArg.model) ||
typeof firstArg.model.provider !== "string") {
return undefined;
}
let providerName = firstArg.model.provider;
if (providerName.includes(".")) {
// e.g. google.generativeai
providerName = providerName.split(".")[0];
}
if (providerName === "amazon-bedrock") {
return "bedrock"; // Normalize amazon-bedrock to bedrock
}
if (providerName.includes("-")) {
// e.g. azure-openai
providerName = providerName.split("-")[0];
}
if (providerName === "google") {
return "gemini"; // Normalize google to gemini
}
return providerName;
}
getModelName(response) {
let modelName = response.response.modelId;
if (modelName.startsWith("models/")) {
modelName = modelName.slice(7); // Remove "models/" prefix
}
return modelName;
}
getUsage(usage) {
var _a, _b, _c, _d, _e;
const inputTokens = ((_a = usage.inputTokens) !== null && _a !== void 0 ? _a : 0) +
((_b = usage.promptTokens) !== null && _b !== void 0 ? _b : 0) +
((_c = usage.reasoningTokens) !== null && _c !== void 0 ? _c : 0);
const outputTokens = ((_d = usage.outputTokens) !== null && _d !== void 0 ? _d : 0) + ((_e = usage.completionTokens) !== null && _e !== void 0 ? _e : 0);
if (inputTokens === 0 && outputTokens === 0) {
return undefined;
}
return {
inputTokens,
outputTokens,
};
}
getInterceptors(methodName) {
return {
kind: "ai_op",
modifyReturnValue: (args, returnValue, agent) => {
if (returnValue instanceof Promise) {
// Inspect the response after the promise resolves, it won't change the original promise
returnValue
.then((response) => {
this.inspectAiCall(agent, args, response);
})
.catch((error) => {
agent.onErrorThrownByInterceptor({
error: error,
method: `${methodName}.<promise>`,
module: "ai",
});
});
}
else {
try {
this.inspectAiCall(agent, args, returnValue);
}
catch (error) {
agent.onErrorThrownByInterceptor({
error: error instanceof Error ? error : new Error(String(error)),
method: methodName,
module: "ai",
});
}
}
return returnValue;
},
};
}
getStreamInterceptors(methodName) {
return {
kind: "ai_op",
modifyReturnValue: (args, returnValue, agent) => {
if (!returnValue ||
typeof returnValue !== "object" ||
!("response" in returnValue) ||
!(returnValue.response instanceof Promise) ||
!("usage" in returnValue) ||
!(returnValue.usage instanceof Promise)) {
return returnValue;
}
Promise.allSettled([returnValue.response, returnValue.usage])
.then((promiseResults) => {
const response = promiseResults[0].status === "fulfilled"
? promiseResults[0].value
: undefined;
const usage = promiseResults[1].status === "fulfilled"
? promiseResults[1].value
: undefined;
if (!response || !usage) {
return;
}
this.inspectAiCall(agent, args, {
response,
usage,
});
})
.catch((error) => {
agent.onErrorThrownByInterceptor({
error: error,
method: `${methodName}.<promise>`,
module: "ai",
});
});
return returnValue;
},
};
}
wrap(hooks) {
hooks
.addPackage("ai")
.withVersion("^5.0.0 || ^4.0.0")
.onRequire((exports, pkgInfo) => {
// Can't wrap it directly because it's a readonly proxy
const generateTextFunc = exports.generateText;
const generateObjectFunc = exports.generateObject;
const streamTextFunc = exports.streamText;
const streamObjectFunc = exports.streamObject;
return {
...exports,
generateText: (0, wrapExport_1.wrapExport)(generateTextFunc, undefined, pkgInfo, this.getInterceptors("generateText")),
generateObject: (0, wrapExport_1.wrapExport)(generateObjectFunc, undefined, pkgInfo, this.getInterceptors("generateObject")),
streamText: (0, wrapExport_1.wrapExport)(streamTextFunc, undefined, pkgInfo, this.getStreamInterceptors("streamText")),
streamObject: (0, wrapExport_1.wrapExport)(streamObjectFunc, undefined, pkgInfo, this.getStreamInterceptors("streamObject")),
};
})
.addMultiFileInstrumentation(["dist/index.js", "dist/index.mjs"], [
{
name: "generateText",
nodeType: "FunctionDeclaration",
operationKind: "ai_op",
modifyReturnValue: this.getInterceptors("generateText").modifyReturnValue,
},
{
name: "generateObject",
nodeType: "FunctionDeclaration",
operationKind: "ai_op",
modifyReturnValue: this.getInterceptors("generateObject").modifyReturnValue,
},
{
name: "streamText",
nodeType: "FunctionDeclaration",
operationKind: "ai_op",
modifyReturnValue: this.getStreamInterceptors("streamText").modifyReturnValue,
},
{
name: "streamObject",
nodeType: "FunctionDeclaration",
operationKind: "ai_op",
modifyReturnValue: this.getStreamInterceptors("streamObject").modifyReturnValue,
},
]);
}
}
exports.AiSDK = AiSDK;