feat: 新增数据看板

Qianxia666 · Jan 7, 2024 · bf8794d · bf8794d
1 parent c09df83
commit bf8794d
Show file tree

Hide file tree

Showing 16 changed files with 455 additions and 6 deletions.
diff --git a/common/constants.go b/common/constants.go
@@ -24,6 +24,9 @@ var ChatLink = ""
 var QuotaPerUnit = 500 * 1000.0 // $0.002 / 1K tokens
 var DisplayInCurrencyEnabled = true
 var DisplayTokenStatEnabled = true
+var DrawingEnabled = true
+var DataExportEnabled = true
+var DataExportInterval = 5 // unit: minute
 
 // Any options with "Secret", "Token" in its key won't be return by GetOptions
 

diff --git a/controller/misc.go b/controller/misc.go
@@ -34,6 +34,8 @@ func GetStatus(c *gin.Context) {
 			"quota_per_unit":      common.QuotaPerUnit,
 			"display_in_currency": common.DisplayInCurrencyEnabled,
 			"enable_batch_update": common.BatchUpdateEnabled,
+			"enable_drawing":      common.DrawingEnabled,
+			"enable_data_export":  common.DataExportEnabled,
 		},
 	})
 	return

diff --git a/controller/usedata.go b/controller/usedata.go
@@ -0,0 +1,24 @@
+package controller
+
+import (
+	"github.com/gin-gonic/gin"
+	"net/http"
+	"one-api/model"
+)
+
+func GetAllQuotaDates(c *gin.Context) {
+	dates, err := model.GetAllQuotaDates()
+	if err != nil {
+		c.JSON(http.StatusOK, gin.H{
+			"success": false,
+			"message": err.Error(),
+		})
+		return
+	}
+	c.JSON(http.StatusOK, gin.H{
+		"success": true,
+		"message": "",
+		"data":    dates,
+	})
+	return
+}
diff --git a/main.go b/main.go
@@ -67,6 +67,9 @@ func main() {
 		go model.SyncOptions(common.SyncFrequency)
 		go model.SyncChannelCache(common.SyncFrequency)
 	}
+	if common.DataExportEnabled {
+		go model.UpdateQuotaData(common.DataExportInterval)
+	}
 	if os.Getenv("CHANNEL_UPDATE_FREQUENCY") != "" {
 		frequency, err := strconv.Atoi(os.Getenv("CHANNEL_UPDATE_FREQUENCY"))
 		if err != nil {

diff --git a/model/log.go b/model/log.go
@@ -59,9 +59,10 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
 	if !common.LogConsumeEnabled {
 		return
 	}
+	username := GetUsernameById(userId)
 	log := &Log{
 		UserId:           userId,
-		Username:         GetUsernameById(userId),
+		Username:         username,
 		CreatedAt:        common.GetTimestamp(),
 		Type:             LogTypeConsume,
 		Content:          content,
@@ -77,6 +78,9 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
 	if err != nil {
 		common.LogError(ctx, "failed to record log: "+err.Error())
 	}
+	if common.DataExportEnabled {
+		LogQuotaData(userId, username, modelName, quota, common.GetTimestamp())
+	}
 }
 
 func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int) (logs []*Log, err error) {

diff --git a/model/main.go b/model/main.go
@@ -127,6 +127,10 @@ func InitDB() (err error) {
 		if err != nil {
 			return err
 		}
+		err = db.AutoMigrate(&QuotaData{})
+		if err != nil {
+			return err
+		}
 		common.SysLog("database migrated")
 		err = createRootAccountIfNeed()
 		return err

diff --git a/model/option.go b/model/option.go
@@ -37,6 +37,8 @@ func InitOptionMap() {
 	common.OptionMap["LogConsumeEnabled"] = strconv.FormatBool(common.LogConsumeEnabled)
 	common.OptionMap["DisplayInCurrencyEnabled"] = strconv.FormatBool(common.DisplayInCurrencyEnabled)
 	common.OptionMap["DisplayTokenStatEnabled"] = strconv.FormatBool(common.DisplayTokenStatEnabled)
+	common.OptionMap["DrawingEnabled"] = strconv.FormatBool(common.DrawingEnabled)
+	common.OptionMap["DataExportEnabled"] = strconv.FormatBool(common.DataExportEnabled)
 	common.OptionMap["ChannelDisableThreshold"] = strconv.FormatFloat(common.ChannelDisableThreshold, 'f', -1, 64)
 	common.OptionMap["EmailDomainRestrictionEnabled"] = strconv.FormatBool(common.EmailDomainRestrictionEnabled)
 	common.OptionMap["EmailDomainWhitelist"] = strings.Join(common.EmailDomainWhitelist, ",")
@@ -76,6 +78,7 @@ func InitOptionMap() {
 	common.OptionMap["ChatLink"] = common.ChatLink
 	common.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(common.QuotaPerUnit, 'f', -1, 64)
 	common.OptionMap["RetryTimes"] = strconv.Itoa(common.RetryTimes)
+	common.OptionMap["DataExportInterval"] = strconv.Itoa(common.DataExportInterval)
 
 	common.OptionMapRWMutex.Unlock()
 	loadOptionsFromDatabase()
@@ -157,6 +160,12 @@ func updateOptionMap(key string, value string) (err error) {
 			common.LogConsumeEnabled = boolValue
 		case "DisplayInCurrencyEnabled":
 			common.DisplayInCurrencyEnabled = boolValue
+		case "DisplayTokenStatEnabled":
+			common.DisplayTokenStatEnabled = boolValue
+		case "DrawingEnabled":
+			common.DrawingEnabled = boolValue
+		case "DataExportEnabled":
+			common.DataExportEnabled = boolValue
 		}
 	}
 	switch key {
@@ -217,6 +226,8 @@ func updateOptionMap(key string, value string) (err error) {
 		common.PreConsumedQuota, _ = strconv.Atoi(value)
 	case "RetryTimes":
 		common.RetryTimes, _ = strconv.Atoi(value)
+	case "DataExportInterval":
+		common.DataExportInterval, _ = strconv.Atoi(value)
 	case "ModelRatio":
 		err = common.UpdateModelRatioByJSONString(value)
 	case "GroupRatio":

diff --git a/model/usedata.go b/model/usedata.go
@@ -0,0 +1,87 @@
+package model
+
+import (
+	"fmt"
+	"one-api/common"
+	"time"
+)
+
+// QuotaData 柱状图数据
+type QuotaData struct {
+	Id        int    `json:"id"`
+	UserID    int    `json:"user_id" gorm:"index"`
+	Username  string `json:"username" gorm:"index:index_quota_data_model_user_name,priority:2;default:''"`
+	ModelName string `json:"model_name" gorm:"index;index:index_quota_data_model_user_name,priority:1;default:''"`
+	CreatedAt int64  `json:"created_at" gorm:"bigint;index:index_quota_data_created_at,priority:2"`
+	Count     int    `json:"count" gorm:"default:0"`
+	Quota     int    `json:"quota" gorm:"default:0"`
+}
+
+func UpdateQuotaData(frequency int) {
+	for {
+		common.SysLog("正在更新数据看板数据...")
+		SaveQuotaDataCache()
+		time.Sleep(time.Duration(frequency) * time.Minute)
+	}
+}
+
+var CacheQuotaData = make(map[string]*QuotaData)
+
+func LogQuotaDataCache(userId int, username string, modelName string, quota int, createdAt int64) {
+	// 只精确到小时
+	createdAt = createdAt - (createdAt % 3600)
+	key := fmt.Sprintf("%d-%s-%s-%d", userId, username, modelName, createdAt)
+	quotaData, ok := CacheQuotaData[key]
+	if ok {
+		quotaData.Count += 1
+		quotaData.Quota += quota
+	} else {
+		quotaData = &QuotaData{
+			UserID:    userId,
+			Username:  username,
+			ModelName: modelName,
+			CreatedAt: createdAt,
+			Count:     1,
+			Quota:     quota,
+		}
+	}
+	CacheQuotaData[key] = quotaData
+}
+
+func LogQuotaData(userId int, username string, modelName string, quota int, createdAt int64) {
+	LogQuotaDataCache(userId, username, modelName, quota, createdAt)
+}
+
+func SaveQuotaDataCache() {
+	// 如果缓存中有数据，就保存到数据库中
+	// 1. 先查询数据库中是否有数据
+	// 2. 如果有数据，就更新数据
+	// 3. 如果没有数据，就插入数据
+	for _, quotaData := range CacheQuotaData {
+		quotaDataDB := &QuotaData{}
+		DB.Table("quota_data").Where("user_id = ? and username = ? and model_name = ? and created_at = ?",
+			quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt).First(quotaDataDB)
+		if quotaDataDB.Id > 0 {
+			quotaDataDB.Count += quotaData.Count
+			quotaDataDB.Quota += quotaData.Quota
+			DB.Table("quota_data").Save(quotaDataDB)
+		} else {
+			DB.Table("quota_data").Create(quotaData)
+		}
+	}
+	CacheQuotaData = make(map[string]*QuotaData)
+}
+
+func GetQuotaDataByUsername(username string) (quotaData []*QuotaData, err error) {
+	var quotaDatas []*QuotaData
+	// 从quota_data表中查询数据
+	err = DB.Table("quota_data").Where("username = ?", username).Find(&quotaDatas).Error
+	return quotaDatas, err
+}
+
+func GetAllQuotaDates() (quotaData []*QuotaData, err error) {
+	var quotaDatas []*QuotaData
+	// 从quota_data表中查询数据
+	err = DB.Table("quota_data").Find(&quotaDatas).Error
+	return quotaDatas, err
+}
diff --git a/router/api-router.go b/router/api-router.go
@@ -114,6 +114,9 @@ func SetApiRouter(router *gin.Engine) {
 		logRoute.GET("/self", middleware.UserAuth(), controller.GetUserLogs)
 		logRoute.GET("/self/search", middleware.UserAuth(), controller.SearchUserLogs)
 
+		dataRoute := apiRouter.Group("/data")
+		dataRoute.GET("/", middleware.AdminAuth(), controller.GetAllQuotaDates)
+
 		logRoute.Use(middleware.CORS())
 		{
 			logRoute.GET("/token", controller.GetLogByKey)

diff --git a/web/package.json b/web/package.json
@@ -3,7 +3,10 @@
   "version": "0.1.0",
   "private": true,
   "dependencies": {
-    "@douyinfe/semi-ui": "^2.45.2",
+    "@douyinfe/semi-ui": "^2.46.1",
+    "@visactor/vchart": "~1.7.2",
+    "@visactor/react-vchart": "~1.7.2",
+    "@visactor/vchart-semi-theme": "~1.7.2",
     "axios": "^0.27.2",
     "history": "^5.3.0",
     "marked": "^4.1.1",
@@ -44,7 +47,8 @@
     ]
   },
   "devDependencies": {
-    "prettier": "^2.7.1"
+    "prettier": "^2.7.1",
+    "typescript": "4.4.2"
   },
   "prettier": {
     "singleQuote": true,

diff --git a/web/src/App.js b/web/src/App.js
@@ -23,10 +23,10 @@ import Log from './pages/Log';
 import Chat from './pages/Chat';
 import {Layout} from "@douyinfe/semi-ui";
 import Midjourney from "./pages/Midjourney";
+import Detail from "./pages/Detail";
 
 const Home = lazy(() => import('./pages/Home'));
 const About = lazy(() => import('./pages/About'));
-
 function App() {
   const [userState, userDispatch] = useContext(UserContext);
   const [statusState, statusDispatch] = useContext(StatusContext);
@@ -49,6 +49,8 @@ function App() {
       localStorage.setItem('footer_html', data.footer_html);
       localStorage.setItem('quota_per_unit', data.quota_per_unit);
       localStorage.setItem('display_in_currency', data.display_in_currency);
+      localStorage.setItem('enable_drawing', data.enable_drawing);
+      localStorage.setItem('enable_data_export', data.enable_data_export);
       if (data.chat_link) {
         localStorage.setItem('chat_link', data.chat_link);
       } else {
@@ -228,6 +230,14 @@ function App() {
                         </PrivateRoute>
                     }
                 />
+                <Route
+                    path='/detail'
+                    element={
+                        <PrivateRoute>
+                            <Detail />
+                        </PrivateRoute>
+                    }
+                />
                 <Route
                     path='/midjourney'
                     element={

diff --git a/web/src/helpers/render.js b/web/src/helpers/render.js
@@ -76,6 +76,12 @@ export function getQuotaPerUnit() {
   return quotaPerUnit;
 }
 
+export function getQuotaWithUnit(quota, digits = 6) {
+  let quotaPerUnit = localStorage.getItem('quota_per_unit');
+  quotaPerUnit = parseFloat(quotaPerUnit);
+  return (quota / quotaPerUnit).toFixed(digits);
+}
+
 export function renderQuota(quota, digits = 2) {
   let quotaPerUnit = localStorage.getItem('quota_per_unit');
   let displayInCurrency = localStorage.getItem('display_in_currency');

diff --git a/web/src/helpers/utils.js b/web/src/helpers/utils.js
@@ -171,6 +171,32 @@ export function timestamp2string(timestamp) {
   );
 }
 
+export function timestamp2string1(timestamp) {
+  let date = new Date(timestamp * 1000);
+  // let year = date.getFullYear().toString();
+  let month = (date.getMonth() + 1).toString();
+  let day = date.getDate().toString();
+  let hour = date.getHours().toString();
+  if (month.length === 1) {
+    month = '0' + month;
+  }
+  if (day.length === 1) {
+    day = '0' + day;
+  }
+  if (hour.length === 1) {
+    hour = '0' + hour;
+  }
+  return (
+      // year +
+      // '-' +
+      month +
+      '-' +
+      day +
+      ' ' +
+      hour + ":00"
+  );
+}
+
 export function downloadTextAsFile(text, filename) {
   let blob = new Blob([text], { type: 'text/plain;charset=utf-8' });
   let url = URL.createObjectURL(blob);

diff --git a/web/src/index.js b/web/src/index.js
@@ -1,7 +1,8 @@
+import { initVChartSemiTheme } from '@visactor/vchart-semi-theme';
+import VChart from "@visactor/vchart";
 import React from 'react';
 import ReactDOM from 'react-dom/client';
 import {BrowserRouter} from 'react-router-dom';
-import {Container} from 'semantic-ui-react';
 import App from './App';
 import HeaderBar from './components/HeaderBar';
 import Footer from './components/Footer';
@@ -14,6 +15,11 @@ import {StatusProvider} from './context/Status';
 import {Layout} from "@douyinfe/semi-ui";
 import SiderBar from "./components/SiderBar";
 
+// initialization
+initVChartSemiTheme({
+    isWatchingThemeSwitch: true,
+});
+
 const root = ReactDOM.createRoot(document.getElementById('root'));
 const {Sider, Content, Header} = Layout;
 root.render(

diff --git a/web/src/pages/About/index.js b/web/src/pages/About/index.js
@@ -1,5 +1,4 @@
 import React, { useEffect, useState } from 'react';
-import { Header, Segment } from 'semantic-ui-react';
 import { API, showError } from '../../helpers';
 import { marked } from 'marked';
 import {Layout} from "@douyinfe/semi-ui";