Added streaming chat for openai

Author: djthorpe

client.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"encoding/xml"
-	"errors"
 	"fmt"
 	"io"
 	"mime"
@@ -57,7 +56,6 @@ const (
 	PathSeparator             = string(os.PathSeparator)
 	ContentTypeAny            = "*/*"
 	ContentTypeJson           = "application/json"
-	ContentTypeJsonStream     = "application/x-ndjson"
 	ContentTypeTextXml        = "text/xml"
 	ContentTypeApplicationXml = "application/xml"
 	ContentTypeTextPlain      = "text/plain"
@@ -301,31 +299,20 @@ func do(client *http.Client, req *http.Request, accept string, strict bool, out
 		return nil
 	}
 
-	// Decode the body - and call any callback once the body has been decoded
+	// Decode the body
 	switch mimetype {
-	case ContentTypeJson, ContentTypeJsonStream:
-		dec := json.NewDecoder(response.Body)
-		for {
-			if err := dec.Decode(out); errors.Is(err, io.EOF) {
-				break
-			} else if err != nil {
-				return err
-			}
-			if reqopts.callback != nil {
-				if err := reqopts.callback(); err != nil {
-					return err
-				}
-			}
+	case ContentTypeJson:
+		if err := json.NewDecoder(response.Body).Decode(out); err != nil {
+			return err
+		}
+	case ContentTypeTextStream:
+		if err := NewTextStream().Decode(response.Body, reqopts.textStreamCallback); err != nil {
+			return err
 		}
 	case ContentTypeTextXml, ContentTypeApplicationXml:
 		if err := xml.NewDecoder(response.Body).Decode(out); err != nil {
 			return err
 		}
-		if reqopts.callback != nil {
-			if err := reqopts.callback(); err != nil {
-				return err
-			}
-		}
 	default:
 		if v, ok := out.(Unmarshaler); ok {
 			return v.Unmarshal(mimetype, response.Body)
@@ -336,11 +323,6 @@ func do(client *http.Client, req *http.Request, accept string, strict bool, out
 		} else {
 			return ErrInternalAppError.Withf("do: response does not implement Unmarshaler for %q", mimetype)
 		}
-		if reqopts.callback != nil {
-			if err := reqopts.callback(); err != nil {
-				return err
-			}
-		}
 	}
 
 	// Return success

clientopts.go

@@ -56,7 +56,7 @@ func OptUserAgent(value string) ClientOpt {
 // Setting verbose to true also displays the JSON response
 func OptTrace(w io.Writer, verbose bool) ClientOpt {
 	return func(client *Client) error {
-		client.Client.Transport = NewLogTransport(w, client.Client.Transport, verbose)
+		client.Client.Transport = newLogTransport(w, client.Client.Transport, verbose)
 		return nil
 	}
 }

cmd/api/openai.go

@@ -221,7 +221,21 @@ func openaiChat(ctx context.Context, w *tablewriter.Writer, args []string) error
 		opts = append(opts, openai.OptResponseFormat(openaiResponseFormat))
 	}
 	if openaiStream {
-		opts = append(opts, openai.OptStream())
+		opts = append(opts, openai.OptStream(func(choice schema.MessageChoice) {
+			w := w.Output()
+			if choice.Delta == nil {
+				return
+			}
+			if choice.Delta.Role != "" {
+				fmt.Fprintf(w, "\nrole: %q\n", choice.Delta.Role)
+			}
+			if choice.Delta.Content != "" {
+				fmt.Fprintf(w, "%v", choice.Delta.Content)
+			}
+			if choice.FinishReason != "" {
+				fmt.Printf("\nfinish_reason: %q\n", choice.FinishReason)
+			}
+		}))
 	}
 	if openaiUser != "" {
 		opts = append(opts, openai.OptUser(openaiUser))
@@ -243,7 +257,12 @@ func openaiChat(ctx context.Context, w *tablewriter.Writer, args []string) error
 		return err
 	}
 
-	return w.Write(responses)
+	// Write table (if not streaming)
+	if !openaiStream {
+		return w.Write(responses)
+	} else {
+		return nil
+	}
 }
 
 func openaiTranscribe(ctx context.Context, w *tablewriter.Writer, args []string) error {

pkg/openai/chat.go

@@ -3,6 +3,7 @@ package openai
 import (
 	"context"
 	"encoding/json"
+	"io"
 	"reflect"
 
 	// Packages
@@ -30,19 +31,15 @@ type respChat struct {
 	Model             string                  `json:"model"`
 	Choices           []*schema.MessageChoice `json:"choices"`
 	SystemFingerprint string                  `json:"system_fingerprint,omitempty"`
-
-	Usage struct {
-		PromptTokens     int `json:"prompt_tokens"`
-		CompletionTokens int `json:"completion_tokens"`
-		TotalTokens      int `json:"total_tokens"`
-	} `json:"usage"`
+	TokenUsage        schema.TokenUsage       `json:"usage,omitempty"`
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // GLOBALS
 
 const (
 	defaultChatCompletion = "gpt-3.5-turbo"
+	endOfStreamToken      = "[DONE]"
 )
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -73,10 +70,20 @@ func (c *Client) Chat(ctx context.Context, messages []*schema.Message, opts ...O
 		}
 	}
 
-	// Return the response
+	// Set up the request
+	reqopts := []client.RequestOpt{
+		client.OptPath("chat/completions"),
+	}
+	if request.Stream {
+		reqopts = append(reqopts, client.OptTextStreamCallback(func(event client.TextStreamEvent) error {
+			return response.streamCallback(event, request.StreamCallback)
+		}))
+	}
+
+	// Request->Response
 	if payload, err := client.NewJSONRequest(request); err != nil {
 		return nil, err
-	} else if err := c.DoWithContext(ctx, payload, &response, client.OptPath("chat/completions")); err != nil {
+	} else if err := c.DoWithContext(ctx, payload, &response, reqopts...); err != nil {
 		return nil, err
 	}
 
@@ -101,3 +108,81 @@ func (c *Client) Chat(ctx context.Context, messages []*schema.Message, opts ...O
 	// Return success
 	return result, nil
 }
+
+///////////////////////////////////////////////////////////////////////////////
+// PRIVATE METHODS
+
+func (response *respChat) streamCallback(v client.TextStreamEvent, fn Callback) error {
+	var delta schema.MessageChunk
+
+	// [DONE] indicates the end of the stream, return io.EOF
+	// or decode the data into a MessageChunk
+	if v.Data == endOfStreamToken {
+		return io.EOF
+	} else if err := v.Json(&delta); err != nil {
+		return err
+	}
+
+	// Set the response fields
+	if delta.Id != "" {
+		response.Id = delta.Id
+	}
+	if delta.Model != "" {
+		response.Model = delta.Model
+	}
+	if delta.Created != 0 {
+		response.Created = delta.Created
+	}
+	if delta.SystemFingerprint != "" {
+		response.SystemFingerprint = delta.SystemFingerprint
+	}
+	if delta.TokenUsage != nil {
+		response.TokenUsage = *delta.TokenUsage
+	}
+
+	// With no choices, return success
+	if len(delta.Choices) == 0 {
+		return nil
+	}
+
+	// Append choices
+	for _, choice := range delta.Choices {
+		// Sanity check the choice index
+		if choice.Index < 0 || choice.Index >= 6 {
+			continue
+		}
+		// Ensure message has the choice
+		for {
+			if choice.Index < len(response.Choices) {
+				break
+			}
+			response.Choices = append(response.Choices, new(schema.MessageChoice))
+		}
+		// Append the choice data onto the messahe
+		if response.Choices[choice.Index].Message == nil {
+			response.Choices[choice.Index].Message = new(schema.Message)
+		}
+		if choice.Index != 0 {
+			response.Choices[choice.Index].Index = choice.Index
+		}
+		if choice.FinishReason != "" {
+			response.Choices[choice.Index].FinishReason = choice.FinishReason
+		}
+		if choice.Delta != nil {
+			if choice.Delta.Role != "" {
+				response.Choices[choice.Index].Message.Role = choice.Delta.Role
+			}
+			if choice.Delta.Content != "" {
+				response.Choices[choice.Index].Message.Add(choice.Delta.Content)
+			}
+		}
+
+		// Callback to the client
+		if fn != nil {
+			fn(choice)
+		}
+	}
+
+	// Return success
+	return nil
+}

pkg/openai/opts.go

@@ -27,6 +27,7 @@ type options struct {
 	User             string         `json:"user,omitempty"`
 	Stream           bool           `json:"stream,omitempty"`
 	StreamOptions    *streamoptions `json:"stream_options,omitempty"`
+	StreamCallback   Callback       `json:"-"`
 
 	// Options for audio
 	Speed    *float32 `json:"speed,omitempty"`
@@ -45,6 +46,9 @@ type streamoptions struct {
 
 type Opt func(*options) error
 
+// Callback when new stream data is received
+type Callback func(schema.MessageChoice)
+
 ///////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
@@ -127,14 +131,14 @@ func OptStop(value ...string) Opt {
 	}
 }
 
-// Partial message deltas will be sent, like in ChatGPT. Tokens will be sent as data-only
-// server-sent events as they become available, with the stream terminated by a data: [DONE]
-func OptStream() Opt {
+// Stream the response, which will be returned as a series of message chunks.
+func OptStream(fn Callback) Opt {
 	return func(o *options) error {
 		o.Stream = true
 		o.StreamOptions = &streamoptions{
 			IncludeUsage: true,
 		}
+		o.StreamCallback = fn
 		return nil
 	}
 }

pkg/openai/schema/message.go

@@ -35,6 +35,23 @@ type Message struct {
 	Created int64 `json:"created,omitempty"`
 }
 
+// Chat completion chunk
+type MessageChunk struct {
+	Id                string          `json:"id,omitempty"`
+	Model             string          `json:"model,omitempty"`
+	Created           int64           `json:"created,omitempty"`
+	SystemFingerprint string          `json:"system_fingerprint,omitempty"`
+	TokenUsage        *TokenUsage     `json:"usage,omitempty"`
+	Choices           []MessageChoice `json:"choices,omitempty"`
+}
+
+// Token usage
+type TokenUsage struct {
+	PromptTokens     int `json:"prompt_tokens,omitempty"`
+	CompletionTokens int `json:"completion_tokens,omitempty"`
+	TotalTokens      int `json:"total_tokens,omitempty"`
+}
+
 // One choice of chat completion messages
 type MessageChoice struct {
 	Message      *Message      `json:"message,omitempty"`
@@ -149,6 +166,11 @@ func (m MessageChoice) String() string {
 	return string(data)
 }
 
+func (m MessageChunk) String() string {
+	data, _ := json.MarshalIndent(m, "", "  ")
+	return string(data)
+}
+
 func (m MessageDelta) String() string {
 	data, _ := json.MarshalIndent(m, "", "  ")
 	return string(data)

requestopts.go

@@ -15,12 +15,8 @@ import (
 
 type requestOpts struct {
 	*http.Request
-
-	// OptResponse
-	callback func() error
-
-	// OptNoTimeout
-	noTimeout bool
+	noTimeout          bool               // OptNoTimeout
+	textStreamCallback TextStreamCallback // OptTextStreamCallback
 }
 
 type RequestOpt func(*requestOpts) error
@@ -91,21 +87,19 @@ func OptReqHeader(name, value string) RequestOpt {
 	}
 }
 
-// OptResponse calls a function when a response has been decoded,
-// used for streaming responses. The function can return an error to
-// stop the request immediately
-func OptResponse(fn func() error) RequestOpt {
+// OptNoTimeout disables the timeout for this request, useful for long-running
+// requests. The context can be used instead for cancelling requests
+func OptNoTimeout() RequestOpt {
 	return func(r *requestOpts) error {
-		r.callback = fn
+		r.noTimeout = true
 		return nil
 	}
 }
 
-// OptNoTimeout disables the timeout for this request, useful for long-running
-// requests. The context can be used instead for cancelling requests
-func OptNoTimeout() RequestOpt {
+// OptTextStreamCallback is called for each event in a text stream
+func OptTextStreamCallback(fn TextStreamCallback) RequestOpt {
 	return func(r *requestOpts) error {
-		r.noTimeout = true
+		r.textStreamCallback = fn
 		return nil
 	}
 }