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

fix: scanner issue

Browse Source
This commit is contained in:
ckt1031 2023-07-15 23:05:43 +08:00
parent 0c175b4e44
commit 68abcd48ab
1 changed files with 39 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
controller/channel-test.go
View File

@ -122,13 +122,14 @@ func testChannel(channel *model.Channel, request ChatRequest) error {
var done = false var done = false
var streamResponseText = "" var streamResponseText = ""
if channel.Type != common.ChannelTypeChatGPTWeb {
scanner := bufio.NewScanner(resp.Body) scanner := bufio.NewScanner(resp.Body)
scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) { scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
if atEOF && len(data) == 0 { if atEOF && len(data) == 0 {
return 0, nil, nil return 0, nil, nil
} }
if i := strings.Index(string(data), "\n\n"); i >= 0 { if i := strings.Index(string(data), "\n"); i >= 0 {
return i + 2, data[0:i], nil return i + 2, data[0:i], nil
} }
@ -141,10 +142,9 @@ func testChannel(channel *model.Channel, request ChatRequest) error {
for scanner.Scan() { for scanner.Scan() {
data := scanner.Text() data := scanner.Text()
if len(data) < 6 { // must be something wrong! if len(data) < 6 { // must be something wrong!
common.SysError("invalid stream response: " + data)
continue continue
} }
if channel.Type != common.ChannelTypeChatGPTWeb {
// If data has event: event content inside, remove it, it can be prefix or inside the data // If data has event: event content inside, remove it, it can be prefix or inside the data
if strings.HasPrefix(data, "event:") || strings.Contains(data, "event:") { if strings.HasPrefix(data, "event:") || strings.Contains(data, "event:") {
// Remove event: event in the front or back // Remove event: event in the front or back
@ -185,8 +185,10 @@ func testChannel(channel *model.Channel, request ChatRequest) error {
done = true done = true
break break
} }
}
} else if channel.Type == common.ChannelTypeChatGPTWeb { } else if channel.Type == common.ChannelTypeChatGPTWeb {
scanner := bufio.NewScanner(resp.Body) scanner := bufio.NewScanner(resp.Body)
go func() { go func() {
for scanner.Scan() { for scanner.Scan() {
var chatResponse ChatGptWebChatResponse var chatResponse ChatGptWebChatResponse
@ -206,7 +208,6 @@ func testChannel(channel *model.Channel, request ChatRequest) error {
} }
}() }()
} }
}
defer resp.Body.Close() defer resp.Body.Close()