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

refactor: improve relay's implementation

Browse Source
This commit is contained in:
JustSong 2023-04-28 17:11:57 +08:00
parent 053bb85a1c
commit 23ec541ba6
1 changed files with 36 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

107
controller/relay.go
View File

@ -45,6 +45,18 @@ type StreamResponse struct {
}
func Relay(c *gin.Context) {
err := relayHelper(c)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
}
}
func relayHelper(c *gin.Context) error {
channelType := c.GetInt("channel")
tokenId := c.GetInt("token_id")
consumeQuota := c.GetBool("consume_quota")
@ -54,47 +66,27 @@ func Relay(c *gin.Context) {
}
requestBody, err := io.ReadAll(c.Request.Body)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
err = c.Request.Body.Close()
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
var textRequest TextRequest
err = json.Unmarshal(requestBody, &textRequest)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
// 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 {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
err = c.Request.Body.Close()
if err != nil {
return err
}
req.Header.Set("Authorization", c.Request.Header.Get("Authorization"))
req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
@ -103,23 +95,11 @@ func Relay(c *gin.Context) {
client := &http.Client{}
resp, err := client.Do(req)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
err = req.Body.Close()
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
var textResponse TextResponse
@ -192,53 +172,38 @@ func Relay(c *gin.Context) {
return false
}
})
return
err = resp.Body.Close()
if err != nil {
return err
}
return nil
} else {
for k, v := range resp.Header {
c.Writer.Header().Set(k, v[0])
}
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
err = resp.Body.Close()
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
err = json.Unmarshal(responseBody, &textResponse)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
// Reset response body
resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
_, err = io.Copy(c.Writer, resp.Body)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"error": gin.H{
"message": err.Error(),
"type": "one_api_error",
},
})
return
return err
}
err = resp.Body.Close()
if err != nil {
return err
}
return nil
}
}