@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.
145 lines (144 loc) • 5.85 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.OpenAI = void 0;
const wrapExport_1 = require("../agent/hooks/wrapExport");
const isPlainObject_1 = require("../helpers/isPlainObject");
function isResponse(response) {
return ((0, isPlainObject_1.isPlainObject)(response) &&
"model" in response &&
typeof response.model === "string");
}
function isCompletionResponse(response) {
return ((0, isPlainObject_1.isPlainObject)(response) &&
"model" in response &&
typeof response.model === "string");
}
class OpenAI {
inspectResponse(agent, response, provider) {
var _a;
if (!isResponse(response)) {
return;
}
let inputTokens = 0;
let outputTokens = 0;
if (response.usage) {
if (typeof response.usage.input_tokens === "number") {
inputTokens = response.usage.input_tokens;
}
if (typeof response.usage.output_tokens === "number") {
outputTokens = response.usage.output_tokens;
}
}
const aiStats = agent.getAIStatistics();
aiStats.onAICall({
provider: provider,
model: (_a = response.model) !== null && _a !== void 0 ? _a : "",
inputTokens: inputTokens,
outputTokens: outputTokens,
});
}
inspectCompletionResponse(agent, response, provider) {
var _a;
if (!isCompletionResponse(response)) {
return;
}
let inputTokens = 0;
let outputTokens = 0;
if (response.usage) {
if (typeof response.usage.prompt_tokens === "number") {
inputTokens = response.usage.prompt_tokens;
}
if (typeof response.usage.completion_tokens === "number") {
outputTokens = response.usage.completion_tokens;
}
}
const aiStats = agent.getAIStatistics();
aiStats.onAICall({
provider: provider,
model: (_a = response.model) !== null && _a !== void 0 ? _a : "",
inputTokens: inputTokens,
outputTokens: outputTokens,
});
}
// The _client property is used to determine if the OpenAI client is an Azure OpenAI client or not.
// See https://github.com/openai/openai-node/blob/master/src/core/resource.ts
getProvider(exports, subject) {
if (
// @ts-expect-error We don't know the type of exports
exports.AzureOpenAI &&
// @ts-expect-error We don't know the type of subject
subject._client &&
// @ts-expect-error We don't know the type of subject
subject._client instanceof exports.AzureOpenAI) {
return "azure";
}
return "openai";
}
getResponsesClass(exports) {
if (exports.Responses) {
return exports.Responses; // v4
}
if (exports.OpenAI && exports.OpenAI.Responses) {
return exports.OpenAI.Responses; // v5
}
}
getCompletionsClass(exports) {
if (exports.Chat && exports.Chat.Completions) {
return exports.Chat.Completions; // v4
}
if (exports.OpenAI &&
exports.OpenAI.Chat &&
exports.OpenAI.Chat.Completions) {
return exports.OpenAI.Chat.Completions; // v5
}
}
wrap(hooks) {
// Note: Streaming is not supported yet
hooks
.addPackage("openai")
.withVersion("^5.0.0 || ^4.0.0")
.onRequire((exports, pkgInfo) => {
const responsesClass = this.getResponsesClass(exports);
if (responsesClass) {
(0, wrapExport_1.wrapExport)(responsesClass.prototype, "create", pkgInfo, {
kind: "ai_op",
modifyReturnValue: (_, returnValue, agent, subject) => {
if (returnValue instanceof Promise) {
// Inspect the response after the promise resolves, it won't change the original promise
returnValue.then((response) => {
try {
this.inspectResponse(agent, response, this.getProvider(exports, subject));
}
catch {
// If we don't catch these errors, it will result in an unhandled promise rejection!
}
});
}
return returnValue;
},
});
}
const completionsClass = this.getCompletionsClass(exports);
if (completionsClass) {
(0, wrapExport_1.wrapExport)(completionsClass.prototype, "create", pkgInfo, {
kind: "ai_op",
modifyReturnValue: (_, returnValue, agent, subject) => {
if (returnValue instanceof Promise) {
// Inspect the response after the promise resolves, it won't change the original promise
returnValue.then((response) => {
try {
this.inspectCompletionResponse(agent, response, this.getProvider(exports, subject));
}
catch {
// If we don't catch these errors, it will result in an unhandled promise rejection!
}
});
}
return returnValue;
},
});
}
});
}
}
exports.OpenAI = OpenAI;