0

Merge pull request #216 from mendableai/nsc/new-pricing

feat: New pricing/limits changes
This commit is contained in:
Nicolas 2024-05-30 15:36:59 -07:00 committed by GitHub
commit 0c115c6181
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 87 additions and 127 deletions

View File

@ -31,6 +31,13 @@ POSTHOG_HOST= # set if you'd like to send posthog events like job logs
STRIPE_PRICE_ID_STANDARD= STRIPE_PRICE_ID_STANDARD=
STRIPE_PRICE_ID_SCALE= STRIPE_PRICE_ID_SCALE=
STRIPE_PRICE_ID_STARTER=
STRIPE_PRICE_ID_HOBBY=
STRIPE_PRICE_ID_HOBBY_YEARLY=
STRIPE_PRICE_ID_STANDARD_NEW=
STRIPE_PRICE_ID_STANDARD_NEW_YEARLY=
STRIPE_PRICE_ID_GROWTH=
STRIPE_PRICE_ID_GROWTH_YEARLY=
HYPERDX_API_KEY= HYPERDX_API_KEY=
HDX_NODE_BETA_MODE=1 HDX_NODE_BETA_MODE=1

View File

@ -1004,7 +1004,7 @@ describe("E2E Tests for API Routes", () => {
describe("Rate Limiter", () => { describe("Rate Limiter", () => {
it("should return 429 when rate limit is exceeded for preview token", async () => { it("should return 429 when rate limit is exceeded for preview token", async () => {
for (let i = 0; i < 5; i++) { for (let i = 0; i < 4; i++) {
const response = await request(TEST_URL) const response = await request(TEST_URL)
.post("/v0/scrape") .post("/v0/scrape")
.set("Authorization", `Bearer this_is_just_a_preview_token`) .set("Authorization", `Bearer this_is_just_a_preview_token`)

View File

@ -29,6 +29,7 @@ export async function supaAuthenticateUser(
team_id?: string; team_id?: string;
error?: string; error?: string;
status?: number; status?: number;
plan?: string;
}> { }> {
const authHeader = req.headers.authorization; const authHeader = req.headers.authorization;
if (!authHeader) { if (!authHeader) {
@ -104,12 +105,13 @@ export async function supaAuthenticateUser(
case RateLimiterMode.Scrape: case RateLimiterMode.Scrape:
rateLimiter = getRateLimiter(RateLimiterMode.Scrape, token, subscriptionData.plan); rateLimiter = getRateLimiter(RateLimiterMode.Scrape, token, subscriptionData.plan);
break; break;
case RateLimiterMode.Search:
rateLimiter = getRateLimiter(RateLimiterMode.Search, token, subscriptionData.plan);
break;
case RateLimiterMode.CrawlStatus: case RateLimiterMode.CrawlStatus:
rateLimiter = getRateLimiter(RateLimiterMode.CrawlStatus, token); rateLimiter = getRateLimiter(RateLimiterMode.CrawlStatus, token);
break; break;
case RateLimiterMode.Search:
rateLimiter = getRateLimiter(RateLimiterMode.Search, token);
break;
case RateLimiterMode.Preview: case RateLimiterMode.Preview:
rateLimiter = getRateLimiter(RateLimiterMode.Preview, token); rateLimiter = getRateLimiter(RateLimiterMode.Preview, token);
break; break;
@ -172,16 +174,24 @@ export async function supaAuthenticateUser(
subscriptionData = data[0]; subscriptionData = data[0];
} }
return { success: true, team_id: subscriptionData.team_id }; return { success: true, team_id: subscriptionData.team_id, plan: subscriptionData.plan ?? ""};
} }
function getPlanByPriceId(price_id: string) { function getPlanByPriceId(price_id: string) {
switch (price_id) { switch (price_id) {
case process.env.STRIPE_PRICE_ID_STARTER:
return 'starter';
case process.env.STRIPE_PRICE_ID_STANDARD: case process.env.STRIPE_PRICE_ID_STANDARD:
return 'standard'; return 'standard';
case process.env.STRIPE_PRICE_ID_SCALE: case process.env.STRIPE_PRICE_ID_SCALE:
return 'scale'; return 'scale';
case process.env.STRIPE_PRICE_ID_HOBBY || process.env.STRIPE_PRICE_ID_HOBBY_YEARLY:
return 'hobby';
case process.env.STRIPE_PRICE_ID_STANDARD_NEW || process.env.STRIPE_PRICE_ID_STANDARD_NEW_YEARLY:
return 'standard-new';
case process.env.STRIPE_PRICE_ID_GROWTH || process.env.STRIPE_PRICE_ID_GROWTH_YEARLY:
return 'growth';
default: default:
return 'starter'; return 'free';
} }
} }

View File

@ -15,7 +15,8 @@ export async function scrapeHelper(
crawlerOptions: any, crawlerOptions: any,
pageOptions: PageOptions, pageOptions: PageOptions,
extractorOptions: ExtractorOptions, extractorOptions: ExtractorOptions,
timeout: number timeout: number,
plan?: string
): Promise<{ ): Promise<{
success: boolean; success: boolean;
error?: string; error?: string;
@ -64,7 +65,9 @@ export async function scrapeHelper(
} }
let creditsToBeBilled = filteredDocs.length; let creditsToBeBilled = filteredDocs.length;
const creditsPerLLMExtract = 5; const creditsPerLLMExtract = plan === "starter" ? 5 : 50;
if (extractorOptions.mode === "llm-extraction") { if (extractorOptions.mode === "llm-extraction") {
creditsToBeBilled = creditsToBeBilled + (creditsPerLLMExtract * filteredDocs.length); creditsToBeBilled = creditsToBeBilled + (creditsPerLLMExtract * filteredDocs.length);
@ -93,7 +96,7 @@ export async function scrapeHelper(
export async function scrapeController(req: Request, res: Response) { export async function scrapeController(req: Request, res: Response) {
try { try {
// make sure to authenticate user first, Bearer <token> // make sure to authenticate user first, Bearer <token>
const { success, team_id, error, status } = await authenticateUser( const { success, team_id, error, status, plan } = await authenticateUser(
req, req,
res, res,
RateLimiterMode.Scrape RateLimiterMode.Scrape
@ -129,7 +132,8 @@ export async function scrapeController(req: Request, res: Response) {
crawlerOptions, crawlerOptions,
pageOptions, pageOptions,
extractorOptions, extractorOptions,
timeout timeout,
plan
); );
const endTime = new Date().getTime(); const endTime = new Date().getTime();
const timeTakenInSeconds = (endTime - startTime) / 1000; const timeTakenInSeconds = (endTime - startTime) / 1000;

View File

@ -168,3 +168,6 @@ app.get('/serverHealthCheck/notify', async (req, res) => {
app.get("/is-production", (req, res) => { app.get("/is-production", (req, res) => {
res.send({ isProduction: global.isProduction }); res.send({ isProduction: global.isProduction });
}); });
// /workers health check, cant act as load balancer, just has to be a pre deploy thing

View File

@ -1,7 +1,7 @@
import { withAuth } from "../../lib/withAuth"; import { withAuth } from "../../lib/withAuth";
import { supabase_service } from "../supabase"; import { supabase_service } from "../supabase";
const FREE_CREDITS = 300; const FREE_CREDITS = 500;
export async function billTeam(team_id: string, credits: number) { export async function billTeam(team_id: string, credits: number) {
return withAuth(supaBillTeam)(team_id, credits); return withAuth(supaBillTeam)(team_id, credits);

View File

@ -2,133 +2,68 @@ import { RateLimiterRedis } from "rate-limiter-flexible";
import * as redis from "redis"; import * as redis from "redis";
import { RateLimiterMode } from "../../src/types"; import { RateLimiterMode } from "../../src/types";
const MAX_CRAWLS_PER_MINUTE_STARTER = 3; const RATE_LIMITS = {
const MAX_CRAWLS_PER_MINUTE_STANDARD = 5; crawl: {
const MAX_CRAWLS_PER_MINUTE_SCALE = 20; free: 1,
starter: 3,
const MAX_SCRAPES_PER_MINUTE_STARTER = 20; standard: 5,
const MAX_SCRAPES_PER_MINUTE_STANDARD = 40; scale: 20,
const MAX_SCRAPES_PER_MINUTE_SCALE = 50; hobby: 3,
standardNew: 10,
const MAX_SEARCHES_PER_MINUTE_STARTER = 20; growth: 50,
const MAX_SEARCHES_PER_MINUTE_STANDARD = 40; },
const MAX_SEARCHES_PER_MINUTE_SCALE = 50; scrape: {
free: 5,
const MAX_REQUESTS_PER_MINUTE_PREVIEW = 5; starter: 20,
const MAX_REQUESTS_PER_MINUTE_ACCOUNT = 20; standardOld: 40,
const MAX_REQUESTS_PER_MINUTE_CRAWL_STATUS = 150; scale: 50,
hobby: 10,
standardNew: 50,
growth: 500,
},
search: {
free: 5,
starter: 20,
standard: 40,
scale: 50,
hobby: 10,
standardNew: 50,
growth: 500,
},
preview: 5,
account: 20,
crawlStatus: 150,
testSuite: 10000,
};
export const redisClient = redis.createClient({ export const redisClient = redis.createClient({
url: process.env.REDIS_URL, url: process.env.REDIS_URL,
legacyMode: true, legacyMode: true,
}); });
export const previewRateLimiter = new RateLimiterRedis({ const createRateLimiter = (keyPrefix, points) => new RateLimiterRedis({
storeClient: redisClient, storeClient: redisClient,
keyPrefix: "preview", keyPrefix,
points: MAX_REQUESTS_PER_MINUTE_PREVIEW, points,
duration: 60, // Duration in seconds
});
export const serverRateLimiter = new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "server",
points: MAX_REQUESTS_PER_MINUTE_ACCOUNT,
duration: 60, // Duration in seconds
});
export const crawlStatusRateLimiter = new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "crawl-status",
points: MAX_REQUESTS_PER_MINUTE_CRAWL_STATUS,
duration: 60, // Duration in seconds
});
export const testSuiteRateLimiter = new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "test-suite",
points: 10000,
duration: 60, // Duration in seconds duration: 60, // Duration in seconds
}); });
export const previewRateLimiter = createRateLimiter("preview", RATE_LIMITS.preview);
export const serverRateLimiter = createRateLimiter("server", RATE_LIMITS.account);
export const crawlStatusRateLimiter = createRateLimiter("crawl-status", RATE_LIMITS.crawlStatus);
export const testSuiteRateLimiter = createRateLimiter("test-suite", RATE_LIMITS.testSuite);
export function getRateLimiter(mode: RateLimiterMode, token: string, plan?: string) { export function getRateLimiter(mode: RateLimiterMode, token: string, plan?: string) {
// Special test suite case. TODO: Change this later.
if (token.includes("5089cefa58") || token.includes("6254cf9")) { if (token.includes("5089cefa58") || token.includes("6254cf9")) {
return testSuiteRateLimiter; return testSuiteRateLimiter;
} }
switch (mode) {
case RateLimiterMode.Preview:
return previewRateLimiter; const rateLimitConfig = RATE_LIMITS[mode];
case RateLimiterMode.CrawlStatus: if (!rateLimitConfig) return serverRateLimiter;
return crawlStatusRateLimiter;
case RateLimiterMode.Crawl: const planKey = plan ? plan.replace("-", "") : "starter";
if (plan === "standard"){ const points = rateLimitConfig[planKey] || rateLimitConfig.preview;
return new RateLimiterRedis({
storeClient: redisClient, return createRateLimiter(`${mode}-${planKey}`, points);
keyPrefix: "crawl-standard",
points: MAX_CRAWLS_PER_MINUTE_STANDARD,
duration: 60, // Duration in seconds
});
} else if (plan === "scale"){
return new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "crawl-scale",
points: MAX_CRAWLS_PER_MINUTE_SCALE,
duration: 60, // Duration in seconds
});
}
return new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "crawl-starter",
points: MAX_CRAWLS_PER_MINUTE_STARTER,
duration: 60, // Duration in seconds
});
case RateLimiterMode.Scrape:
if (plan === "standard"){
return new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "scrape-standard",
points: MAX_SCRAPES_PER_MINUTE_STANDARD,
duration: 60, // Duration in seconds
});
} else if (plan === "scale"){
return new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "scrape-scale",
points: MAX_SCRAPES_PER_MINUTE_SCALE,
duration: 60, // Duration in seconds
});
}
return new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "scrape-starter",
points: MAX_SCRAPES_PER_MINUTE_STARTER,
duration: 60, // Duration in seconds
});
case RateLimiterMode.Search:
if (plan === "standard"){
return new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "search-standard",
points: MAX_SEARCHES_PER_MINUTE_STANDARD,
duration: 60, // Duration in seconds
});
} else if (plan === "scale"){
return new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "search-scale",
points: MAX_SEARCHES_PER_MINUTE_SCALE,
duration: 60, // Duration in seconds
});
}
return new RateLimiterRedis({
storeClient: redisClient,
keyPrefix: "search-starter",
points: MAX_SEARCHES_PER_MINUTE_STARTER,
duration: 60, // Duration in seconds
});
default:
return serverRateLimiter;
}
} }

View File

@ -57,6 +57,7 @@ export interface AuthResponse {
team_id?: string; team_id?: string;
error?: string; error?: string;
status?: number; status?: number;
plan?: string;
} }