@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.
151 lines (150 loc) • 6.08 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) => {
this.inspectResponse(agent, response, this.getProvider(exports, subject));
})
.catch((error) => {
agent.onErrorThrownByInterceptor({
error: error,
method: "create.<promise>",
module: "openai",
});
});
}
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) => {
this.inspectCompletionResponse(agent, response, this.getProvider(exports, subject));
})
.catch((error) => {
agent.onErrorThrownByInterceptor({
error: error,
method: "create.<promise>",
module: "openai",
});
});
}
return returnValue;
},
});
}
});
}
}
exports.OpenAI = OpenAI;