ChatGPT/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-10-16 08:04:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: only reduce remain times when request /v1/chat/completions (close #15)

Browse Source 
BREAKING CHANGE: now remain_times is -1 doesn't mean unlimited times anymore!
This commit is contained in:
JustSong
2023-04-26 10:45:34 +08:00
parent eb8f43acb5
commit 109736cc05
6 changed files with 83 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
controller/relay.go
View File

@@ -7,16 +7,20 @@ import (
"io"
"net/http"
"one-api/common"
"one-api/model"
"strings"
)
func Relay(c *gin.Context) {
channelType := c.GetInt("channel")
tokenId := c.GetInt("token_id")
isUnlimitedTimes := c.GetBool("unlimited_times")
baseURL := common.ChannelBaseURLs[channelType]
if channelType == common.ChannelTypeCustom {
baseURL = c.GetString("base_url")
}
req, err := http.NewRequest(c.Request.Method, fmt.Sprintf("%s%s", baseURL, c.Request.URL.String()), c.Request.Body)
requestURL := c.Request.URL.String()
req, err := http.NewRequest(c.Request.Method, fmt.Sprintf("%s%s", baseURL, requestURL), c.Request.Body)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
@@ -46,7 +50,19 @@ func Relay(c *gin.Context) {
})
return
}
defer resp.Body.Close()
defer func() {
err := req.Body.Close()
if err != nil {
common.SysError("Error closing request body: " + err.Error())
}
if !isUnlimitedTimes && requestURL == "/v1/chat/completions" {
err := model.DecreaseTokenRemainTimesById(tokenId)
if err != nil {
common.SysError("Error decreasing token remain times: " + err.Error())
}
}
}()
isStream := resp.Header.Get("Content-Type") == "text/event-stream"
if isStream {
scanner := bufio.NewScanner(resp.Body)