From fb19d8a353c6b4ddbe10a686f3fd07721fca2360 Mon Sep 17 00:00:00 2001
From: heimoshuiyu <heimoshuiyu@gmail.com>
Date: Mon, 27 Nov 2023 17:19:08 +0800
Subject: [PATCH] upstreams.yaml

---
 .gitignore   |  1 +
 main.go      |  4 ++--
 process.go   |  2 +-
 record.go    | 23 ++++++++++++-----------
 structure.go | 30 +++++++++++++++++++++++++-----
 5 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/.gitignore b/.gitignore
index a838100..caa5548 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,3 @@
 openai-api-route
 db.sqlite
+/upstreams.yaml
diff --git a/main.go b/main.go
index 86f8f5f..c6c9eb0 100644
--- a/main.go
+++ b/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)
diff --git a/process.go b/process.go
index e9901c3..f53dd85 100644
--- a/process.go
+++ b/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
 
diff --git a/record.go b/record.go
index fa0e4ef..0c31ada 100644
--- a/record.go
+++ b/record.go
@@ -11,17 +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
-	ResponseTime  time.Duration
-	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 {
diff --git a/structure.go b/structure.go
index c47c116..5ff7b3b 100644
--- a/structure.go
+++ b/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
 }