Skip to content
Toggle navigation
P
Projects
G
Groups
S
Snippets
Help
phsl
/
new-api
This project
Loading...
Sign in
Toggle navigation
Go to a project
Project
Repository
Issues
0
Merge Requests
0
Pipelines
Wiki
Snippets
Members
Activity
Graph
Charts
Create a new issue
Jobs
Commits
Issue Boards
Files
Commits
Branches
Tags
Contributors
Graph
Compare
Charts
Commit
7e29eaa0
authored
Nov 23, 2023
by
CaIon
Browse files
Options
Browse Files
Download
Email Patches
Plain Diff
删除relay-text中的consumeQuota变量
parent
5f48b808
Hide whitespace changes
Inline
Side-by-side
Showing
2 changed files
with
55 additions
and
60 deletions
+55
-60
controller/relay-openai.go
+21
-23
controller/relay-text.go
+34
-37
No files found.
controller/relay-openai.go
View file @
7e29eaa0
...
@@ -88,30 +88,28 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
...
@@ -88,30 +88,28 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
return
nil
,
responseText
return
nil
,
responseText
}
}
func
openaiHandler
(
c
*
gin
.
Context
,
resp
*
http
.
Response
,
consumeQuota
bool
,
promptTokens
int
,
model
string
)
(
*
OpenAIErrorWithStatusCode
,
*
Usage
)
{
func
openaiHandler
(
c
*
gin
.
Context
,
resp
*
http
.
Response
,
promptTokens
int
,
model
string
)
(
*
OpenAIErrorWithStatusCode
,
*
Usage
)
{
var
textResponse
TextResponse
var
textResponse
TextResponse
if
consumeQuota
{
responseBody
,
err
:=
io
.
ReadAll
(
resp
.
Body
)
responseBody
,
err
:=
io
.
ReadAll
(
resp
.
Body
)
if
err
!=
nil
{
if
err
!=
nil
{
return
errorWrapper
(
err
,
"read_response_body_failed"
,
http
.
StatusInternalServerError
),
nil
return
errorWrapper
(
err
,
"read_response_body_failed"
,
http
.
StatusInternalServerError
),
nil
}
}
err
=
resp
.
Body
.
Close
()
err
=
resp
.
Body
.
Close
()
if
err
!=
nil
{
if
err
!=
nil
{
return
errorWrapper
(
err
,
"close_response_body_failed"
,
http
.
StatusInternalServerError
),
nil
return
errorWrapper
(
err
,
"close_response_body_failed"
,
http
.
StatusInternalServerError
),
nil
}
}
err
=
json
.
Unmarshal
(
responseBody
,
&
textResponse
)
err
=
json
.
Unmarshal
(
responseBody
,
&
textResponse
)
if
err
!=
nil
{
if
err
!=
nil
{
return
errorWrapper
(
err
,
"unmarshal_response_body_failed"
,
http
.
StatusInternalServerError
),
nil
return
errorWrapper
(
err
,
"unmarshal_response_body_failed"
,
http
.
StatusInternalServerError
),
nil
}
}
if
textResponse
.
Error
.
Type
!=
""
{
if
textResponse
.
Error
.
Type
!=
""
{
return
&
OpenAIErrorWithStatusCode
{
return
&
OpenAIErrorWithStatusCode
{
OpenAIError
:
textResponse
.
Error
,
OpenAIError
:
textResponse
.
Error
,
StatusCode
:
resp
.
StatusCode
,
StatusCode
:
resp
.
StatusCode
,
},
nil
},
nil
}
// Reset response body
resp
.
Body
=
io
.
NopCloser
(
bytes
.
NewBuffer
(
responseBody
))
}
}
// Reset response body
resp
.
Body
=
io
.
NopCloser
(
bytes
.
NewBuffer
(
responseBody
))
// We shouldn't set the header before we parse the response body, because the parse part may fail.
// We shouldn't set the header before we parse the response body, because the parse part may fail.
// And then we will have to send an error response, but in this case, the header has already been set.
// And then we will have to send an error response, but in this case, the header has already been set.
// So the httpClient will be confused by the response.
// So the httpClient will be confused by the response.
...
@@ -120,7 +118,7 @@ func openaiHandler(c *gin.Context, resp *http.Response, consumeQuota bool, promp
...
@@ -120,7 +118,7 @@ func openaiHandler(c *gin.Context, resp *http.Response, consumeQuota bool, promp
c
.
Writer
.
Header
()
.
Set
(
k
,
v
[
0
])
c
.
Writer
.
Header
()
.
Set
(
k
,
v
[
0
])
}
}
c
.
Writer
.
WriteHeader
(
resp
.
StatusCode
)
c
.
Writer
.
WriteHeader
(
resp
.
StatusCode
)
_
,
err
:
=
io
.
Copy
(
c
.
Writer
,
resp
.
Body
)
_
,
err
=
io
.
Copy
(
c
.
Writer
,
resp
.
Body
)
if
err
!=
nil
{
if
err
!=
nil
{
return
errorWrapper
(
err
,
"copy_response_body_failed"
,
http
.
StatusInternalServerError
),
nil
return
errorWrapper
(
err
,
"copy_response_body_failed"
,
http
.
StatusInternalServerError
),
nil
}
}
...
...
controller/relay-text.go
View file @
7e29eaa0
...
@@ -50,10 +50,9 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
...
@@ -50,10 +50,9 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
channelId
:=
c
.
GetInt
(
"channel_id"
)
channelId
:=
c
.
GetInt
(
"channel_id"
)
tokenId
:=
c
.
GetInt
(
"token_id"
)
tokenId
:=
c
.
GetInt
(
"token_id"
)
userId
:=
c
.
GetInt
(
"id"
)
userId
:=
c
.
GetInt
(
"id"
)
consumeQuota
:=
c
.
GetBool
(
"consume_quota"
)
group
:=
c
.
GetString
(
"group"
)
group
:=
c
.
GetString
(
"group"
)
var
textRequest
GeneralOpenAIRequest
var
textRequest
GeneralOpenAIRequest
if
c
onsumeQuota
||
c
hannelType
==
common
.
ChannelTypeAzure
||
channelType
==
common
.
ChannelTypePaLM
{
if
channelType
==
common
.
ChannelTypeAzure
||
channelType
==
common
.
ChannelTypePaLM
{
err
:=
common
.
UnmarshalBodyReusable
(
c
,
&
textRequest
)
err
:=
common
.
UnmarshalBodyReusable
(
c
,
&
textRequest
)
if
err
!=
nil
{
if
err
!=
nil
{
return
errorWrapper
(
err
,
"bind_request_body_failed"
,
http
.
StatusBadRequest
)
return
errorWrapper
(
err
,
"bind_request_body_failed"
,
http
.
StatusBadRequest
)
...
@@ -236,7 +235,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
...
@@ -236,7 +235,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
preConsumedQuota
=
0
preConsumedQuota
=
0
//common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", userId, userQuota))
//common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", userId, userQuota))
}
}
if
consumeQuota
&&
preConsumedQuota
>
0
{
if
preConsumedQuota
>
0
{
userQuota
,
err
=
model
.
PreConsumeTokenQuota
(
tokenId
,
preConsumedQuota
)
userQuota
,
err
=
model
.
PreConsumeTokenQuota
(
tokenId
,
preConsumedQuota
)
if
err
!=
nil
{
if
err
!=
nil
{
return
errorWrapper
(
err
,
"pre_consume_token_quota_failed"
,
http
.
StatusForbidden
)
return
errorWrapper
(
err
,
"pre_consume_token_quota_failed"
,
http
.
StatusForbidden
)
...
@@ -420,41 +419,39 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
...
@@ -420,41 +419,39 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
defer
func
(
ctx
context
.
Context
)
{
defer
func
(
ctx
context
.
Context
)
{
// c.Writer.Flush()
// c.Writer.Flush()
go
func
()
{
go
func
()
{
if
consumeQuota
{
quota
:=
0
quota
:=
0
completionRatio
:=
common
.
GetCompletionRatio
(
textRequest
.
Model
)
completionRatio
:=
common
.
GetCompletionRatio
(
textRequest
.
Model
)
promptTokens
=
textResponse
.
Usage
.
PromptTokens
promptTokens
=
textResponse
.
Usage
.
PromptTokens
completionTokens
=
textResponse
.
Usage
.
CompletionTokens
completionTokens
=
textResponse
.
Usage
.
CompletionTokens
quota
=
promptTokens
+
int
(
float64
(
completionTokens
)
*
completionRatio
)
quota
=
promptTokens
+
int
(
float64
(
completionTokens
)
*
completionRatio
)
quota
=
int
(
float64
(
quota
)
*
ratio
)
quota
=
int
(
float64
(
quota
)
*
ratio
)
if
ratio
!=
0
&&
quota
<=
0
{
if
ratio
!=
0
&&
quota
<=
0
{
quota
=
1
quota
=
1
}
}
totalTokens
:=
promptTokens
+
completionTokens
totalTokens
:=
promptTokens
+
completionTokens
if
totalTokens
==
0
{
if
totalTokens
==
0
{
// in this case, must be some error happened
// in this case, must be some error happened
// we cannot just return, because we may have to return the pre-consumed quota
// we cannot just return, because we may have to return the pre-consumed quota
quota
=
0
quota
=
0
}
}
quotaDelta
:=
quota
-
preConsumedQuota
quotaDelta
:=
quota
-
preConsumedQuota
err
:=
model
.
PostConsumeTokenQuota
(
tokenId
,
userQuota
,
quotaDelta
,
preConsumedQuota
,
true
)
err
:=
model
.
PostConsumeTokenQuota
(
tokenId
,
userQuota
,
quotaDelta
,
preConsumedQuota
,
true
)
if
err
!=
nil
{
if
err
!=
nil
{
common
.
LogError
(
ctx
,
"error consuming token remain quota: "
+
err
.
Error
())
common
.
LogError
(
ctx
,
"error consuming token remain quota: "
+
err
.
Error
())
}
err
=
model
.
CacheUpdateUserQuota
(
userId
)
if
err
!=
nil
{
common
.
LogError
(
ctx
,
"error update user quota cache: "
+
err
.
Error
())
}
// record all the consume log even if quota is 0
logContent
:=
fmt
.
Sprintf
(
"模型倍率 %.2f,分组倍率 %.2f"
,
modelRatio
,
groupRatio
)
model
.
RecordConsumeLog
(
ctx
,
userId
,
channelId
,
promptTokens
,
completionTokens
,
textRequest
.
Model
,
tokenName
,
quota
,
logContent
,
tokenId
)
model
.
UpdateUserUsedQuotaAndRequestCount
(
userId
,
quota
)
model
.
UpdateChannelUsedQuota
(
channelId
,
quota
)
//if quota != 0 {
//
//}
}
}
err
=
model
.
CacheUpdateUserQuota
(
userId
)
if
err
!=
nil
{
common
.
LogError
(
ctx
,
"error update user quota cache: "
+
err
.
Error
())
}
// record all the consume log even if quota is 0
logContent
:=
fmt
.
Sprintf
(
"模型倍率 %.2f,分组倍率 %.2f"
,
modelRatio
,
groupRatio
)
model
.
RecordConsumeLog
(
ctx
,
userId
,
channelId
,
promptTokens
,
completionTokens
,
textRequest
.
Model
,
tokenName
,
quota
,
logContent
,
tokenId
)
model
.
UpdateUserUsedQuotaAndRequestCount
(
userId
,
quota
)
model
.
UpdateChannelUsedQuota
(
channelId
,
quota
)
//if quota != 0 {
//
//}
}()
}()
}(
c
.
Request
.
Context
())
}(
c
.
Request
.
Context
())
switch
apiType
{
switch
apiType
{
...
@@ -468,7 +465,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
...
@@ -468,7 +465,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
textResponse
.
Usage
.
CompletionTokens
=
countTokenText
(
responseText
,
textRequest
.
Model
)
textResponse
.
Usage
.
CompletionTokens
=
countTokenText
(
responseText
,
textRequest
.
Model
)
return
nil
return
nil
}
else
{
}
else
{
err
,
usage
:=
openaiHandler
(
c
,
resp
,
consumeQuota
,
promptTokens
,
textRequest
.
Model
)
err
,
usage
:=
openaiHandler
(
c
,
resp
,
promptTokens
,
textRequest
.
Model
)
if
err
!=
nil
{
if
err
!=
nil
{
return
err
return
err
}
}
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment