Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: relay bug fix

Browse Source
This commit is contained in:
JustSong 2023-04-28 18:16:59 +08:00
parent 480e789cd8
commit aea6c859e7
1 changed files with 31 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
controller/relay.go
View File

@ -64,6 +64,7 @@ func relayHelper(c *gin.Context) error {
if channelType == common.ChannelTypeCustom {
baseURL = c.GetString("base_url")
}
if consumeQuota {
requestBody, err := io.ReadAll(c.Request.Body)
if err != nil {
return err
@ -79,6 +80,7 @@ func relayHelper(c *gin.Context) error {
}
// Reset request body
c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
}
requestURL := c.Request.URL.String()
req, err := http.NewRequest(c.Request.Method, fmt.Sprintf("%s%s", baseURL, requestURL), c.Request.Body)
if err != nil {
@ -181,6 +183,7 @@ func relayHelper(c *gin.Context) error {
for k, v := range resp.Header {
c.Writer.Header().Set(k, v[0])
}
if consumeQuota {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return err
@ -195,6 +198,7 @@ func relayHelper(c *gin.Context) error {
}
// Reset response body
resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
}
_, err = io.Copy(c.Writer, resp.Body)
if err != nil {
return err