@aikidosec/firewall/sinks/Mistral.js

Zen by Aikido is an embedded Web Application Firewall that autonomously protects Node.js apps against common and critical attacks

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.Mistral = void 0;
const wrapExport_1 = require("../agent/hooks/wrapExport");
const isPlainObject_1 = require("../helpers/isPlainObject");
function isMistralChatCompletionResponse(response) {
    return ((0, isPlainObject_1.isPlainObject)(response) &&
        "model" in response &&
        typeof response.model === "string");
}
class Mistral {
    inspectResponse(agent, response) {
        var _a;
        if (!isMistralChatCompletionResponse(response)) {
            return;
        }
        let inputTokens = 0;
        let outputTokens = 0;
        if (response.usage) {
            if (typeof response.usage.promptTokens === "number") {
                inputTokens = response.usage.promptTokens;
            }
            if (typeof response.usage.completionTokens === "number") {
                outputTokens = response.usage.completionTokens;
            }
        }
        const aiStats = agent.getAIStatistics();
        aiStats.onAICall({
            provider: "mistral",
            model: (_a = response.model) !== null && _a !== void 0 ? _a : "",
            inputTokens: inputTokens,
            outputTokens: outputTokens,
        });
    }
    wrap(hooks) {
        hooks
            .addPackage("@mistralai/mistralai")
            .withVersion("^1.0.0")
            .onRequire((exports, pkgInfo) => {
            if (exports.Mistral &&
                exports.Mistral.prototype &&
                exports.Mistral.prototype.chat) {
                (0, wrapExport_1.wrapExport)(exports.Mistral.prototype.chat, "complete", pkgInfo, {
                    kind: "ai_op",
                    modifyReturnValue: (_, returnValue, agent) => {
                        if (returnValue instanceof Promise) {
                            // Inspect the response after the promise resolves
                            returnValue.then((response) => {
                                try {
                                    this.inspectResponse(agent, response);
                                }
                                catch {
                                    // If we don't catch these errors, it will result in an unhandled promise rejection!
                                }
                            });
                        }
                        return returnValue;
                    },
                });
            }
        });
    }
}
exports.Mistral = Mistral;