saas: new tier policy

src/api/crawler.ts

@@ -240,6 +240,7 @@ export class CrawlerHost extends RPCHost {
         const uid = await auth.solveUID();
         let chargeAmount = 0;
         const crawlerOptions = ctx.method === 'GET' ? crawlerOptionsHeaderOnly : crawlerOptionsParamsAllowed;
+        const tierPolicy = await this.saasAssertTierPolicy(crawlerOptions, auth);
 
         // Use koa ctx.URL, a standard URL object to avoid node.js framework prop naming confusion
         const targetUrl = await this.getTargetUrl(tryDecodeURIComponent(`${ctx.URL.pathname}${ctx.URL.search}`), crawlerOptions);
@@ -298,15 +299,13 @@ export class CrawlerHost extends RPCHost {
                 if (crawlerOptions.tokenBudget && chargeAmount > crawlerOptions.tokenBudget) {
                     return;
                 }
-                if (chargeAmount) {
-                    apiRoll._ref?.set({
-                        chargeAmount,
-                    }, { merge: true }).catch((err) => this.logger.warn(`Failed to log charge amount in apiRoll`, { err }));
-                }
+                apiRoll.chargeAmount = chargeAmount;
             });
         }
 
         if (!uid) {
+            // Enforce no proxy is allocated for anonymous users due to abuse.
+            crawlerOptions.proxy = 'none';
             const blockade = (await DomainBlockade.fromFirestoreQuery(
                 DomainBlockade.COLLECTION
                     .where('domain', '==', targetUrl.hostname.toLowerCase())
@@ -338,10 +337,7 @@ export class CrawlerHost extends RPCHost {
                     }
 
                     const formatted = await this.formatSnapshot(crawlerOptions, scrapped, targetUrl, this.urlValidMs, crawlOpts);
-                    chargeAmount = this.assignChargeAmount(formatted, crawlerOptions);
-                    if (crawlerOptions.tokenBudget && chargeAmount > crawlerOptions.tokenBudget) {
-                        throw new BudgetExceededError(`Token budget (${crawlerOptions.tokenBudget}) exceeded, intended charge amount ${chargeAmount}.`);
-                    }
+                    chargeAmount = this.assignChargeAmount(formatted, tierPolicy);
                     sseStream.write({
                         event: 'data',
                         data: formatted,
@@ -379,11 +375,7 @@ export class CrawlerHost extends RPCHost {
                     }
 
                     const formatted = await this.formatSnapshot(crawlerOptions, scrapped, targetUrl, this.urlValidMs, crawlOpts);
-                    chargeAmount = this.assignChargeAmount(formatted, crawlerOptions);
-
-                    if (crawlerOptions.tokenBudget && chargeAmount > crawlerOptions.tokenBudget) {
-                        throw new BudgetExceededError(`Token budget (${crawlerOptions.tokenBudget}) exceeded, intended charge amount ${chargeAmount}.`);
-                    }
+                    chargeAmount = this.assignChargeAmount(formatted, tierPolicy);
 
                     if (scrapped?.pdfs?.length && !chargeAmount) {
                         continue;
@@ -405,10 +397,7 @@ export class CrawlerHost extends RPCHost {
             }
 
             const formatted = await this.formatSnapshot(crawlerOptions, lastScrapped, targetUrl, this.urlValidMs, crawlOpts);
-            chargeAmount = this.assignChargeAmount(formatted, crawlerOptions);
-            if (crawlerOptions.tokenBudget && chargeAmount > crawlerOptions.tokenBudget) {
-                throw new BudgetExceededError(`Token budget (${crawlerOptions.tokenBudget}) exceeded, intended charge amount ${chargeAmount}.`);
-            }
+            chargeAmount = this.assignChargeAmount(formatted, tierPolicy);
 
             return formatted;
         }
@@ -434,10 +423,7 @@ export class CrawlerHost extends RPCHost {
                 }
 
                 const formatted = await this.formatSnapshot(crawlerOptions, scrapped, targetUrl, this.urlValidMs, crawlOpts);
-                chargeAmount = this.assignChargeAmount(formatted, crawlerOptions);
-                if (crawlerOptions.tokenBudget && chargeAmount > crawlerOptions.tokenBudget) {
-                    throw new BudgetExceededError(`Token budget (${crawlerOptions.tokenBudget}) exceeded, intended charge amount ${chargeAmount}.`);
-                }
+                chargeAmount = this.assignChargeAmount(formatted, tierPolicy);
 
                 if (crawlerOptions.respondWith === 'screenshot' && Reflect.get(formatted, 'screenshotUrl')) {
                     return assignTransferProtocolMeta(`${formatted.textRepresentation}`,
@@ -465,10 +451,7 @@ export class CrawlerHost extends RPCHost {
             throw new AssertionFailureError(`No content available for URL ${targetUrl}`);
         }
         const formatted = await this.formatSnapshot(crawlerOptions, lastScrapped, targetUrl, this.urlValidMs, crawlOpts);
-        chargeAmount = this.assignChargeAmount(formatted, crawlerOptions);
-        if (crawlerOptions.tokenBudget && chargeAmount > crawlerOptions.tokenBudget) {
-            throw new BudgetExceededError(`Token budget (${crawlerOptions.tokenBudget}) exceeded, intended charge amount ${chargeAmount}.`);
-        }
+        chargeAmount = this.assignChargeAmount(formatted, tierPolicy);
 
         if (crawlerOptions.respondWith === 'screenshot' && Reflect.get(formatted, 'screenshotUrl')) {
 
@@ -840,7 +823,8 @@ export class CrawlerHost extends RPCHost {
                     yield this.jsdomControl.narrowSnapshot(draftSnapshot, crawlOpts);
                 }
                 let fallbackProxyIsUsed = false;
-                if (((!crawlOpts?.allocProxy || crawlOpts.allocProxy === 'none') && !crawlOpts?.proxyUrl) &&
+                if (
+                    ((!crawlOpts?.allocProxy || crawlOpts.allocProxy !== 'none') && !crawlOpts?.proxyUrl) &&
                     (analyzed.tokens < 42 || sideLoaded.status !== 200)
                 ) {
                     const proxyLoaded = await this.sideLoadWithAllocatedProxy(urlToCrawl, altOpts);
@@ -911,18 +895,14 @@ export class CrawlerHost extends RPCHost {
         }
     }
 
-    assignChargeAmount(formatted: FormattedPage, crawlerOptions?: CrawlerOptions) {
+    assignChargeAmount(formatted: FormattedPage, saasTierPolicy?: Parameters<typeof this.saasApplyTierPolicy>[0]) {
         if (!formatted) {
             return 0;
         }
 
         let amount = 0;
         if (formatted.content) {
-            const x1 = estimateToken(formatted.content);
-            if (crawlerOptions?.respondWith?.toLowerCase().includes('lm')) {
-                amount += x1 * 2;
-            }
-            amount += x1;
+            amount = estimateToken(formatted.content);
         } else if (formatted.description) {
             amount += estimateToken(formatted.description);
         }
@@ -939,6 +919,10 @@ export class CrawlerHost extends RPCHost {
             amount += 765;
         }
 
+        if (saasTierPolicy) {
+            amount = this.saasApplyTierPolicy(saasTierPolicy, amount);
+        }
+
         Object.assign(formatted, { usage: { tokens: amount } });
         assignMeta(formatted, { usage: { tokens: amount } });
 
@@ -1312,4 +1296,54 @@ export class CrawlerHost extends RPCHost {
 
         return false;
     }
+
+    async saasAssertTierPolicy(opts: CrawlerOptions, auth: JinaEmbeddingsAuthDTO) {
+        let chargeScalar = 1;
+        let minimalCharge = 0;
+
+        if (opts.injectPageScript || opts.injectFrameScript) {
+            await auth.assertTier(0, 'Script injection');
+            minimalCharge = 4_000;
+        }
+
+        if (opts.withGeneratedAlt) {
+            await auth.assertTier(0, 'Alt text generation');
+            minimalCharge = 4_000;
+        }
+
+        if (opts.withIframe) {
+            await auth.assertTier(0, 'Iframe');
+        }
+
+        if (opts.engine === ENGINE_TYPE.CF_BROWSER_RENDERING) {
+            await auth.assertTier(0, 'Cloudflare browser rendering');
+            minimalCharge = 4_000;
+        }
+
+        if (opts.respondWith.includes('lm') || opts.engine?.includes('lm')) {
+            await auth.assertTier(0, 'Language model');
+            minimalCharge = 4_000;
+            chargeScalar = 3;
+        }
+
+        if (opts.proxy && opts.proxy !== 'none') {
+            await auth.assertTier(['auto', 'any'].includes(opts.proxy) ? 0 : 2, 'Proxy allocation');
+            chargeScalar = 5;
+        }
+
+        return {
+            budget: opts.tokenBudget || 0,
+            chargeScalar,
+            minimalCharge,
+        };
+    }
+
+    saasApplyTierPolicy(policy: Awaited<ReturnType<typeof this.saasAssertTierPolicy>>, chargeAmount: number) {
+        const effectiveChargeAmount = policy.chargeScalar * Math.max(chargeAmount, policy.minimalCharge);
+        if (policy.budget && policy.budget < effectiveChargeAmount) {
+            throw new BudgetExceededError(`Token budget (${policy.budget}) exceeded, intended charge amount ${effectiveChargeAmount}`);
+        }
+
+        return effectiveChargeAmount;
+    }
 }

src/dto/jina-embeddings-auth.ts

@@ -17,6 +17,7 @@ import { AsyncLocalContext } from '../services/async-context';
 import envConfig from '../shared/services/secrets';
 import { JinaEmbeddingsDashboardHTTP } from '../shared/3rd-party/jina-embeddings';
 import { JinaEmbeddingsTokenAccount } from '../shared/db/jina-embeddings-token-account';
+import { TierFeatureConstraintError } from '../services/errors';
 
 const authDtoLogger = logger.child({ service: 'JinaAuthDTO' });
 
@@ -236,6 +237,30 @@ export class JinaEmbeddingsAuthDTO extends AutoCastable {
         return this.user!;
     }
 
+    async assertTier(n: number, feature?: string) {
+        let user;
+        try {
+            user = await this.assertUser();
+        } catch (err) {
+            if (err instanceof AuthenticationRequiredError) {
+                throw new AuthenticationRequiredError({
+                    message: `Authentication is required to use this feature${feature ? ` (${feature})` : ''}. Please provide a valid API key.`
+                });
+            }
+
+            throw err;
+        }
+
+        const tier = parseInt(user.metadata?.speed_level);
+        if (isNaN(tier) || tier < n) {
+            throw new TierFeatureConstraintError({
+                message: `Your current plan does not support this feature${feature ? ` (${feature})` : ''}. Please upgrade your plan.`
+            });
+        }
+
+        return true;
+    }
+
     getRateLimits(...tags: string[]) {
         const descs = tags.map((x) => this.user?.customRateLimits?.[x] || []).flat().filter((x) => x.isEffective());
 

src/services/errors.ts

@@ -27,7 +27,7 @@ export class EmailUnverifiedError extends ApplicationError { }
 export class InsufficientCreditsError extends ApplicationError { }
 
 @StatusCode(40202)
-export class FreeFeatureLimitError extends ApplicationError { }
+export class TierFeatureConstraintError extends ApplicationError { }
 
 @StatusCode(40203)
 export class InsufficientBalanceError extends ApplicationError { }