@aikidosec/firewall
Version:
Zen by Aikido is an embedded Web Application Firewall that autonomously protects Node.js apps against common and critical attacks
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;