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

fix: fix Azure channel not working in stream mode (#57)

Browse Source
This commit is contained in:
JustSong 2023-05-14 09:39:42 +08:00
parent e123c66bc7
commit 2edd52e851
1 changed files with 9 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
controller/relay.go
View File

@ -94,10 +94,12 @@ func relayHelper(c *gin.Context) error {
if channelType == common.ChannelTypeAzure { if channelType == common.ChannelTypeAzure {
// https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api // https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api
query := c.Request.URL.Query() query := c.Request.URL.Query()
if query.Get("api-version") == "" { apiVersion := query.Get("api-version")
apiVersion := c.GetString("api_version") if apiVersion == "" {
requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, apiVersion) apiVersion = c.GetString("api_version")
} }
requestURL := strings.Split(requestURL, "?")[0]
requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, apiVersion)
baseURL = c.GetString("base_url") baseURL = c.GetString("base_url")
task := strings.TrimPrefix(requestURL, "/v1/") task := strings.TrimPrefix(requestURL, "/v1/")
model_ := textRequest.Model model_ := textRequest.Model
@ -186,7 +188,7 @@ func relayHelper(c *gin.Context) error {
data := scanner.Text() data := scanner.Text()
dataChan <- data dataChan <- data
data = data[6:] data = data[6:]
if data != "[DONE]" { if !strings.HasPrefix(data, "[DONE]") {
var streamResponse StreamResponse var streamResponse StreamResponse
err = json.Unmarshal([]byte(data), &streamResponse) err = json.Unmarshal([]byte(data), &streamResponse)
if err != nil { if err != nil {
@ -207,6 +209,9 @@ func relayHelper(c *gin.Context) error {
c.Stream(func(w io.Writer) bool { c.Stream(func(w io.Writer) bool {
select { select {
case data := <-dataChan: case data := <-dataChan:
if strings.HasPrefix(data, "data: [DONE]") {
data = data[:len(data)-1]
}
c.Render(-1, common.CustomEvent{Data: data}) c.Render(-1, common.CustomEvent{Data: data})
return true return true
case <-stopChan: case <-stopChan: