upstreams.yaml

.gitignore

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

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)
+	upstreams := readUpstreams(*upstreamsFile)
-	db.Find(&upstreams)
 	log.Println("Load upstreams number:", len(upstreams))
 
 	err = initconfig(db)

process.go

@@ -20,7 +20,7 @@ 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 = ""
 	// [TODO] record request body
 

record.go

@@ -11,17 +11,18 @@ import (
 )
 
 type Record struct {
-	ID            int64 `gorm:"primaryKey,autoIncrement"`
+	ID               int64 `gorm:"primaryKey,autoIncrement"`
-	CreatedAt     time.Time
+	UpstreamEndpoint string
-	IP            string
+	CreatedAt        time.Time
-	Body          string `gorm:"serializer:json"`
+	IP               string
-	Model         string
+	Body             string `gorm:"serializer:json"`
-	Response      string
+	Model            string
-	ResponseTime  time.Duration
+	Response         string
-	ElapsedTime   time.Duration
+	ResponseTime     time.Duration
-	Status        int
+	ElapsedTime      time.Duration
-	UpstreamID    uint
+	Status           int
-	Authorization string
+	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 `yaml:"sk"`
-	SK       string `gorm:"index:idx_sk_endpoint,unique"` // key
+	Endpoint string `yaml:"endpoint"`
-	Endpoint string `gorm:"index:idx_sk_endpoint,unique"` // endpoint
+	Timeout  int64  `yaml:"timeout"`
-	Timeout  int64  // timeout in seconds
+}
+
+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
 }