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

58
controller/relay.go
View File

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