♻️ refactor: Refactor price module (#123) (#109) (#128)

cli/export.go

@@ -0,0 +1,48 @@
+package cli
+
+import (
+ "encoding/json"
+ "one-api/common"
+ "one-api/relay/util"
+ "os"
+ "sort"
+)
+
+func ExportPrices() {
+ prices := util.GetPricesList("default")
+
+ if len(prices) == 0 {
+ common.SysError("No prices found")
+ return
+ }
+
+ // Sort prices by ChannelType
+ sort.Slice(prices, func(i, j int) bool {
+ if prices[i].ChannelType == prices[j].ChannelType {
+ return prices[i].Model < prices[j].Model
+ }
+ return prices[i].ChannelType < prices[j].ChannelType
+ })
+
+ // 导出到当前目录下的 prices.json 文件
+ file, err := os.Create("prices.json")
+ if err != nil {
+ common.SysError("Failed to create file: " + err.Error())
+ return
+ }
+ defer file.Close()
+
+ jsonData, err := json.MarshalIndent(prices, "", " ")
+ if err != nil {
+ common.SysError("Failed to encode prices: " + err.Error())
+ return
+ }
+
+ _, err = file.Write(jsonData)
+ if err != nil {
+ common.SysError("Failed to write to file: " + err.Error())
+ return
+ }
+
+ common.SysLog("Prices exported to prices.json")
+}

common/config/flag.go → cli/flag.go

@@ -1,4 +1,4 @@
-package config
+package cli
 
 import (
  "flag"
@@ -14,10 +14,11 @@ var (
  printVersion = flag.Bool("version", false, "print version and exit")
  printHelp = flag.Bool("help", false, "print help and exit")
  logDir = flag.String("log-dir", "", "specify the log directory")
- config = flag.String("config", "config.yaml", "specify the config.yaml path")
+ Config = flag.String("config", "config.yaml", "specify the config.yaml path")
+ export = flag.Bool("export", false, "Exports prices to a JSON file.")
 )
 
-func flagConfig() {
+func FlagConfig() {
  flag.Parse()
 
  if *printVersion {
@@ -38,6 +39,11 @@ func flagConfig() {
  viper.Set("log_dir", *logDir)
  }
 
+ if *export {
+ ExportPrices()
+ os.Exit(0)
+ }
+
 }
 
 func help() {

common/config/config.go

@@ -4,13 +4,14 @@ import (
  "strings"
  "time"
 
+ "one-api/cli"
  "one-api/common"
 
  "github.com/spf13/viper"
 )
 
 func InitConf() {
- flagConfig()
+ cli.FlagConfig()
  defaultConfig()
  setConfigFile()
  setEnv()
@@ -25,11 +26,11 @@ func InitConf() {
 }
 
 func setConfigFile() {
- if !common.IsFileExist(*config) {
+ if !common.IsFileExist(*cli.Config) {
  return
  }
 
- viper.SetConfigFile(*config)
+ viper.SetConfigFile(*cli.Config)
  if err := viper.ReadInConfig(); err != nil {
  panic(err)
  }
@@ -51,4 +52,5 @@ func defaultConfig() {
  viper.SetDefault("global.api_rate_limit", 180)
  viper.SetDefault("global.web_rate_limit", 100)
  viper.SetDefault("connect_timeout", 5)
+ viper.SetDefault("auto_price_updates", true)
 }