@aikidosec/firewall
Version:
Zen by Aikido is an embedded Web Application Firewall that autonomously protects Node.js apps against common and critical attacks
116 lines (115 loc) • 4.61 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) => {
try {
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 {
// If we don't catch these errors, it will result in an unhandled promise rejection!
}
});
}
}
return returnValue;
},
});
}
});
}
}
exports.AwsSDKVersion3 = AwsSDKVersion3;