convert config to yaml file

.dockerignore

@@ -0,0 +1,3 @@
+openai-api-route
+db.sqlite
+/config.yaml

.gitignore

@@ -1,2 +1,3 @@
 openai-api-route
 db.sqlite
+/config.yaml

Dockerfile

@@ -0,0 +1,17 @@
+FROM golang:1.21 as builder
+
+WORKDIR /app
+
+COPY go.mod go.sum ./
+
+RUN go mod download
+
+COPY . .
+
+RUN make
+
+FROM alpine
+
+COPY --from=builder /app/openai-api-route /openai-api-route
+
+ENTRYPOINT ["/openai-api-route"]

README.md

@@ -8,7 +8,7 @@
 - 支持所有类型的接口 (`/v1/*`)
 - 提供 Prometheus Metrics 统计接口 (`/v1/metrics`)
 - 按照定义顺序请求 OpenAI 上游
-- 识别 ChatCompletions Stream 请求，针对 Stream 请求使用 5 秒超时。对于其他请求使用60秒超时。
+- 识别 ChatCompletions Stream 请求，针对 Stream 请求使用 5 秒超时。对于其他请求使用 60 秒超时。
 - 记录完整的请求内容、使用的上游、IP 地址、响应时间以及 GPT 回复文本
 - 请求出错时发送 飞书 或 Matrix 消息通知
 
@@ -27,37 +27,37 @@
 3. 打开终端，并进入到仓库目录中。
 
 4. 在终端中执行以下命令来编译代码：
-   
+
    ```
    make
    ```
-   
+
    这将会编译代码并生成可执行文件。
 
 5. 编译成功后，您可以直接运行以下命令来启动负载均衡 API：
-   
+
    ```
    ./openai-api-route
    ```
-   
+
    默认情况下，API 将会在本地的 8888 端口进行监听。
-   
+
    如果您希望使用不同的监听地址，可以使用 `-addr` 参数来指定，例如：
-   
+
    ```
    ./openai-api-route -addr 0.0.0.0:8080
    ```
-   
+
    这将会将监听地址设置为 0.0.0.0:8080。
 
 6. 如果数据库不存在，系统会自动创建一个名为 `db.sqlite` 的数据库文件。
-   
+
    如果您希望使用不同的数据库地址，可以使用 `-database` 参数来指定，例如：
-   
+
    ```
    ./openai-api-route -database /path/to/database.db
    ```
-   
+
    这将会将数据库地址设置为 `/path/to/database.db`。
 
 7. 现在，您已经成功编译并运行了负载均衡和能力 API。您可以根据需要添加上游、管理上游，并使用 API 进行相关操作。
@@ -68,30 +68,30 @@
 
 ```
 Usage of ./openai-api-route:
-  -add
-        添加一个 OpenAI 上游
   -addr string
         监听地址（默认为 ":8888"）
+  -upstreams string
+        上游配置文件（默认为 "./upstreams.yaml"）
   -database string
         数据库地址（默认为 "./db.sqlite"）
-  -endpoint string
-        OpenAI API 基地址（默认为 "https://api.openai.com/v1"）
   -list
         列出所有上游
   -noauth
         不检查传入的授权头
-  -sk string
-        OpenAI API 密钥（sk-xxxxx）
 ```
 
+以下是一个 `./upstreams.yaml` 文件配置示例
+
+```yaml
+authorization: woshimima
+upstreams:
+  - sk: "secret_key_1"
+    endpoint: "https://api.openai.com/v2"
+  - sk: "secret_key_2"
+    endpoint: "https://api.openai.com/v1"
+    timeout: 30
+```
+
+请注意，程序会根据情况修改 timeout 的值
+
 您可以直接运行 `./openai-api-route` 命令，如果数据库不存在，系统会自动创建。
-
-### 上游管理
-
-您可以使用以下命令添加一个上游：
-
-```bash
-./openai-api-route -add -sk sk-xxxxx -endpoint https://api.openai.com/v1
-```
-
-另外，您还可以直接编辑数据库中的 `openai_upstreams` 表进行 OpenAI 上游的增删改查管理。改动的上游需要重启负载均衡服务后才能生效。

auth.go

@@ -21,10 +21,12 @@ func handleAuth(c *gin.Context) error {
 	authorization = strings.Trim(authorization[len("Bearer"):], " ")
 	log.Println("Received authorization", authorization)
 
-	if authorization != authConfig.Value {
-		err = errors.New("wrong authorization header")
-		c.AbortWithError(403, err)
-		return err
+	for _, auth := range strings.Split(config.Authorization, ",") {
+		if authorization != strings.Trim(auth, " ") {
+			err = errors.New("wrong authorization header")
+			c.AbortWithError(403, err)
+			return err
+		}
 	}
 
 	return nil

config.go

@@ -1,49 +0,0 @@
-package main
-
-import (
-	"errors"
-	"log"
-
-	"gorm.io/gorm"
-)
-
-// K-V struct to store program's config
-type ConfigKV struct {
-	gorm.Model
-	Key   string `gorm:"unique"`
-	Value string
-}
-
-// init db
-func initconfig(db *gorm.DB) error {
-	var err error
-
-	err = db.AutoMigrate(&ConfigKV{})
-	if err != nil {
-		return err
-	}
-
-	// config list and their default values
-	configs := make(map[string]string)
-	configs["authorization"] = "woshimima"
-	configs["policy"] = "main"
-
-	for key, value := range configs {
-		kv := ConfigKV{}
-		err = db.Take(&kv, "key = ?", key).Error
-		if err != nil {
-			if errors.Is(err, gorm.ErrRecordNotFound) {
-				log.Println("Missing config", key, "creating with value", value)
-				kv.Key = key
-				kv.Value = value
-				if err = db.Create(&kv).Error; err != nil {
-					return err
-				}
-			} else {
-				return err
-			}
-		}
-	}
-
-	return nil
-}

global.go

@@ -1,5 +0,0 @@
-package main
-
-// declare global variable
-
-var authConfig ConfigKV

main.go

@@ -13,14 +13,15 @@ import (
 	"gorm.io/gorm"
 )
 
+// global config
+var config Config
+
 func main() {
 	dbAddr := flag.String("database", "./db.sqlite", "Database address")
+	configFile := flag.String("config", "./config.yaml", "Config file")
 	listenAddr := flag.String("addr", ":8888", "Listening address")
-	addMode := flag.Bool("add", false, "Add an OpenAI upstream")
 	listMode := flag.Bool("list", false, "List all upstream")
-	sk := flag.String("sk", "", "OpenAI API key (sk-xxxxx)")
 	noauth := flag.Bool("noauth", false, "Do not check incoming authorization header")
-	endpoint := flag.String("endpoint", "https://api.openai.com/v1", "OpenAI API base")
 	flag.Parse()
 
 	log.Println("Service starting")
@@ -35,39 +36,15 @@ func main() {
 	}
 
 	// load all upstreams
-	upstreams := make([]OPENAI_UPSTREAM, 0)
-	db.Find(&upstreams)
-	log.Println("Load upstreams number:", len(upstreams))
+	config = readConfig(*configFile)
+	log.Println("Load upstreams number:", len(config.Upstreams))
 
-	err = initconfig(db)
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	db.AutoMigrate(&OPENAI_UPSTREAM{})
 	db.AutoMigrate(&Record{})
 	log.Println("Auto migrate database done")
 
-	if *addMode {
-		if *sk == "" {
-			log.Fatal("Missing --sk flag")
-		}
-		newUpstream := OPENAI_UPSTREAM{}
-		newUpstream.SK = *sk
-		newUpstream.Endpoint = *endpoint
-		err = db.Create(&newUpstream).Error
-		if err != nil {
-			log.Fatal("Can not add upstream", err)
-		}
-		log.Println("Successuflly add upstream", *sk, *endpoint)
-		return
-	}
-
 	if *listMode {
-		result := make([]OPENAI_UPSTREAM, 0)
-		db.Find(&result)
 		fmt.Println("SK\tEndpoint")
-		for _, upstream := range result {
+		for _, upstream := range config.Upstreams {
 			fmt.Println(upstream.SK, upstream.Endpoint)
 		}
 		return
@@ -102,9 +79,6 @@ func main() {
 		ctx.AbortWithStatus(200)
 	})
 
-	// get authorization config from db
-	db.Take(&authConfig, "key = ?", "authorization")
-
 	engine.POST("/v1/*any", func(c *gin.Context) {
 		record := Record{
 			IP:            c.ClientIP(),
@@ -125,15 +99,15 @@ func main() {
 			}
 		}
 
-		for index, upstream := range upstreams {
+		for index, upstream := range config.Upstreams {
 			if upstream.Endpoint == "" || upstream.SK == "" {
 				c.AbortWithError(500, fmt.Errorf("invaild upstream '%s' '%s'", upstream.SK, upstream.Endpoint))
 				continue
 			}
 
-			shouldResponse := index == len(upstreams)-1
+			shouldResponse := index == len(config.Upstreams)-1
 
-			if len(upstreams) == 1 {
+			if len(config.Upstreams) == 1 {
 				upstream.Timeout = 120
 			}
 

process.go

@@ -20,9 +20,8 @@ import (
 func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, shouldResponse bool) error {
 	var errCtx error
 
-	record.UpstreamID = upstream.ID
+	record.UpstreamEndpoint = upstream.Endpoint
 	record.Response = ""
-	record.Authorization = upstream.SK
 	// [TODO] record request body
 
 	// reverse proxy
@@ -134,6 +133,7 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s
 	}
 	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
 		haveResponse = true
+		record.ResponseTime = time.Now().Sub(record.CreatedAt)
 		log.Println("Error", err, upstream.SK, upstream.Endpoint)
 
 		errCtx = err

record.go

@@ -11,16 +11,18 @@ import (
 )
 
 type Record struct {
-	ID            int64 `gorm:"primaryKey,autoIncrement"`
-	CreatedAt     time.Time
-	IP            string
-	Body          string `gorm:"serializer:json"`
-	Model         string
-	Response      string
-	ElapsedTime   time.Duration
-	Status        int
-	UpstreamID    uint
-	Authorization string
+	ID               int64 `gorm:"primaryKey,autoIncrement"`
+	UpstreamEndpoint string
+	CreatedAt        time.Time
+	IP               string
+	Body             string `gorm:"serializer:json"`
+	Model            string
+	Response         string
+	ResponseTime     time.Duration
+	ElapsedTime      time.Duration
+	Status           int
+	UpstreamID       uint
+	Authorization    string
 }
 
 type StreamModeChunk struct {

structure.go

@@ -1,13 +1,36 @@
 package main
 
 import (
-	"gorm.io/gorm"
+	"log"
+	"os"
+
+	"gopkg.in/yaml.v3"
 )
 
-// one openai upstream contain a pair of key and endpoint
-type OPENAI_UPSTREAM struct {
-	gorm.Model
-	SK       string `gorm:"index:idx_sk_endpoint,unique"` // key
-	Endpoint string `gorm:"index:idx_sk_endpoint,unique"` // endpoint
-	Timeout  int64  // timeout in seconds
+type Config struct {
+	Authorization string            `yaml:"authorization"`
+	Upstreams     []OPENAI_UPSTREAM `yaml:"upstreams"`
+}
+type OPENAI_UPSTREAM struct {
+	SK       string `yaml:"sk"`
+	Endpoint string `yaml:"endpoint"`
+	Timeout  int64  `yaml:"timeout"`
+}
+
+func readConfig(filepath string) Config {
+	var config Config
+
+	// read yaml file
+	data, err := os.ReadFile(filepath)
+	if err != nil {
+		log.Fatalf("Error reading YAML file: %s", err)
+	}
+
+	// Unmarshal the YAML into the upstreams slice
+	err = yaml.Unmarshal(data, &config)
+	if err != nil {
+		log.Fatalf("Error unmarshaling YAML: %s", err)
+	}
+
+	return config
 }