@convo-lang/convo-lang-aws-cdk/src/lib/price-capping.js

The language of AI

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.storeTokenUsageAsync = exports.checkTokenQuotaAsync = exports.getTokenQuotaAsync = void 0;
const aws_dynamo_1 = require("@iyio/aws-dynamo");
const convo_lang_aws_cdk_deps_1 = require("./convo-lang-aws-cdk.deps");
const CapsTable = {
    name: "Caps",
    primaryKey: "id",
    tableIdParam: convo_lang_aws_cdk_deps_1.convoLangCapsTableParam,
};
const _getCapAsync = async (id) => {
    let cap = await (0, aws_dynamo_1.dynamoClient)().getFromTableAsync(CapsTable, { id });
    if (cap) {
        return cap;
    }
    cap = { id, usage: 0 };
    await (0, aws_dynamo_1.dynamoClient)().putIntoTable(CapsTable, cap);
    return cap;
};
const getTokenQuotaAsync = (caller) => {
    const id = caller.remoteAddress || '0.0.0.0';
    return _getCapAsync(id);
};
exports.getTokenQuotaAsync = getTokenQuotaAsync;
const checkTokenQuotaAsync = async (caller) => {
    const id = caller.remoteAddress || '0.0.0.0';
    let singleCap = convo_lang_aws_cdk_deps_1.convoLangAnonUsdCapParam.get();
    let totalCap = convo_lang_aws_cdk_deps_1.convoLangAnonUsdCapTotalParam.get();
    if (singleCap === undefined && totalCap === undefined) {
        return { allow: true };
    }
    const [single, total] = await Promise.all([
        singleCap !== undefined ? _getCapAsync(id) : undefined,
        totalCap !== undefined ? _getCapAsync('_') : undefined,
    ]);
    if (single?.cap !== undefined) {
        singleCap = single.cap;
    }
    if (total?.cap !== undefined) {
        totalCap = total.cap;
    }
    const allow = ((singleCap === undefined ? true : single ? single.usage < singleCap : false) &&
        (totalCap === undefined ? true : total ? total.usage < totalCap : false));
    return {
        allow,
        singleUsage: (singleCap === undefined || !single) ? undefined : singleCap - single.usage,
        totalUsage: (totalCap === undefined || !total) ? undefined : totalCap - total.usage,
        singleCap,
        totalCap
    };
};
exports.checkTokenQuotaAsync = checkTokenQuotaAsync;
const storeTokenUsageAsync = async (usage, caller) => {
    const id = caller.remoteAddress || '0.0.0.0';
    const singleCap = convo_lang_aws_cdk_deps_1.convoLangAnonUsdCapParam.get();
    const totalCap = convo_lang_aws_cdk_deps_1.convoLangAnonUsdCapTotalParam.get();
    if (singleCap === undefined && totalCap === undefined) {
        return;
    }
    await Promise.all([
        singleCap !== undefined ? (0, aws_dynamo_1.dynamoClient)().patchTableItem(CapsTable, {
            id,
            usage: (0, aws_dynamo_1.createUpdateExpression)({
                add: usage
            })
        }) : undefined,
        totalCap !== undefined ? (0, aws_dynamo_1.dynamoClient)().patchTableItem(CapsTable, {
            id: '_',
            usage: (0, aws_dynamo_1.createUpdateExpression)({
                add: usage
            })
        }) : undefined,
    ]);
};
exports.storeTokenUsageAsync = storeTokenUsageAsync;
//# sourceMappingURL=price-capping.js.map