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>
2023-02-07 20:42:53 +04:00
parent 3695eb3ade
commit 6758ec4d96
4 changed files with 296 additions and 1 deletions
--- 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)
+
+	}
+}
+```
--- 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
--- a/stream.go
+++ 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
+}
--- a/stream_test.go
+++ 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
+}