saas: new tier policy

This commit is contained in:
Yanlong Wang 2025-04-18 15:32:43 +08:00
parent 33ca16405e
commit 5a81177c56
No known key found for this signature in database
GPG Key ID: C0A623C0BADF9F37
3 changed files with 93 additions and 34 deletions

View File

@ -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;
}
}

View File

@ -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());

View File

@ -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 { }