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

feat: strict testing

Browse Source
This commit is contained in:
ckt1031 2023-07-11 23:01:36 +08:00
parent 3baad1d926
commit 9ef8167e5d
1 changed files with 11 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
controller/channel-test.go
View File

@ -1,7 +1,6 @@
package controller package controller
import ( import (
"bufio"
"bytes" "bytes"
"encoding/json" "encoding/json"
"errors" "errors"
@ -60,68 +59,23 @@ func testChannel(channel *model.Channel, request ChatRequest) error {
} }
var response TextResponse var response TextResponse
isStream := strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
var streamResponseText string
if isStream { body, err := io.ReadAll(resp.Body)
scanner := bufio.NewScanner(resp.Body) if err != nil {
scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) { return err
if atEOF && len(data) == 0 { }
return 0, nil, nil err = json.Unmarshal(body, &response)
} if err != nil {
return err
}
if i := strings.Index(string(data), "\n\n"); i >= 0 { // channel.BaseURL starts with https://api.openai.com
return i + 2, data[0:i], nil if response.Usage.CompletionTokens == 0 && strings.HasPrefix(channel.BaseURL, "https://api.openai.com") {
} return errors.New(fmt.Sprintf("type %s, code %v, message %s", response.Error.Type, response.Error.Code, response.Error.Message))
if atEOF {
return len(data), data, nil
}
return 0, nil, nil
})
for scanner.Scan() {
data := scanner.Text()
if len(data) < 6 { // must be something wrong!
common.SysError("invalid stream response: " + data)
continue
}
data = data[6:]
if !strings.HasPrefix(data, "[DONE]") {
var streamResponse ChatCompletionsStreamResponse
err = json.Unmarshal([]byte(data), &streamResponse)
if err != nil {
common.SysError("error unmarshalling stream response: " + err.Error())
return err
}
for _, choice := range streamResponse.Choices {
streamResponseText += choice.Delta.Content
}
}
}
} else {
body, err := io.ReadAll(resp.Body)
if err != nil {
return err
}
err = json.Unmarshal(body, &response)
if err != nil {
return err
}
// channel.BaseURL starts with https://api.openai.com
if response.Usage.CompletionTokens == 0 && strings.HasPrefix(channel.BaseURL, "https://api.openai.com") {
return errors.New(fmt.Sprintf("type %s, code %v, message %s", response.Error.Type, response.Error.Code, response.Error.Message))
}
} }
defer resp.Body.Close() defer resp.Body.Close()
// Check if streaming is complete and streamResponseText is populated
if isStream && streamResponseText == "" {
return errors.New("Streaming not complete")
}
return nil return nil
} }