@aikidosec/firewall/sinks/Anthropic.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.Anthropic = void 0;
const wrapExport_1 = require("../agent/hooks/wrapExport");
const isPlainObject_1 = require("../helpers/isPlainObject");
function isAnthropicMessageResponse(obj) {
    if (!(0, isPlainObject_1.isPlainObject)(obj)) {
        return false;
    }
    if (typeof obj.model !== "string") {
        return false;
    }
    if (obj.usage !== undefined) {
        if (!(0, isPlainObject_1.isPlainObject)(obj.usage)) {
            return false;
        }
        if (typeof obj.usage.input_tokens !== "number" ||
            typeof obj.usage.output_tokens !== "number") {
            return false;
        }
    }
    return true;
}
class Anthropic {
    inspectResponse(agent, response) {
        var _a;
        if (!isAnthropicMessageResponse(response)) {
            return;
        }
        let inputTokens = 0;
        let outputTokens = 0;
        if (response.usage) {
            inputTokens = response.usage.input_tokens;
            outputTokens = response.usage.output_tokens;
        }
        const aiStats = agent.getAIStatistics();
        aiStats.onAICall({
            provider: "anthropic",
            model: (_a = response.model) !== null && _a !== void 0 ? _a : "",
            inputTokens: inputTokens,
            outputTokens: outputTokens,
        });
    }
    wrap(hooks) {
        hooks
            .addPackage("@anthropic-ai/sdk")
            .withVersion("^0.40.0")
            .onFileRequire("resources/messages/messages.js", (exports, pkgInfo) => {
            if (exports.Messages &&
                exports.Messages.prototype &&
                exports.Messages.prototype) {
                (0, wrapExport_1.wrapExport)(exports.Messages.prototype, "create", 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.Anthropic = Anthropic;