fix: record resp time

.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 秒超时。具体超时策略请参阅 [超时策略](#超时策略) 一节
 - 记录完整的请求内容、使用的上游、IP 地址、响应时间以及 GPT 回复文本
 - 请求出错时发送 飞书 或 Matrix 消息通知
 
@@ -68,30 +68,67 @@
 
 ```
 Usage of ./openai-api-route:
-  -add
-        添加一个 OpenAI 上游
   -addr string
         监听地址（默认为 ":8888"）
+  -upstreams string
+        上游配置文件（默认为 "./upstreams.yaml"）
+  -dbtype
+        数据库类型 (sqlite 或 postgres，默认为 sqlite)
   -database string
         数据库地址（默认为 "./db.sqlite"）
-  -endpoint string
+        如果数据库为 postgres ，则此值应 PostgreSQL DSN 格式
-        OpenAI API 基地址（默认为 "https://api.openai.com/v1"）
+        例如 "host=127.0.0.1 port=5432 user=postgres dbname=openai_api_route sslmode=disable password=woshimima"
   -list
         列出所有上游
   -noauth
         不检查传入的授权头
-  -sk string
-        OpenAI API 密钥（sk-xxxxx）
 ```
 
+以下是一个 `./upstreams.yaml` 文件配置示例
+
+```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"
+  - sk: "secret_key_2"
+    endpoint: "https://api.openai.com/v1"
+    timeout: 30
+```
+
+请注意，程序会根据情况修改 timeout 的值
+
 您可以直接运行 `./openai-api-route` 命令，如果数据库不存在，系统会自动创建。
 
-### 上游管理
+## 超时策略
 
-您可以使用以下命令添加一个上游：
+在处理上游请求时，超时策略是确保服务稳定性和响应性的关键因素。本服务通过配置文件中的 `Upstreams` 部分来定义多个上游服务器。每个上游服务器都有自己的 `Endpoint` 和 `SK`（可能是密钥或特殊标识）。服务会按照配置文件中的顺序依次尝试每个上游服务器，直到请求成功或所有上游服务器都已尝试。
 
-```bash
+### 单一上游配置
-./openai-api-route -add -sk sk-xxxxx -endpoint https://api.openai.com/v1
-```
 
-另外，您还可以直接编辑数据库中的 `openai_upstreams` 表进行 OpenAI 上游的增删改查管理。改动的上游需要重启负载均衡服务后才能生效。
+当配置文件中只定义了一个上游服务器时，该上游的超时时间将被设置为 120 秒。这意味着，如果请求没有在 120 秒内得到上游服务器的响应，服务将会中止该请求并可能返回错误。
+
+### 多上游配置
+
+如果配置文件中定义了多个上游服务器，服务将会按照定义的顺序依次尝试每个上游。对于每个上游服务器，服务会检查其 `Endpoint` 和 `SK` 是否有效。如果任一字段为空，服务将返回 500 错误，并记录无效的上游信息。
+
+### 超时策略细节
+
+服务在处理请求时会根据不同的条件设置不同的超时时间。超时时间是指服务等待上游服务器响应的最大时间。以下是超时时间的设置规则：
+
+1. **默认超时时间**：如果没有特殊条件，服务将使用默认的超时时间，即 60 秒。
+
+2. **流式请求**：如果请求体被识别为流式（`requestBody.Stream` 为 `true`），并且请求体检查（`requestBodyOK`）没有发现问题，超时时间将被设置为 5 秒。这适用于那些预期会快速响应的流式请求。
+
+3. **大请求体**：如果请求体的大小超过 128KB（即 `len(inBody) > 1024*128`），超时时间将被设置为 20 秒。这考虑到了处理大型数据可能需要更长的时间。
+
+4. **上游超时配置**：如果上游服务器在配置中指定了超时时间（`upstream.Timeout` 大于 0），服务将使用该值作为超时时间。这个值是以秒为单位的。

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 {
+	for _, auth := range strings.Split(config.Authorization, ",") {
-		err = errors.New("wrong authorization header")
+		if authorization != strings.Trim(auth, " ") {
-		c.AbortWithError(403, err)
+			err = errors.New("wrong authorization header")
-		return err
+			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
-}

config.sample.yaml

@@ -0,0 +1,16 @@
+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"
+  - sk: "secret_key_2"
+    endpoint: "https://api.openai.com/v1"
+    timeout: 30

global.go

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

go.mod

@@ -4,9 +4,12 @@ go 1.20
 
 require (
 	github.com/gin-gonic/gin v1.9.1
-	github.com/google/uuid v1.3.0
+	github.com/penglongli/gin-metrics v0.1.10
+	golang.org/x/net v0.10.0
+	gopkg.in/yaml.v3 v3.0.1
+	gorm.io/driver/postgres v1.5.4
 	gorm.io/driver/sqlite v1.5.2
-	gorm.io/gorm v1.25.2
+	gorm.io/gorm v1.25.5
 )
 
 require (
@@ -22,6 +25,9 @@ require (
 	github.com/go-playground/validator/v10 v10.14.0 // indirect
 	github.com/goccy/go-json v0.10.2 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+	github.com/jackc/pgx/v5 v5.4.3 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
@@ -33,19 +39,17 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/pelletier/go-toml/v2 v2.0.8 // indirect
-	github.com/penglongli/gin-metrics v0.1.10 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.12.0 // indirect
 	github.com/prometheus/client_model v0.2.0 // indirect
 	github.com/prometheus/common v0.32.1 // indirect
 	github.com/prometheus/procfs v0.7.3 // indirect
+	github.com/rogpeppe/go-internal v1.11.0 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.2.11 // indirect
 	golang.org/x/arch v0.3.0 // indirect
-	golang.org/x/crypto v0.9.0 // indirect
+	golang.org/x/crypto v0.14.0 // indirect
-	golang.org/x/net v0.10.0 // indirect
+	golang.org/x/sys v0.13.0 // indirect
-	golang.org/x/sys v0.8.0 // indirect
+	golang.org/x/text v0.13.0 // indirect
-	golang.org/x/text v0.9.0 // indirect
 	google.golang.org/protobuf v1.30.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -148,13 +148,17 @@ github.com/google/pprof v0.0.0-20200229191704-1ebb73c60ed3/go.mod h1:ZgVRPoUq/hf
 github.com/google/pprof v0.0.0-20200430221834-fc25d7d30c6d/go.mod h1:ZgVRPoUq/hfqzAqh7sHMqb3I9Rq5C59dIz2SbBwJ4eM=
 github.com/google/pprof v0.0.0-20200708004538-1a94d8640e99/go.mod h1:ZgVRPoUq/hfqzAqh7sHMqb3I9Rq5C59dIz2SbBwJ4eM=
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
-github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
-github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
+github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
+github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
+github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
+github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
+github.com/jackc/pgx/v5 v5.4.3 h1:cxFyXhxlvAifxnkKKdlxv8XqUf59tDlYjnV5YYfsJJY=
+github.com/jackc/pgx/v5 v5.4.3/go.mod h1:Ig06C2Vu0t5qXC60W8sqIthScaEnFvojjj9dSljmHRA=
 github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
 github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
 github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
@@ -178,7 +182,9 @@ github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxv
 github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII=
 github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q=
@@ -232,6 +238,8 @@ github.com/prometheus/procfs v0.6.0/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1
 github.com/prometheus/procfs v0.7.3 h1:4jVXhlkAyzOScmCkXBTOLRLTz8EeU+eyjrwB/EPq0VU=
 github.com/prometheus/procfs v0.7.3/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
+github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
+github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
 github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
 github.com/sirupsen/logrus v1.6.0/go.mod h1:7uNnSEd1DgxDLC74fIahvMZmmYsHGZGEOFrfsX/uA88=
@@ -272,8 +280,8 @@ golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8U
 golang.org/x/crypto v0.0.0-20190605123033-f99c8df09eb5/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
-golang.org/x/crypto v0.9.0 h1:LF6fAI+IutBocDJ2OT0Q1g8plpYljMZ4+lty+dsqw3g=
+golang.org/x/crypto v0.14.0 h1:wBqGXzWJW6m1XrIKlAH0Hs1JJ7+9KBwnIO8v66Q9cHc=
-golang.org/x/crypto v0.9.0/go.mod h1:yrmDGqONDYtNj3tH8X9dzUun2m2lzPa9ngI6/RUPGR0=
+golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20190510132918-efd6b22b2522/go.mod h1:ZjyILWgesfNpC6sMxTJOJm9Kp84zZh5NQWvqDGG3Qr8=
@@ -389,8 +397,8 @@ golang.org/x/sys v0.0.0-20210603081109-ebe580a85c40/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20220114195835-da31bd327af9/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU=
+golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
-golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
@@ -398,8 +406,8 @@ golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
+golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k=
-golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
+golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
 golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
@@ -445,7 +453,6 @@ golang.org/x/tools v0.0.0-20200804011535-6c149bb5ef0d/go.mod h1:njjCfa9FT2d7l9Bc
 golang.org/x/tools v0.0.0-20200825202427-b303f430e36d/go.mod h1:njjCfa9FT2d7l9Bc6FUM5FLjQPp3cFF28FI3qnDFljA=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
@@ -527,11 +534,10 @@ google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQ
 google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng=
 google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
-gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
@@ -543,10 +549,12 @@ gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gorm.io/driver/postgres v1.5.4 h1:Iyrp9Meh3GmbSuyIAGyjkN+n9K+GHX9b9MqsTL4EJCo=
+gorm.io/driver/postgres v1.5.4/go.mod h1:Bgo89+h0CRcdA33Y6frlaHHVuTdOf87pmyzwW9C/BH0=
 gorm.io/driver/sqlite v1.5.2 h1:TpQ+/dqCY4uCigCFyrfnrJnrW9zjpelWVoEVNy5qJkc=
 gorm.io/driver/sqlite v1.5.2/go.mod h1:qxAuCol+2r6PannQDpOP1FP6ag3mKi4esLnB/jHed+4=
-gorm.io/gorm v1.25.2 h1:gs1o6Vsa+oVKG/a9ElL3XgyGfghFfkKA2SInQaCyMho=
+gorm.io/gorm v1.25.5 h1:zR9lOiiYf09VNh5Q1gphfyia1JpiClIWG9hQaxB/mls=
-gorm.io/gorm v1.25.2/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k=
+gorm.io/gorm v1.25.5/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=

main.go

@@ -4,70 +4,62 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"os"
 	"strings"
 	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/penglongli/gin-metrics/ginmetrics"
+	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"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")
 
-	// connect to database
-	db, err := gorm.Open(sqlite.Open(*dbAddr), &gorm.Config{
-		PrepareStmt:            true,
-		SkipDefaultTransaction: true,
-	})
-	if err != nil {
-		log.Fatal("Failed to connect to database")
-	}
-
 	// load all upstreams
-	upstreams := make([]OPENAI_UPSTREAM, 0)
+	config = readConfig(*configFile)
-	db.Find(&upstreams)
+	log.Println("Load upstreams number:", len(config.Upstreams))
-	log.Println("Load upstreams number:", len(upstreams))
 
-	err = initconfig(db)
+	// connect to database
-	if err != nil {
+	var db *gorm.DB
-		log.Fatal(err)
+	var err error
+	switch config.DBType {
+	case "sqlite":
+		db, err = gorm.Open(sqlite.Open(config.DBAddr), &gorm.Config{
+			PrepareStmt:            true,
+			SkipDefaultTransaction: true,
+		})
+		if err != nil {
+			log.Fatalf("Error to connect sqlite database: %s", err)
+		}
+	case "postgres":
+		db, err = gorm.Open(postgres.Open(config.DBAddr), &gorm.Config{
+			PrepareStmt:            true,
+			SkipDefaultTransaction: true,
+		})
+		if err != nil {
+			log.Fatalf("Error to connect postgres database: %s", err)
+		}
+	default:
+		log.Fatalf("Unsupported database type: '%s'", config.DBType)
 	}
 
-	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,14 +94,17 @@ func main() {
 		ctx.AbortWithStatus(200)
 	})
 
-	// get authorization config from db
-	db.Take(&authConfig, "key = ?", "authorization")
-
 	engine.POST("/v1/*any", func(c *gin.Context) {
+		hostname, err := os.Hostname()
+		if config.Hostname != "" {
+			hostname = config.Hostname
+		}
 		record := Record{
 			IP:            c.ClientIP(),
+			Hostname:      hostname,
 			CreatedAt:     time.Now(),
 			Authorization: c.Request.Header.Get("Authorization"),
+			UserAgent:     c.Request.Header.Get("User-Agent"),
 		}
 		defer func() {
 			if err := recover(); err != nil {
@@ -125,15 +120,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
 			}
 
@@ -158,5 +153,5 @@ func main() {
 		}
 	})
 
-	engine.Run(*listenAddr)
+	engine.Run(config.Address)
 }

process.go

@@ -20,9 +20,9 @@ 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.UpstreamSK = upstream.SK
 	record.Response = ""
-	record.Authorization = upstream.SK
 	// [TODO] record request body
 
 	// reverse proxy
@@ -105,6 +105,7 @@ func processRequest(c *gin.Context, upstream *OPENAI_UPSTREAM, record *Record, s
 	var contentType string
 	proxy.ModifyResponse = func(r *http.Response) error {
 		haveResponse = true
+		record.ResponseTime = time.Now().Sub(record.CreatedAt)
 		record.Status = r.StatusCode
 		if !shouldResponse && r.StatusCode != 200 {
 			log.Println("upstream return not 200 and should not response", r.StatusCode)
@@ -134,6 +135,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

@@ -1,26 +1,24 @@
 package main
 
 import (
-	"encoding/json"
-	"log"
-	"strings"
 	"time"
-
-	"github.com/google/uuid"
-	"gorm.io/gorm"
 )
 
 type Record struct {
-	ID            int64 `gorm:"primaryKey,autoIncrement"`
+	ID               int64 `gorm:"primaryKey,autoIncrement"`
-	CreatedAt     time.Time
+	Hostname         string
-	IP            string
+	UpstreamEndpoint string
-	Body          string `gorm:"serializer:json"`
+	UpstreamSK       string
-	Model         string
+	CreatedAt        time.Time
-	Response      string
+	IP               string
-	ElapsedTime   time.Duration
+	Body             string
-	Status        int
+	Model            string
-	UpstreamID    uint
+	Response         string
-	Authorization string
+	ResponseTime     time.Duration
+	ElapsedTime      time.Duration
+	Status           int
+	Authorization    string // the autorization header send by client
+	UserAgent        string
 }
 
 type StreamModeChunk struct {
@@ -52,61 +50,3 @@ type FetchModeUsage struct {
 	CompletionTokens int64 `json:"completion_tokens"`
 	TotalTokens      int64 `json:"total_tokens"`
 }
-
-func recordAssistantResponse(contentType string, db *gorm.DB, trackID uuid.UUID, body []byte, elapsedTime time.Duration) {
-	result := ""
-	// stream mode
-	if strings.HasPrefix(contentType, "text/event-stream") {
-		resp := string(body)
-		for _, line := range strings.Split(resp, "\n") {
-			chunk := StreamModeChunk{}
-			line = strings.TrimPrefix(line, "data:")
-			line = strings.TrimSpace(line)
-			if line == "" {
-				continue
-			}
-
-			err := json.Unmarshal([]byte(line), &chunk)
-			if err != nil {
-				log.Println(err)
-				continue
-			}
-
-			if len(chunk.Choices) == 0 {
-				continue
-			}
-			result += chunk.Choices[0].Delta.Content
-		}
-	} else if strings.HasPrefix(contentType, "application/json") {
-		var fetchResp FetchModeResponse
-		err := json.Unmarshal(body, &fetchResp)
-		if err != nil {
-			log.Println("Error parsing fetch response:", err)
-			return
-		}
-		if !strings.HasPrefix(fetchResp.Model, "gpt-") {
-			log.Println("Not GPT model, skip recording response:", fetchResp.Model)
-			return
-		}
-		if len(fetchResp.Choices) == 0 {
-			log.Println("Error: fetch response choice length is 0")
-			return
-		}
-		result = fetchResp.Choices[0].Message.Content
-	} else {
-		log.Println("Unknown content type", contentType)
-		return
-	}
-	log.Println("Record result:", result)
-	record := Record{}
-	if db.Find(&record, "id = ?", trackID).Error != nil {
-		log.Println("Error find request record with trackID:", trackID)
-		return
-	}
-	record.Response = result
-	record.ElapsedTime = elapsedTime
-	if db.Save(&record).Error != nil {
-		log.Println("Error to save record:", record)
-		return
-	}
-}

structure.go

@@ -1,13 +1,54 @@
 package main
 
 import (
-	"gorm.io/gorm"
+	"log"
+	"os"
+
+	"gopkg.in/yaml.v3"
 )
 
-// one openai upstream contain a pair of key and endpoint
+type Config struct {
-type OPENAI_UPSTREAM struct {
+	Address       string            `yaml:"address"`
-	gorm.Model
+	Hostname      string            `yaml:"hostname"`
-	SK       string `gorm:"index:idx_sk_endpoint,unique"` // key
+	DBType        string            `yaml:"dbtype"`
-	Endpoint string `gorm:"index:idx_sk_endpoint,unique"` // endpoint
+	DBAddr        string            `yaml:"dbaddr"`
-	Timeout  int64  // timeout in seconds
+	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)
+	}
+
+	// set default value
+	if config.Address == "" {
+		log.Println("Address not set, use default value: :8888")
+		config.Address = ":8888"
+	}
+	if config.DBType == "" {
+		log.Println("DBType not set, use default value: sqlite")
+		config.DBType = "sqlite"
+	}
+	if config.DBAddr == "" {
+		log.Println("DBAddr not set, use default value: ./db.sqlite")
+		config.DBAddr = "./db.sqlite"
+	}
+
+	return config
 }