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

fix: ignore empty choice response for azure (close #1324)

Browse Source
This commit is contained in:
JustSong 2024-04-21 16:22:28 +08:00
parent b2679cca65
commit 541182102e
1 changed files with 42 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
relay/adaptor/openai/main.go
View File

@ -15,6 +15,12 @@ import (
"strings" "strings"
) )
const (
dataPrefix = "data: "
done = "[DONE]"
dataPrefixLength = len(dataPrefix)
)
func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.ErrorWithStatusCode, string, *model.Usage) { func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.ErrorWithStatusCode, string, *model.Usage) {
responseText := "" responseText := ""
scanner := bufio.NewScanner(resp.Body) scanner := bufio.NewScanner(resp.Body)
@ -36,39 +42,46 @@ func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.E
go func() { go func() {
for scanner.Scan() { for scanner.Scan() {
data := scanner.Text() data := scanner.Text()
if len(data) < 6 { // ignore blank line or wrong format if len(data) < dataPrefixLength { // ignore blank line or wrong format
continue continue
} }
if data[:6] != "data: " && data[:6] != "[DONE]" { if data[:dataPrefixLength] != dataPrefix && data[:dataPrefixLength] != done {
continue continue
} }
dataChan <- data if strings.HasPrefix(data[dataPrefixLength:], done) {
data = data[6:] dataChan <- data
if !strings.HasPrefix(data, "[DONE]") { continue
switch relayMode { }
case relaymode.ChatCompletions: switch relayMode {
var streamResponse ChatCompletionsStreamResponse case relaymode.ChatCompletions:
err := json.Unmarshal([]byte(data), &streamResponse) var streamResponse ChatCompletionsStreamResponse
if err != nil { err := json.Unmarshal([]byte(data[dataPrefixLength:]), &streamResponse)
logger.SysError("error unmarshalling stream response: " + err.Error()) if err != nil {
continue // just ignore the error logger.SysError("error unmarshalling stream response: " + err.Error())
} dataChan <- data // if error happened, pass the data to client
for _, choice := range streamResponse.Choices { continue // just ignore the error
responseText += conv.AsString(choice.Delta.Content) }
} if len(streamResponse.Choices) == 0 {
if streamResponse.Usage != nil { // but for empty choice, we should not pass it to client, this is for azure
usage = streamResponse.Usage continue // just ignore empty choice
} }
case relaymode.Completions: dataChan <- data
var streamResponse CompletionsStreamResponse for _, choice := range streamResponse.Choices {
err := json.Unmarshal([]byte(data), &streamResponse) responseText += conv.AsString(choice.Delta.Content)
if err != nil { }
logger.SysError("error unmarshalling stream response: " + err.Error()) if streamResponse.Usage != nil {
continue usage = streamResponse.Usage
} }
for _, choice := range streamResponse.Choices { case relaymode.Completions:
responseText += choice.Text dataChan <- data
} var streamResponse CompletionsStreamResponse
err := json.Unmarshal([]byte(data[dataPrefixLength:]), &streamResponse)
if err != nil {
logger.SysError("error unmarshalling stream response: " + err.Error())
continue
}
for _, choice := range streamResponse.Choices {
responseText += choice.Text
} }
} }
} }