handle stream completion (#86)

* handle stream completion

* fix tests
This commit is contained in:
sashabaranov
2023-02-22 12:33:25 +04:00
committed by GitHub
parent 1eb5d625f8
commit ae05ed976f
2 changed files with 16 additions and 1 deletions

View File

@@ -7,6 +7,7 @@ import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
)
@@ -16,12 +17,18 @@ var (
type CompletionStream struct {
emptyMessagesLimit uint
isFinished bool
reader *bufio.Reader
response *http.Response
}
func (stream *CompletionStream) Recv() (response CompletionResponse, err error) {
if stream.isFinished {
err = io.EOF
return
}
var emptyMessagesCount uint
waitForData:
@@ -44,6 +51,8 @@ waitForData:
line = bytes.TrimPrefix(line, headerData)
if string(line) == "[DONE]" {
stream.isFinished = true
err = io.EOF
return
}