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

fix: wait goroutine

Browse Source
This commit is contained in:
Xyfacai 2023-11-26 14:02:08 +08:00
parent 622ff85e94
commit 6d2210ad5c
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
controller/relay-openai.go
View File

@ -9,6 +9,7 @@ import (
"net/http" "net/http"
"one-api/common" "one-api/common"
"strings" "strings"
"sync"
) )
func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*OpenAIErrorWithStatusCode, string) { func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*OpenAIErrorWithStatusCode, string) {
@ -28,7 +29,9 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
}) })
dataChan := make(chan string) dataChan := make(chan string)
stopChan := make(chan bool) stopChan := make(chan bool)
var wg sync.WaitGroup
go func() { go func() {
wg.Add(1)
var streamItems []string var streamItems []string
for scanner.Scan() { for scanner.Scan() {
data := scanner.Text() data := scanner.Text()
@ -71,6 +74,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
} }
} }
} }
wg.Done()
stopChan <- true stopChan <- true
}() }()
setEventStreamHeaders(c) setEventStreamHeaders(c)
@ -92,6 +96,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
if err != nil { if err != nil {
return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), "" return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
} }
wg.Wait()
return nil, responseTextBuilder.String() return nil, responseTextBuilder.String()
} }
@ -117,7 +122,6 @@ func openaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model
} }
// Reset response body // Reset response body
resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) 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.