@aikidosec/firewall/sinks/Mistral.js

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.

"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) => {
                                this.inspectResponse(agent, response);
                            })
                                .catch((error) => {
                                agent.onErrorThrownByInterceptor({
                                    error: error,
                                    method: "complete.<promise>",
                                    module: "@mistralai/mistralai",
                                });
                            });
                        }
                        return returnValue;
                    },
                });
            }
        });
    }
}
exports.Mistral = Mistral;