update readme.md

.gitignore

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

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,28 @@
 
 ```
 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
+- 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(authConfig.Value, ",") {
+		if authorization != strings.Trim(auth, " ") {
+			err = errors.New("wrong authorization header")
+			c.AbortWithError(403, err)
+			return err
+		}
 	}
 
 	return nil

main.go

@@ -15,12 +15,10 @@ import (
 
 func main() {
 	dbAddr := flag.String("database", "./db.sqlite", "Database address")
+	upstreamsFile := flag.String("upstreams", "./upstreams.yaml", "Upstreams 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,8 +33,7 @@ func main() {
 	}
 
 	// load all upstreams
-	upstreams := make([]OPENAI_UPSTREAM, 0)
-	db.Find(&upstreams)
+	upstreams := readUpstreams(*upstreamsFile)
 	log.Println("Load upstreams number:", len(upstreams))
 
 	err = initconfig(db)
@@ -48,26 +45,9 @@ func main() {
 	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 upstreams {
 			fmt.Println(upstream.SK, upstream.Endpoint)
 		}
 		return

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,33 @@
 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
+	SK       string `yaml:"sk"`
+	Endpoint string `yaml:"endpoint"`
+	Timeout  int64  `yaml:"timeout"`
+}
+
+func readUpstreams(filepath string) []OPENAI_UPSTREAM {
+	var upstreams []OPENAI_UPSTREAM
+
+	// 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, &upstreams)
+	if err != nil {
+		log.Fatalf("Error unmarshaling YAML: %s", err)
+	}
+
+	return upstreams
 }