AI/firecrawl
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/mendableai/firecrawl synced 2025-08-05 15:50:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(db): implement bill_team RPC

Browse Source
This commit is contained in:
Gergő Móricz 2024-09-25 20:57:45 +02:00
parent c1f68c3e0a
commit f22ab5ffaf
7 changed files with 23 additions and 182 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
apps/api/src/controllers/v0/scrape.ts
View File

@ -244,7 +244,7 @@ export async function scrapeController(req: Request, res: Response) {
}
if (creditsToBeBilled > 0) {
// billing for doc done on queue end, bill only for llm extraction
billTeam(team_id, creditsToBeBilled).catch(error => {
billTeam(team_id, chunk.sub_id, creditsToBeBilled).catch(error => {
Logger.error(`Failed to bill team ${team_id} for ${creditsToBeBilled} credits: ${error}`);
// Optionally, you could notify an admin or add to a retry queue here
});

4
apps/api/src/controllers/v0/search.ts
View File

@ -18,6 +18,7 @@ export async function searchHelper(
jobId: string,
req: Request,
team_id: string,
subscription_id: string,
crawlerOptions: any,
pageOptions: PageOptions,
searchOptions: SearchOptions,
@ -54,7 +55,7 @@ export async function searchHelper(
if (justSearch) {
billTeam(team_id, res.length).catch(error => {
billTeam(team_id, subscription_id, res.length).catch(error => {
Logger.error(`Failed to bill team ${team_id} for ${res.length} credits: ${error}`);
// Optionally, you could notify an admin or add to a retry queue here
});
@ -169,6 +170,7 @@ export async function searchController(req: Request, res: Response) {
jobId,
req,
team_id,
chunk.sub_id,
crawlerOptions,
pageOptions,
searchOptions,

2
apps/api/src/controllers/v1/map.ts
View File

@ -152,7 +152,7 @@ export async function mapController(
// remove duplicates that could be due to http/https or www
links = removeDuplicateUrls(links);
billTeam(req.auth.team_id, 1).catch((error) => {
billTeam(req.auth.team_id, req.acuc.sub_id, 1).catch((error) => {
Logger.error(
`Failed to bill team ${req.auth.team_id} for 1 credit: ${error}`
);

2
apps/api/src/controllers/v1/scrape.ts
View File

@ -108,7 +108,7 @@ export async function scrapeController(
creditsToBeBilled = 5;
}
billTeam(req.auth.team_id, creditsToBeBilled).catch(error => {
billTeam(req.auth.team_id, req.acuc.sub_id, creditsToBeBilled).catch(error => {
Logger.error(`Failed to bill team ${req.auth.team_id} for ${creditsToBeBilled} credits: ${error}`);
// Optionally, you could notify an admin or add to a retry queue here
});

8
apps/api/src/controllers/v1/types.ts
View File

@ -318,10 +318,10 @@ type Account = {
export type AuthCreditUsageChunk = {
api_key: string;
team_id: string;
sub_id: string;
sub_current_period_start: string;
sub_current_period_end: string;
price_id: string;
sub_id: string | null;
sub_current_period_start: string | null;
sub_current_period_end: string | null;
price_id: string | null;
price_credits: number; // credit limit with assoicated price, or free_credits (500) if free plan
credits_used: number;
coupon_credits: number;

2
apps/api/src/main/runWebScraper.ts
View File

@ -120,7 +120,7 @@ export async function runWebScraper({
: docs;
if(is_scrape === false) {
billTeam(team_id, filteredDocs.length).catch(error => {
billTeam(team_id, undefined, filteredDocs.length).catch(error => {
Logger.error(`Failed to bill team ${team_id} for ${filteredDocs.length} credits: ${error}`);
// Optionally, you could notify an admin or add to a retry queue here
});

185
apps/api/src/services/billing/credit_billing.ts
View File

@ -3,168 +3,30 @@ import { withAuth } from "../../lib/withAuth";
import { sendNotification } from "../notification/email_notification";
import { supabase_service } from "../supabase";
import { Logger } from "../../lib/logger";
import { getValue, setValue } from "../redis";
import { redlock } from "../redlock";
import * as Sentry from "@sentry/node";
import { AuthCreditUsageChunk } from "../../controllers/v1/types";
const FREE_CREDITS = 500;
export async function billTeam(team_id: string, credits: number) {
return withAuth(supaBillTeam)(team_id, credits);
/**
* If you do not know the subscription_id in the current context, pass subscription_id as undefined.
*/
export async function billTeam(team_id: string, subscription_id: string | null | undefined, credits: number) {
return withAuth(supaBillTeam)(team_id, subscription_id, credits);
}
export async function supaBillTeam(team_id: string, credits: number) {
export async function supaBillTeam(team_id: string, subscription_id: string, credits: number) {
if (team_id === "preview") {
return { success: true, message: "Preview team, no credits used" };
}
Logger.info(`Billing team ${team_id} for ${credits} credits`);
// When the API is used, you can log the credit usage in the credit_usage table:
// team_id: The ID of the team using the API.
// subscription_id: The ID of the team's active subscription.
// credits_used: The number of credits consumed by the API call.
// created_at: The timestamp of the API usage.
// 1. get the subscription and check for available coupons concurrently
const [{ data: subscription }, { data: coupons }] = await Promise.all([
supabase_service
.from("subscriptions")
.select("*")
.eq("team_id", team_id)
.eq("status", "active")
.single(),
supabase_service
.from("coupons")
.select("id, credits")
.eq("team_id", team_id)
.eq("status", "active"),
]);
let couponCredits = 0;
let sortedCoupons = [];
if (coupons && coupons.length > 0) {
couponCredits = coupons.reduce(
(total, coupon) => total + coupon.credits,
0
);
sortedCoupons = [...coupons].sort((a, b) => b.credits - a.credits);
const { error } =
await supabase_service.rpc("bill_team", { _team_id: team_id, sub_id: subscription_id ?? null, fetch_subscription: subscription_id === undefined, credits });
if (error) {
Sentry.captureException(error);
Logger.error("Failed to bill team: " + JSON.stringify(error));
}
// using coupon credits:
if (couponCredits > 0) {
// if there is no subscription and they have enough coupon credits
if (!subscription) {
// using only coupon credits:
// if there are enough coupon credits
if (couponCredits >= credits) {
// remove credits from coupon credits
let usedCredits = credits;
while (usedCredits > 0) {
// update coupons
if (sortedCoupons[0].credits < usedCredits) {
usedCredits = usedCredits - sortedCoupons[0].credits;
// update coupon credits
await supabase_service
.from("coupons")
.update({
credits: 0,
})
.eq("id", sortedCoupons[0].id);
sortedCoupons.shift();
} else {
// update coupon credits
await supabase_service
.from("coupons")
.update({
credits: sortedCoupons[0].credits - usedCredits,
})
.eq("id", sortedCoupons[0].id);
usedCredits = 0;
}
}
return await createCreditUsage({ team_id, credits: 0 });
// not enough coupon credits and no subscription
} else {
// update coupon credits
const usedCredits = credits - couponCredits;
for (let i = 0; i < sortedCoupons.length; i++) {
await supabase_service
.from("coupons")
.update({
credits: 0,
})
.eq("id", sortedCoupons[i].id);
}
return await createCreditUsage({ team_id, credits: usedCredits });
}
}
// with subscription
// using coupon + subscription credits:
if (credits > couponCredits) {
// update coupon credits
for (let i = 0; i < sortedCoupons.length; i++) {
await supabase_service
.from("coupons")
.update({
credits: 0,
})
.eq("id", sortedCoupons[i].id);
}
const usedCredits = credits - couponCredits;
return await createCreditUsage({
team_id,
subscription_id: subscription.id,
credits: usedCredits,
});
} else {
// using only coupon credits
let usedCredits = credits;
while (usedCredits > 0) {
// update coupons
if (sortedCoupons[0].credits < usedCredits) {
usedCredits = usedCredits - sortedCoupons[0].credits;
// update coupon credits
await supabase_service
.from("coupons")
.update({
credits: 0,
})
.eq("id", sortedCoupons[0].id);
sortedCoupons.shift();
} else {
// update coupon credits
await supabase_service
.from("coupons")
.update({
credits: sortedCoupons[0].credits - usedCredits,
})
.eq("id", sortedCoupons[0].id);
usedCredits = 0;
}
}
return await createCreditUsage({
team_id,
subscription_id: subscription.id,
credits: 0,
});
}
}
// not using coupon credits
if (!subscription) {
return await createCreditUsage({ team_id, credits });
}
return await createCreditUsage({
team_id,
subscription_id: subscription.id,
credits,
});
}
export async function checkTeamCredits(chunk: AuthCreditUsageChunk, team_id: string, credits: number) {
@ -297,26 +159,3 @@ export async function countCreditsAndRemainingForCurrentBillingPeriod(
totalCredits: price.credits,
};
}
async function createCreditUsage({
team_id,
subscription_id,
credits,
}: {
team_id: string;
subscription_id?: string;
credits: number;
}) {
await supabase_service
.from("credit_usage")
.insert([
{
team_id,
credits_used: credits,
subscription_id: subscription_id || null,
created_at: new Date(),
},
]);
return { success: true };
}