From 1fccb103fd79fad312dbf9403d935e0726da95af Mon Sep 17 00:00:00 2001
From: moonrailgun <moonrailgun@gmail.com>
Date: Sat, 13 Jan 2024 00:55:08 +0800
Subject: [PATCH] fix: fix openai balance calc logic

---
 .../components/monitor/provider/openai.tsx    |  2 +-
 .../monitor/provider/__tests__/openai.spec.ts |  5 +-
 src/server/model/monitor/provider/openai.ts   | 68 +++++++++++++++++--
 3 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/src/client/components/monitor/provider/openai.tsx b/src/client/components/monitor/provider/openai.tsx
index 9cb7455..91665a6 100644
--- a/src/client/components/monitor/provider/openai.tsx
+++ b/src/client/components/monitor/provider/openai.tsx
@@ -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}`}
       />
     );
   }
diff --git a/src/server/model/monitor/provider/__tests__/openai.spec.ts b/src/server/model/monitor/provider/__tests__/openai.spec.ts
index b7e9148..e72e92f 100644
--- a/src/server/model/monitor/provider/__tests__/openai.spec.ts
+++ b/src/server/model/monitor/provider/__tests__/openai.spec.ts
@@ -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');
   });
 });
diff --git a/src/server/model/monitor/provider/openai.ts b/src/server/model/monitor/provider/openai.ts
index 79c7b13..8147c2a 100644
--- a/src/server/model/monitor/provider/openai.ts
+++ b/src/server/model/monitor/provider/openai.ts
@@ -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 };
 }