fix: fix openai balance calc logic

This commit is contained in:
moonrailgun 2024-01-13 00:55:08 +08:00
parent bf887ae8fc
commit 1fccb103fd
3 changed files with 66 additions and 9 deletions

View File

@ -43,7 +43,7 @@ export const MonitorOpenaiOverview: MonitorOverviewComponent = React.memo(
<MonitorStatsBlock
title="Usage"
desc={dayjs(data.updatedAt).format('YYYY-MM-DD')}
text={`$${payload.totalUsed} / $${payload.totalGranted}`}
text={`$${payload.totalUsed} / $${payload.allUSD}`}
/>
);
}

View File

@ -1,4 +1,4 @@
import { describe, test } from 'vitest';
import { describe, expect, test } from 'vitest';
import { getBillingCreditGrants } from '../openai';
describe.runIf(!!process.env.OPENAI_SESS_KEY)('openai', () => {
@ -7,6 +7,7 @@ describe.runIf(!!process.env.OPENAI_SESS_KEY)('openai', () => {
String(process.env.OPENAI_SESS_KEY)
);
console.log(res);
expect(typeof res.allUSD).toBe('number');
expect(typeof res.totalUsed).toBe('number');
});
});

View File

@ -2,6 +2,7 @@ import { z } from 'zod';
import { MonitorProvider } from './type';
import axios from 'axios';
import { saveMonitorStatus } from './_utils';
import _ from 'lodash';
const openaiCreditGrantsSchema = z.object({
object: z.string(),
@ -24,6 +25,47 @@ const openaiCreditGrantsSchema = z.object({
}),
});
const openaiSubscriptionSchema = z.object({
object: z.string(),
has_payment_method: z.boolean(),
canceled: z.boolean(),
canceled_at: z.any(),
delinquent: z.any(),
access_until: z.number(),
soft_limit: z.number(),
hard_limit: z.number(),
system_hard_limit: z.number(),
soft_limit_usd: z.number(),
hard_limit_usd: z.number(),
system_hard_limit_usd: z.number(),
plan: z.object({
title: z.string(),
id: z.string(),
}),
primary: z.boolean(),
billing_mechanism: z.any(),
is_arrears_eligible: z.boolean(),
max_balance: z.number(),
auto_recharge_eligible: z.boolean(),
auto_recharge_enabled: z.boolean(),
auto_recharge_threshold: z.any(),
auto_recharge_to_balance: z.any(),
trust_tier: z.string(),
account_name: z.string(),
po_number: z.any(),
billing_email: z.any(),
tax_ids: z.any(),
billing_address: z.object({
city: z.string(),
line1: z.string(),
line2: z.string(),
state: z.string(),
country: z.string(),
postal_code: z.string(),
}),
business_address: z.any(),
});
export const openai: MonitorProvider<{
sessionKey: string;
}> = {
@ -36,13 +78,11 @@ export const openai: MonitorProvider<{
const res = await getBillingCreditGrants(sessionKey);
const balance = (res.total_granted - res.total_used) * 100;
const balance = (res.allUSD - res.totalUsed) * 100;
await saveMonitorStatus(monitor.id, 'credit', {
totalGranted: res.total_granted,
totalUsed: res.total_used,
totalAvailable: res.total_available,
totalPaidAvailable: res.total_paid_available,
allUSD: res.allUSD,
totalUsed: res.totalUsed,
});
if (balance <= 0) {
@ -54,6 +94,20 @@ export const openai: MonitorProvider<{
};
export async function getBillingCreditGrants(sessionKey: string) {
const { data: count_info } = await axios.get(
'https://api.openai.com/dashboard/billing/subscription',
{
headers: {
'Content-Type': 'application/json',
Authorization: `Bearer ${sessionKey}`,
},
}
);
const allUSD = Number(
openaiSubscriptionSchema.parse(count_info).system_hard_limit_usd.toFixed(2)
);
const { data } = await axios.get(
'https://api.openai.com/dashboard/billing/credit_grants',
{
@ -64,5 +118,7 @@ export async function getBillingCreditGrants(sessionKey: string) {
}
);
return openaiCreditGrantsSchema.parse(data);
const totalUsed = openaiCreditGrantsSchema.parse(data).total_used;
return { allUSD, totalUsed };
}