@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.
119 lines (118 loc) • 4.73 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.AwsSDKVersion3 = void 0;
const wrapExport_1 = require("../agent/hooks/wrapExport");
const isPlainObject_1 = require("../helpers/isPlainObject");
function isUsage(usage) {
return ((0, isPlainObject_1.isPlainObject)(usage) &&
typeof usage.input_tokens === "number" &&
typeof usage.output_tokens === "number");
}
function isInvokeResponse(response) {
return (0, isPlainObject_1.isPlainObject)(response);
}
function isConverseUsage(usage) {
return ((0, isPlainObject_1.isPlainObject)(usage) &&
typeof usage.inputTokens === "number" &&
typeof usage.outputTokens === "number");
}
function isConverseResponse(response) {
return (0, isPlainObject_1.isPlainObject)(response);
}
class AwsSDKVersion3 {
processInvokeModelResponse(response, agent) {
if (!isInvokeResponse(response)) {
return;
}
let body;
try {
const json = new TextDecoder("utf-8", {
fatal: true,
}).decode(response.body);
body = JSON.parse(json);
}
catch {
// Ignore errors when parsing the response body
return;
}
if (typeof body.model === "string") {
let inputTokens = 0;
let outputTokens = 0;
if (isUsage(body.usage)) {
inputTokens = body.usage.input_tokens;
outputTokens = body.usage.output_tokens;
}
const aiStats = agent.getAIStatistics();
aiStats.onAICall({
provider: "bedrock",
model: body.model,
inputTokens: inputTokens,
outputTokens: outputTokens,
});
}
}
processConverseResponse(response, command, agent) {
// @ts-expect-error We don't know the type of command
if (!command || !command.input || !command.input.modelId) {
return;
}
if (!isConverseResponse(response)) {
return;
}
// @ts-expect-error We don't know the type of command
const modelId = command.input.modelId;
let inputTokens = 0;
let outputTokens = 0;
if (isConverseUsage(response.usage)) {
inputTokens = response.usage.inputTokens || 0;
outputTokens = response.usage.outputTokens || 0;
}
const aiStats = agent.getAIStatistics();
aiStats.onAICall({
provider: "bedrock",
model: modelId,
inputTokens: inputTokens,
outputTokens: outputTokens,
});
}
wrap(hooks) {
hooks
.addPackage("@aws-sdk/client-bedrock-runtime")
.withVersion("^3.0.0")
.onRequire((exports, pkgInfo) => {
if (exports.BedrockRuntimeClient) {
(0, wrapExport_1.wrapExport)(exports.BedrockRuntimeClient.prototype, "send", pkgInfo, {
kind: "ai_op",
modifyReturnValue: (args, returnValue, agent) => {
if (args.length > 0) {
const command = args[0];
if (returnValue instanceof Promise) {
// Inspect the response after the promise resolves, it won't change the original promise
returnValue
.then((response) => {
if (exports.InvokeModelCommand &&
command instanceof exports.InvokeModelCommand) {
this.processInvokeModelResponse(response, agent);
}
else if (exports.ConverseCommand &&
command instanceof exports.ConverseCommand) {
this.processConverseResponse(response, command, agent);
}
})
.catch((error) => {
agent.onErrorThrownByInterceptor({
error: error,
method: "send.<promise>",
module: "@aws-sdk/client-bedrock-runtime",
});
});
}
}
return returnValue;
},
});
}
});
}
}
exports.AwsSDKVersion3 = AwsSDKVersion3;