部署文档

README.md

@@ -1,17 +1,100 @@
-# API文档
+# API 文档
 
-本文档提供了使用该负载君恩和能够API的方法和端点的详细说明。
+本文档详细介绍了如何使用负载均衡和能力 API 的方法和端点。
 
+## 部署方法
+
+### 编译
+
+以下是编译和运行该负载均衡 API 的步骤：
+
+1. 首先，确保您已经安装了 golang 和 gcc。
+
+2. 克隆本仓库到您的本地机器上。
+
+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 进行相关操作。
+
+### 运行
+
+以下是运行命令的用法：
+
+```
+Usage of ./openai-api-route:
+  -add
+        添加一个 OpenAI 上游
+  -addr string
+        监听地址（默认为 ":8888"）
+  -database string
+        数据库地址（默认为 "./db.sqlite"）
+  -endpoint string
+        OpenAI API 基地址（默认为 "https://api.openai.com/v1"）
+  -list
+        列出所有上游
+  -noauth
+        不检查传入的授权头
+  -sk string
+        OpenAI API 密钥（sk-xxxxx）
+```
+
+您可以直接运行 `./openai-api-route` 命令，如果数据库不存在，系统会自动创建。
+
+### 上游管理
+
+您可以使用以下命令添加一个上游：
+
+```bash
+./openai-api-route -add -sk sk-xxxxx -endpoint https://api.openai.com/v1
+```
+
+您也可以使用 `/admin/upstreams` 的 HTTP 接口进行控制。
+
+另外，您还可以直接编辑数据库中的 `openai_upstreams` 表。
 ## 身份验证
 
 ### 身份验证中间件流程
 
 1. 从请求头中获取`Authorization`字段的值。
 2. 检查`Authorization`字段的值是否以`"Bearer"`开头。
-   - 如果不是，则返回错误信息："authorization header should start with 'Bearer'"（HTTP状态码403）。
+   - 如果不是，则返回错误信息："authorization header should start with 'Bearer'"（HTTP 状态码 403）。
 3. 去除`Authorization`字段值开头的`"Bearer"`和前后的空格。
 4. 将剩余的值与预先设置的身份验证配置进行比较。
-   - 如果不匹配，则返回错误信息："wrong authorization header"（HTTP状态码403）。
+   - 如果不匹配，则返回错误信息："wrong authorization header"（HTTP 状态码 403）。
 5. 如果身份验证通过，则返回`nil`。
 
 ## 上游管理
@@ -22,7 +105,7 @@
 - 方法: GET
 - 权限要求: 需要进行身份验证
 - 返回数据类型: JSON
-- 请求示例: 
+- 请求示例:
   ```bash
   curl -X GET -H "Authorization: Bearer access_token" http://localhost:8080/admin/upstreams
   ```
@@ -60,7 +143,7 @@
   }
   ```
 
-### 删除指定ID的上游
+### 删除指定 ID 的上游
 
 - URL: `/admin/upstreams/:id`
 - 方法: DELETE
@@ -77,7 +160,7 @@
   }
   ```
 
-### 更新指定ID的上游
+### 更新指定 ID 的上游
 
 - URL: `/admin/upstreams/:id`
 - 方法: PUT
@@ -93,4 +176,4 @@
   {
     "message": "success"
   }
-  ```
+  ```

config.go

@@ -26,7 +26,7 @@ func initconfig(db *gorm.DB) error {
 	// config list and their default values
 	configs := make(map[string]string)
 	configs["authorization"] = "woshimima"
-	configs["policy"] = "random"
+	configs["policy"] = "main"
 
 	for key, value := range configs {
 		kv := ConfigKV{}

cors.go

@@ -0,0 +1,18 @@
+package main
+
+import (
+	"github.com/gin-gonic/gin"
+)
+
+// Middleware function to handle CORS requests
+func handleCORS(c *gin.Context) {
+	c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
+	c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, PATCH")
+	c.Writer.Header().Set("Access-Control-Allow-Headers", "Origin, Authorization, Content-Type")
+
+	if c.Request.Method == "OPTIONS" {
+		c.AbortWithStatus(200)
+		return
+	}
+
+}

go.mod

@@ -17,6 +17,7 @@ require (
 	github.com/go-playground/universal-translator v0.18.1 // indirect
 	github.com/go-playground/validator/v10 v10.14.0 // indirect
 	github.com/goccy/go-json v0.10.2 // indirect
+	github.com/google/uuid v1.3.0 // 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

go.sum

@@ -26,6 +26,8 @@ github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaS
 github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+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/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=

main.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
 )
@@ -46,7 +47,7 @@ func main() {
 	}
 
 	db.AutoMigrate(&OPENAI_UPSTREAM{})
-	db.AutoMigrate(&RequestRecord{})
+	db.AutoMigrate(&Record{})
 	log.Println("Auto migrate database done")
 
 	if *addMode {
@@ -89,10 +90,15 @@ func main() {
 		})
 	})
 
+	// CORS handler
+	engine.Use(handleCORS)
+
 	// get authorization config from db
 	db.Take(&authConfig, "key = ?", "authorization")
 
 	engine.POST("/v1/*any", func(c *gin.Context) {
+		begin := time.Now()
+		trackID := uuid.New()
 		// check authorization header
 		if !*noauth {
 			if handleAuth(c) != nil {
@@ -153,7 +159,7 @@ func main() {
 			}
 
 			// record chat message from user
-			go recordUserMessage(c, db, body)
+			go recordUserMessage(c, db, trackID, body)
 
 			out.Body = io.NopCloser(bytes.NewReader(body))
 
@@ -166,6 +172,8 @@ func main() {
 			out.Header.Set("Authorization", "Bearer "+upstream.SK)
 			out.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
 		}
+		var buf bytes.Buffer
+		var contentType string
 		proxy.ModifyResponse = func(r *http.Response) error {
 			if r.StatusCode != 200 {
 				body, err := io.ReadAll(r.Body)
@@ -179,6 +187,8 @@ func main() {
 				"success_count":          gorm.Expr("success_count + ?", 1),
 				"last_call_success_time": time.Now(),
 			})
+			r.Body = io.NopCloser(io.TeeReader(r.Body, &buf))
+			contentType = r.Header.Get("content-type")
 			return nil
 		}
 		proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
@@ -211,6 +221,11 @@ func main() {
 			log.Println("response is", r.Response)
 		}
 		proxy.ServeHTTP(c.Writer, c.Request)
+		resp, err := io.ReadAll(io.NopCloser(&buf))
+		if err != nil {
+			log.Println("Failed to read from response tee buffer", err)
+		}
+		go recordAssistantResponse(contentType, db, trackID, resp, time.Now().Sub(begin))
 	})
 
 	// ---------------------------------
@@ -308,7 +323,7 @@ func main() {
 		if handleAuth(c) != nil {
 			return
 		}
-		requestRecords := []RequestRecord{}
+		requestRecords := []Record{}
 		err := db.Order("id desc").Limit(100).Find(&requestRecords).Error
 		if err != nil {
 			c.AbortWithError(502, err)

record.go

@@ -1,24 +1,122 @@
 package main
 
 import (
+	"encoding/json"
 	"log"
+	"strings"
+	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
 	"gorm.io/gorm"
 )
 
-type RequestRecord struct {
-	gorm.Model
-	Body string
+type Record struct {
+	ID          uuid.UUID `gorm:"type:uuid"`
+	CreatedAt   time.Time
+	IP          string
+	Body        string
+	Response    string
+	ElapsedTime time.Duration
 }
 
-func recordUserMessage(c *gin.Context, db *gorm.DB, body []byte) {
+func recordUserMessage(c *gin.Context, db *gorm.DB, trackID uuid.UUID, body []byte) {
 	bodyStr := string(body)
-	requestRecord := RequestRecord{
+	requestRecord := Record{
 		Body: bodyStr,
+		ID:   trackID,
+		IP:   c.ClientIP(),
 	}
 	err := db.Create(&requestRecord).Error
 	if err != nil {
 		log.Println("Error record request:", err)
 	}
 }
+
+type StreamModeChunk struct {
+	Choices []StreamModeChunkChoice `json:"choices"`
+}
+type StreamModeChunkChoice struct {
+	Delta        StreamModeDelta `json:"delta"`
+	FinishReason string          `json:"finish_reason"`
+}
+type StreamModeDelta struct {
+	Content string `json:"content"`
+}
+
+type FetchModeResponse struct {
+	Model   string            `json:"model"`
+	Choices []FetchModeChoice `json:"choices"`
+	Usage   FetchModeUsage    `json:"usage"`
+}
+type FetchModeChoice struct {
+	Message      FetchModeMessage `json:"message"`
+	FinishReason string           `json:"finish_reason"`
+}
+type FetchModeMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+type FetchModeUsage struct {
+	PromptTokens     int64 `json:"prompt_tokens"`
+	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
+	}
+}