From 6758ec4d969ae959b09cacf354af845b474ec53d Mon Sep 17 00:00:00 2001 From: sashabaranov <677093+sashabaranov@users.noreply.github.com> Date: Tue, 7 Feb 2023 20:42:53 +0400 Subject: [PATCH] Streaming support (#61) * Add streaming support feature (#54) * Add streaming support feature removes golangci linting deprecation warnings See: [Issue #49](https://github.com/sashabaranov/go-gpt3/issues/49) * remove dead token * Remove the goroutines from previous implementation Set up separate test and file for streaming support Add client code under cmd dir * Supress CI errors Need to update import path to test under feature/streaming-support branch * suppress linting errors --------- Co-authored-by: sashabaranov <677093+sashabaranov@users.noreply.github.com> * remove main.go * remove code duplication * use int64 * finalize streaming support * lint * fix tests --------- Co-authored-by: e. alvarez <55966724+ealvar3z@users.noreply.github.com> --- README.md | 48 ++++++++++++++++++ api_test.go | 32 +++++++++++- stream.go | 83 ++++++++++++++++++++++++++++++ stream_test.go | 134 +++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 296 insertions(+), 1 deletion(-) create mode 100644 stream.go create mode 100644 stream_test.go diff --git a/README.md b/README.md index b02976d..3600b05 100644 --- a/README.md +++ b/README.md @@ -38,3 +38,51 @@ func main() { fmt.Println(resp.Choices[0].Text) } ``` + +Streaming response example: + +```go +package main + +import ( + "errors" + "context" + "fmt" + "io" + gogpt "github.com/sashabaranov/go-gpt3" +) + +func main() { + c := gogpt.NewClient("your token") + ctx := context.Background() + + req := gogpt.CompletionRequest{ + Model: gogpt.GPT3Ada, + MaxTokens: 5, + Prompt: "Lorem ipsum", + Stream: true, + } + stream, err := c.CreateCompletionStream(ctx, req) + if err != nil { + return + } + defer stream.Close() + + for { + response, err := stream.Recv() + if errors.Is(err, io.EOF) { + fmt.Println("Stream finished") + return + } + + if err != nil { + fmt.Printf("Stream error: %v\n", err) + return + } + + + fmt.Printf("Stream response: %v\n", response) + + } +} +``` diff --git a/api_test.go b/api_test.go index f1950df..d0b4d52 100644 --- a/api_test.go +++ b/api_test.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "log" @@ -19,7 +20,7 @@ import ( ) const ( - testAPIToken = "this-is-my-secure-token-do-not-steal!!" + testAPIToken = "this-is-my-secure-token-do-not-steal!" ) func TestAPI(t *testing.T) { @@ -64,6 +65,33 @@ func TestAPI(t *testing.T) { if err != nil { t.Fatalf("Embedding error: %v", err) } + + stream, err := c.CreateCompletionStream(ctx, CompletionRequest{ + Prompt: "Ex falso quodlibet", + Model: GPT3Ada, + MaxTokens: 5, + Stream: true, + }) + if err != nil { + t.Errorf("CreateCompletionStream returned error: %v", err) + } + defer stream.Close() + + counter := 0 + for { + _, err = stream.Recv() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + t.Errorf("Stream error: %v", err) + } else { + counter++ + } + } + if counter == 0 { + t.Error("Stream did not return any responses") + } } // TestCompletions Tests the completions endpoint of the API using the mocked server. @@ -272,6 +300,7 @@ func handleCompletionEndpoint(w http.ResponseWriter, r *http.Request) { http.Error(w, "could not read request", http.StatusInternalServerError) return } + res := CompletionResponse{ ID: strconv.Itoa(int(time.Now().Unix())), Object: "test-object", @@ -281,6 +310,7 @@ func handleCompletionEndpoint(w http.ResponseWriter, r *http.Request) { // would be required / wouldn't make much sense Model: completionReq.Model, } + // create completions for i := 0; i < completionReq.N; i++ { // generate a random string of length completionReq.Length diff --git a/stream.go b/stream.go new file mode 100644 index 0000000..8c87e4c --- /dev/null +++ b/stream.go @@ -0,0 +1,83 @@ +package gogpt + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" +) + +type CompletionStream struct { + reader *bufio.Reader + response *http.Response +} + +func (stream *CompletionStream) Recv() (response CompletionResponse, err error) { +waitForData: + line, err := stream.reader.ReadBytes('\n') + if err != nil { + if errors.Is(err, io.EOF) { + return + } + } + + var headerData = []byte("data: ") + line = bytes.TrimSpace(line) + if !bytes.HasPrefix(line, headerData) { + goto waitForData + } + + line = bytes.TrimPrefix(line, headerData) + if string(line) == "[DONE]" { + return + } + + err = json.Unmarshal(line, &response) + return +} + +func (stream *CompletionStream) Close() { + stream.response.Body.Close() +} + +// CreateCompletionStream — API call to create a completion w/ streaming +// support. It sets whether to stream back partial progress. If set, tokens will be +// sent as data-only server-sent events as they become available, with the +// stream terminated by a data: [DONE] message. +func (c *Client) CreateCompletionStream( + ctx context.Context, + request CompletionRequest, +) (stream *CompletionStream, err error) { + request.Stream = true + reqBytes, err := json.Marshal(request) + if err != nil { + return + } + + urlSuffix := "/completions" + req, err := http.NewRequest("POST", c.fullURL(urlSuffix), bytes.NewBuffer(reqBytes)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("Cache-Control", "no-cache") + req.Header.Set("Connection", "keep-alive") + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.authToken)) + if err != nil { + return + } + + req = req.WithContext(ctx) + resp, err := c.HTTPClient.Do(req) //nolint:bodyclose // body is closed in stream.Close() + if err != nil { + return + } + + stream = &CompletionStream{ + reader: bufio.NewReader(resp.Body), + response: resp, + } + return +} diff --git a/stream_test.go b/stream_test.go new file mode 100644 index 0000000..bd7ddf7 --- /dev/null +++ b/stream_test.go @@ -0,0 +1,134 @@ +package gogpt_test + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + . "github.com/sashabaranov/go-gpt3" +) + +func TestCreateCompletionStream(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + + // Send test responses + dataBytes := []byte{} + dataBytes = append(dataBytes, []byte("event: message\n")...) + //nolint:lll + data := `{"id":"1","object":"completion","created":1598069254,"model":"text-davinci-002","choices":[{"text":"response1","finish_reason":"max_tokens"}]}` + dataBytes = append(dataBytes, []byte("data: "+data+"\n\n")...) + + dataBytes = append(dataBytes, []byte("event: message\n")...) + //nolint:lll + data = `{"id":"2","object":"completion","created":1598069255,"model":"text-davinci-002","choices":[{"text":"response2","finish_reason":"max_tokens"}]}` + dataBytes = append(dataBytes, []byte("data: "+data+"\n\n")...) + + dataBytes = append(dataBytes, []byte("event: done\n")...) + dataBytes = append(dataBytes, []byte("data: [DONE]\n\n")...) + + _, err := w.Write(dataBytes) + if err != nil { + t.Errorf("Write error: %s", err) + } + })) + defer server.Close() + + // Client portion of the test + client := NewClient(testAPIToken) + ctx := context.Background() + client.BaseURL = server.URL + "/v1" + + request := CompletionRequest{ + Prompt: "Ex falso quodlibet", + Model: "text-davinci-002", + MaxTokens: 10, + Stream: true, + } + + client.HTTPClient.Transport = &tokenRoundTripper{ + testAPIToken, + http.DefaultTransport, + } + + stream, err := client.CreateCompletionStream(ctx, request) + if err != nil { + t.Errorf("CreateCompletionStream returned error: %v", err) + } + defer stream.Close() + + expectedResponses := []CompletionResponse{ + { + ID: "1", + Object: "completion", + Created: 1598069254, + Model: "text-davinci-002", + Choices: []CompletionChoice{{Text: "response1", FinishReason: "max_tokens"}}, + }, + { + ID: "2", + Object: "completion", + Created: 1598069255, + Model: "text-davinci-002", + Choices: []CompletionChoice{{Text: "response2", FinishReason: "max_tokens"}}, + }, + {}, + } + + for ix, expectedResponse := range expectedResponses { + receivedResponse, streamErr := stream.Recv() + if streamErr != nil { + t.Errorf("stream.Recv() failed: %v", streamErr) + } + if !compareResponses(expectedResponse, receivedResponse) { + t.Errorf("Stream response %v is %v, expected %v", ix, receivedResponse, expectedResponse) + } + } +} + +// A "tokenRoundTripper" is a struct that implements the RoundTripper +// interface, specifically to handle the authentication token by adding a token +// to the request header. We need this because the API requires that each +// request include a valid API token in the headers for authentication and +// authorization. +type tokenRoundTripper struct { + token string + fallback http.RoundTripper +} + +// RoundTrip takes an *http.Request as input and returns an +// *http.Response and an error. +// +// It is expected to use the provided request to create a connection to an HTTP +// server and return the response, or an error if one occurred. The returned +// Response should have its Body closed. If the RoundTrip method returns an +// error, the Client's Get, Head, Post, and PostForm methods return the same +// error. +func (t *tokenRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + req.Header.Set("Authorization", "Bearer "+t.token) + return t.fallback.RoundTrip(req) +} + +// Helper funcs. +func compareResponses(r1, r2 CompletionResponse) bool { + if r1.ID != r2.ID || r1.Object != r2.Object || r1.Created != r2.Created || r1.Model != r2.Model { + return false + } + if len(r1.Choices) != len(r2.Choices) { + return false + } + for i := range r1.Choices { + if !compareResponseChoices(r1.Choices[i], r2.Choices[i]) { + return false + } + } + return true +} + +func compareResponseChoices(c1, c2 CompletionChoice) bool { + if c1.Text != c2.Text || c1.FinishReason != c2.FinishReason { + return false + } + return true +}