♻️ 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")
+}

cli/flag.go

@@ -0,0 +1,55 @@
+package cli
+
+import (
+	"flag"
+	"fmt"
+	"one-api/common"
+	"os"
+
+	"github.com/spf13/viper"
+)
+
+var (
+	port         = flag.Int("port", 0, "the listening port")
+	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")
+	export       = flag.Bool("export", false, "Exports prices to a JSON file.")
+)
+
+func FlagConfig() {
+	flag.Parse()
+
+	if *printVersion {
+		fmt.Println(common.Version)
+		os.Exit(0)
+	}
+
+	if *printHelp {
+		help()
+		os.Exit(0)
+	}
+
+	if *port != 0 {
+		viper.Set("port", *port)
+	}
+
+	if *logDir != "" {
+		viper.Set("log_dir", *logDir)
+	}
+
+	if *export {
+		ExportPrices()
+		os.Exit(0)
+	}
+
+}
+
+func help() {
+	fmt.Println("One API " + common.Version + " - All in one API service for OpenAI API.")
+	fmt.Println("Copyright (C) 2024 MartialBE. All rights reserved.")
+	fmt.Println("Original copyright holder: JustSong")
+	fmt.Println("GitHub: https://github.com/MartialBE/one-api")
+	fmt.Println("Usage: one-api [--port <port>] [--log-dir <log directory>] [--config <config.yaml path>] [--version] [--help]")
+}