@aikidosec/firewall/agent/AIStatistics.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.AIStatistics = void 0;
class AIStatistics {
    constructor() {
        this.calls = new Map();
    }
    getProviderKey(provider, model) {
        return `${provider}:${model}`;
    }
    getRouteKey(path, method) {
        return `${method}:${path}`;
    }
    ensureProviderStats(provider, model) {
        const key = this.getProviderKey(provider, model);
        if (!this.calls.has(key)) {
            this.calls.set(key, {
                provider,
                model,
                calls: 0,
                tokens: {
                    input: 0,
                    output: 0,
                    total: 0,
                },
            });
        }
        return this.calls.get(key);
    }
    onAICall({ provider, model, inputTokens, outputTokens, }) {
        if (!provider || !model) {
            return;
        }
        const providerStats = this.ensureProviderStats(provider, model);
        providerStats.calls += 1;
        providerStats.tokens.input += inputTokens;
        providerStats.tokens.output += outputTokens;
        providerStats.tokens.total += inputTokens + outputTokens;
    }
    getStats() {
        return Array.from(this.calls.values()).map((stats) => {
            return {
                provider: stats.provider,
                model: stats.model,
                calls: stats.calls,
                tokens: {
                    input: stats.tokens.input,
                    output: stats.tokens.output,
                    total: stats.tokens.total,
                },
            };
        });
    }
    reset() {
        this.calls.clear();
    }
    isEmpty() {
        return this.calls.size === 0;
    }
}
exports.AIStatistics = AIStatistics;