update config to file

README.md

@@ -88,6 +88,15 @@ Usage of ./openai-api-route:
 
 ```yaml
 authorization: woshimima
+
+# 使用 sqlite 作为数据库储存请求记录
+dbtype: sqlite
+dbaddr: ./db.sqlite
+
+# 使用 postgres 作为数据库储存请求记录
+# dbtype: postgres
+# dbaddr: "host=127.0.0.1 port=5432 user=postgres dbname=openai_api_route sslmode=disable password=woshimima"
+
 upstreams:
   - sk: "secret_key_1"
     endpoint: "https://api.openai.com/v2"

main.go

@@ -18,10 +18,7 @@ import (
 var config Config
 
 func main() {
-	dbType := flag.String("dbtype", "sqlite", "Database type (sqlite or postgres)")
-	dbAddr := flag.String("database", "./db.sqlite", "Database address, if dbType is postgres, this is the DSN connection string")
 	configFile := flag.String("config", "./config.yaml", "Config file")
-	listenAddr := flag.String("addr", ":8888", "Listening address")
 	listMode := flag.Bool("list", false, "List all upstream")
 	noauth := flag.Bool("noauth", false, "Do not check incoming authorization header")
 	flag.Parse()
@@ -31,19 +28,19 @@ func main() {
 	// connect to database
 	var db *gorm.DB
 	var err error
-	switch *dbType {
+	switch config.DBType {
 	case "sqlite":
-		db, err = gorm.Open(sqlite.Open(*dbAddr), &gorm.Config{
+		db, err = gorm.Open(sqlite.Open(config.DBAddr), &gorm.Config{
 			PrepareStmt:            true,
 			SkipDefaultTransaction: true,
 		})
 	case "postgres":
-		db, err = gorm.Open(postgres.Open(*dbAddr), &gorm.Config{
+		db, err = gorm.Open(postgres.Open(config.DBAddr), &gorm.Config{
 			PrepareStmt:            true,
 			SkipDefaultTransaction: true,
 		})
 	default:
-		log.Fatalf("Unsupported database type: %s", *dbType)
+		log.Fatalf("Unsupported database type: %s", config.DBType)
 	}
 
 	// load all upstreams
@@ -144,5 +141,5 @@ func main() {
 		}
 	})
 
-	engine.Run(*listenAddr)
+	engine.Run(config.Address)
 }

structure.go

@@ -8,6 +8,9 @@ import (
 )
 
 type Config struct {
+	Address       string            `yaml:"address"`
+	DBType        string            `yaml:"dbtype"`
+	DBAddr        string            `yaml:"dbaddr"`
 	Authorization string            `yaml:"authorization"`
 	Upstreams     []OPENAI_UPSTREAM `yaml:"upstreams"`
 }
@@ -32,5 +35,16 @@ func readConfig(filepath string) Config {
 		log.Fatalf("Error unmarshaling YAML: %s", err)
 	}
 
+	// set default value
+	if config.Address == "" {
+		config.Address = ":8888"
+	}
+	if config.DBType == "" {
+		config.DBType = "sqlite"
+	}
+	if config.DBAddr == "" {
+		config.DBAddr = "./db.sqlite"
+	}
+
 	return config
 }