upstreams.yaml

.gitignore

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

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,6 +15,7 @@ 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")
@@ -35,8 +36,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)

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
 }