upstream asis

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`。
 
 ## 上游管理
@@ -60,7 +143,7 @@
   }
   ```
 
-### 删除指定ID的上游
+### 删除指定 ID 的上游
 
 - URL: `/admin/upstreams/:id`
 - 方法: DELETE
@@ -77,7 +160,7 @@
   }
   ```
 
-### 更新指定ID的上游
+### 更新指定 ID 的上游
 
 - URL: `/admin/upstreams/:id`
 - 方法: PUT

auth.go

@@ -20,7 +20,8 @@ func handleAuth(c *gin.Context) error {
 
 	authorization = strings.Trim(authorization[len("Bearer"):], " ")
 	log.Println("Received authorization", authorization)
-	if authorization != authConfig.Value {
+
+	if authConfig.Value != "asis" && authorization != authConfig.Value {
 		err = errors.New("wrong authorization header")
 		c.AbortWithError(403, err)
 		return err

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{}

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

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
@@ -10,7 +11,6 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
-	"strconv"
 	"strings"
 	"time"
 
@@ -46,7 +46,7 @@ func main() {
 	}
 
 	db.AutoMigrate(&OPENAI_UPSTREAM{})
-	db.AutoMigrate(&RequestRecord{})
+	db.AutoMigrate(&Record{})
 	log.Println("Auto migrate database done")
 
 	if *addMode {
@@ -89,10 +89,30 @@ func main() {
 		})
 	})
 
+	// CORS handler
+	engine.OPTIONS("/v1/*any", func(ctx *gin.Context) {
+		header := ctx.Writer.Header()
+		header.Set("Access-Control-Allow-Origin", "*")
+		header.Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, PATCH")
+		header.Set("Access-Control-Allow-Headers", "Origin, Authorization, Content-Type")
+		ctx.AbortWithStatus(200)
+	})
+
 	// get authorization config from db
 	db.Take(&authConfig, "key = ?", "authorization")
 
 	engine.POST("/v1/*any", func(c *gin.Context) {
+		record := Record{
+			IP:        c.ClientIP(),
+			CreatedAt: time.Now(),
+		}
+		defer func() {
+			if err := recover(); err != nil {
+				log.Println("Error:", err)
+				c.AbortWithError(500, fmt.Errorf("%s", err))
+			}
+		}()
+
 		// check authorization header
 		if !*noauth {
 			if handleAuth(c) != nil {
@@ -133,6 +153,8 @@ func main() {
 			return
 		}
 
+		record.UpstreamID = upstream.ID
+
 		// reverse proxy
 		remote, err := url.Parse(upstream.Endpoint)
 		if err != nil {
@@ -153,7 +175,7 @@ func main() {
 			}
 
 			// record chat message from user
-			go recordUserMessage(c, db, body)
+			record.Body = string(body)
 
 			out.Body = io.NopCloser(bytes.NewReader(body))
 
@@ -163,22 +185,36 @@ func main() {
 			out.URL.Path = in.URL.Path
 			out.Header = http.Header{}
 			out.Header.Set("Host", remote.Host)
-			out.Header.Set("Authorization", "Bearer "+upstream.SK)
+			if upstream.SK == "asis" {
+				out.Header.Set("Authorization", c.Request.Header.Get("Authorization"))
+			} else {
+				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 {
+			record.Status = r.StatusCode
+			r.Header.Del("Access-Control-Allow-Origin")
+			r.Header.Del("Access-Control-Allow-Methods")
+			r.Header.Del("Access-Control-Allow-Headers")
+			r.Header.Set("Access-Control-Allow-Origin", "*")
+			r.Header.Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, PATCH")
+			r.Header.Set("Access-Control-Allow-Headers", "Origin, Authorization, Content-Type")
+
 			if r.StatusCode != 200 {
 				body, err := io.ReadAll(r.Body)
 				if err != nil {
-					return errors.New("failed to read response from upstream " + err.Error())
+					record.Response = "failed to read response from upstream " + err.Error()
+					return errors.New(record.Response)
 				}
-				return fmt.Errorf("upstream return '%s' with '%s'", r.Status, string(body))
+				record.Response = fmt.Sprintf("upstream return '%s' with '%s'", r.Status, string(body))
+				return fmt.Errorf(record.Response)
 			}
 			// count success
-			go db.Model(&upstream).Updates(map[string]interface{}{
+			r.Body = io.NopCloser(io.TeeReader(r.Body, &buf))
-				"success_count":          gorm.Expr("success_count + ?", 1),
+			contentType = r.Header.Get("content-type")
-				"last_call_success_time": time.Now(),
-			})
 			return nil
 		}
 		proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
@@ -198,123 +234,86 @@ func main() {
 			}
 			content := fmt.Sprintf("[%s] OpenAI 转发出错 ID: %d... 密钥: [%s] 上游: [%s] 错误: %s\n---\n%s",
 				c.ClientIP(),
-				upstream.ID, upstream.SK[:10], upstream.Endpoint, err.Error(),
+				upstream.ID, upstream.SK, upstream.Endpoint, err.Error(),
 				strings.Join(upstreamDescriptions, "\n"),
 			)
 			go sendMatrixMessage(content)
 			if err.Error() != "context canceled" && r.Response.StatusCode != 400 {
-				// count failed
-				go db.Model(&upstream).Update("failed_count", gorm.Expr("failed_count + ?", 1))
 				go sendFeishuMessage(content)
 			}
 
 			log.Println("response is", r.Response)
 		}
-		proxy.ServeHTTP(c.Writer, c.Request)
+
+		func() {
+			defer func() {
+				if err := recover(); err != nil {
+					log.Println("Panic recover :", err)
+				}
+			}()
+			proxy.ServeHTTP(c.Writer, c.Request)
+		}()
+
+		resp, err := io.ReadAll(io.NopCloser(&buf))
+		if err != nil {
+			record.Response = "failed to read response from upstream " + err.Error()
+			log.Println(record.Response)
+		} else {
+
+			// record response
+			// stream mode
+			if strings.HasPrefix(contentType, "text/event-stream") {
+				for _, line := range strings.Split(string(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
+					}
+					record.Response += chunk.Choices[0].Delta.Content
+				}
+			} else if strings.HasPrefix(contentType, "application/json") {
+				var fetchResp FetchModeResponse
+				err := json.Unmarshal(resp, &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
+				}
+				record.Response = fetchResp.Choices[0].Message.Content
+			} else {
+				log.Println("Unknown content type", contentType)
+				return
+			}
+		}
+
+		if len(record.Body) > 1024*512 {
+			record.Body = ""
+		}
+
+		log.Println("Record result:", record.Response)
+		record.ElapsedTime = time.Now().Sub(record.CreatedAt)
+		if db.Create(&record).Error != nil {
+			log.Println("Error to save record:", record)
+		}
 	})
 
-	// ---------------------------------
-	// admin APIs
-	engine.POST("/admin/login", func(c *gin.Context) {
-		// check authorization headers
-		if handleAuth(c) != nil {
-			return
-		}
-		c.JSON(200, gin.H{
-			"message": "success",
-		})
-	})
-	engine.GET("/admin/upstreams", func(c *gin.Context) {
-		// check authorization headers
-		if handleAuth(c) != nil {
-			return
-		}
-		upstreams := make([]OPENAI_UPSTREAM, 0)
-		db.Find(&upstreams)
-		c.JSON(200, upstreams)
-	})
-	engine.POST("/admin/upstreams", func(c *gin.Context) {
-		// check authorization headers
-		if handleAuth(c) != nil {
-			return
-		}
-		newUpstream := OPENAI_UPSTREAM{}
-		err := c.BindJSON(&newUpstream)
-		if err != nil {
-			c.AbortWithError(502, errors.New("can't parse OPENAI_UPSTREAM object"))
-			return
-		}
-		if newUpstream.SK == "" || newUpstream.Endpoint == "" {
-			c.AbortWithError(403, errors.New("can't create new OPENAI_UPSTREAM with empty sk or endpoint"))
-			return
-		}
-		log.Println("Saveing new OPENAI_UPSTREAM", newUpstream)
-		err = db.Create(&newUpstream).Error
-		if err != nil {
-			c.AbortWithError(403, err)
-			return
-		}
-	})
-	engine.DELETE("/admin/upstreams/:id", func(ctx *gin.Context) {
-		// check authorization headers
-		if handleAuth(ctx) != nil {
-			return
-		}
-		id, err := strconv.Atoi(ctx.Param("id"))
-		if err != nil {
-			ctx.AbortWithError(502, err)
-			return
-		}
-		upstream := OPENAI_UPSTREAM{}
-		upstream.ID = uint(id)
-		db.Delete(&upstream)
-		ctx.JSON(200, gin.H{
-			"message": "success",
-		})
-	})
-	engine.PUT("/admin/upstreams/:id", func(c *gin.Context) {
-		// check authorization headers
-		if handleAuth(c) != nil {
-			return
-		}
-		upstream := OPENAI_UPSTREAM{}
-		err := c.BindJSON(&upstream)
-		if err != nil {
-			c.AbortWithError(502, errors.New("can't parse OPENAI_UPSTREAM object"))
-			return
-		}
-		if upstream.SK == "" || upstream.Endpoint == "" {
-			c.AbortWithError(403, errors.New("can't create new OPENAI_UPSTREAM with empty sk or endpoint"))
-			return
-		}
-		id, err := strconv.Atoi(c.Param("id"))
-		if err != nil {
-			c.AbortWithError(502, err)
-			return
-		}
-		upstream.ID = uint(id)
-		log.Println("Saveing new OPENAI_UPSTREAM", upstream)
-		err = db.Create(&upstream).Error
-		if err != nil {
-			c.AbortWithError(403, err)
-			return
-		}
-		c.JSON(200, gin.H{
-			"message": "success",
-		})
-	})
-	engine.GET("/admin/request_records", func(c *gin.Context) {
-		// check authorization headers
-		if handleAuth(c) != nil {
-			return
-		}
-		requestRecords := []RequestRecord{}
-		err := db.Order("id desc").Limit(100).Find(&requestRecords).Error
-		if err != nil {
-			c.AbortWithError(502, err)
-			return
-		}
-		c.JSON(200, requestRecords)
-	})
 	engine.Run(*listenAddr)
 }

record.go

@@ -1,24 +1,110 @@
 package main
 
 import (
+	"encoding/json"
 	"log"
+	"strings"
+	"time"
 
-	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
 	"gorm.io/gorm"
 )
 
-type RequestRecord struct {
+type Record struct {
-	gorm.Model
+	ID          int64 `gorm:"primaryKey,autoIncrement"`
-	Body string
+	CreatedAt   time.Time
+	IP          string
+	Body        string `gorm:"serializer:json"`
+	Response    string
+	ElapsedTime time.Duration
+	Status      int
+	UpstreamID  uint
 }
 
-func recordUserMessage(c *gin.Context, db *gorm.DB, body []byte) {
+type StreamModeChunk struct {
-	bodyStr := string(body)
+	Choices []StreamModeChunkChoice `json:"choices"`
-	requestRecord := RequestRecord{
+}
-		Body: bodyStr,
+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
 	}
-	err := db.Create(&requestRecord).Error
+	log.Println("Record result:", result)
-	if err != nil {
+	record := Record{}
-		log.Println("Error record request:", err)
+	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
 	}
 }