Finished the streaming version

Author: djthorpe

cmd/api/mistral.go

@@ -137,8 +137,8 @@ func mistralChat(ctx context.Context, w *tablewriter.Writer, args []string) erro
 		opts = append(opts, mistral.OptMaxTokens(int(*mistralMaxTokens)))
 	}
 	if mistralStream != nil {
-		opts = append(opts, mistral.OptStream(func() {
-			fmt.Println("STREAM")
+		opts = append(opts, mistral.OptStream(func(choice schema.MessageChoice) {
+			fmt.Printf("%v\n", choice)
 		}))
 	}
 	if mistralSafePrompt {

pkg/mistral/chat.go

@@ -1,10 +1,15 @@
 package mistral
 
 import (
-	// Packages
+	"bufio"
+	"bytes"
 	"context"
+	"encoding/json"
+	"io"
 	"reflect"
+	"strings"
 
+	// Packages
 	client "github.com/mutablelogic/go-client"
 	schema "github.com/mutablelogic/go-client/pkg/openai/schema"
 
@@ -25,18 +30,25 @@ type respChat struct {
 	Created int64                  `json:"created"`
 	Model   string                 `json:"model"`
 	Choices []schema.MessageChoice `json:"choices,omitempty"`
-	Usage   struct {
-		PromptTokens     int `json:"prompt_tokens"`
-		CompletionTokens int `json:"completion_tokens"`
-		TotalTokens      int `json:"total_tokens"`
-	} `json:"usage"`
+	Usage   *respUsage             `json:"usage,omitempty"`
+
+	// Private fields
+	callback Callback `json:"-"`
+}
+
+type respUsage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+	TotalTokens      int `json:"total_tokens"`
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // GLOBALS
 
 const (
 	defaultChatCompletionModel = "mistral-small-latest"
+	contentTypeTextStream      = "text/event-stream"
+	endOfStream                = "[DONE]"
 )
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -61,6 +73,9 @@ func (c *Client) Chat(ctx context.Context, messages []*schema.Message, opts ...O
 		}
 	}
 
+	// Set the callback
+	response.callback = request.callback
+
 	// Request->Response
 	if payload, err := client.NewJSONRequest(request); err != nil {
 		return nil, err
@@ -73,13 +88,117 @@ func (c *Client) Chat(ctx context.Context, messages []*schema.Message, opts ...O
 	// Return all choices
 	var result []*schema.Content
 	for _, choice := range response.Choices {
-		if str, ok := choice.Content.(string); ok {
-			result = append(result, schema.Text(str))
-		} else {
-			return nil, ErrUnexpectedResponse.With("unexpected content type", reflect.TypeOf(choice.Content))
+		if choice.Message == nil || choice.Message.Content == nil {
+			continue
+		}
+		switch v := choice.Message.Content.(type) {
+		case []string:
+			for _, v := range v {
+				result = append(result, schema.Text(v))
+			}
+		case string:
+			result = append(result, schema.Text(v))
+		default:
+			return nil, ErrUnexpectedResponse.With("unexpected content type ", reflect.TypeOf(choice.Message.Content))
 		}
 	}
 
 	// Return success
 	return result, nil
 }
+
+///////////////////////////////////////////////////////////////////////////////
+// STRINGIFY
+
+func (s respChat) String() string {
+	data, _ := json.MarshalIndent(s, "", "  ")
+	return string(data)
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// UNMARSHAL TEXT STREAM
+
+func (m *respChat) Unmarshal(mimetype string, r io.Reader) error {
+	switch mimetype {
+	case client.ContentTypeJson:
+		return json.NewDecoder(r).Decode(m)
+	case contentTypeTextStream:
+		return m.decodeTextStream(r)
+	default:
+		return ErrUnexpectedResponse.Withf("%q", mimetype)
+	}
+}
+
+func (m *respChat) decodeTextStream(r io.Reader) error {
+	var stream respChat
+	scanner := bufio.NewScanner(r)
+	buf := new(bytes.Buffer)
+
+FOR_LOOP:
+	for scanner.Scan() {
+		data := scanner.Text()
+		switch {
+		case data == "":
+			continue FOR_LOOP
+		case strings.HasPrefix(data, "data:") && strings.HasSuffix(data, endOfStream):
+			// [DONE] - Set usage from the stream, break the loop
+			m.Usage = stream.Usage
+			break FOR_LOOP
+		case strings.HasPrefix(data, "data:"):
+			// Reset
+			stream.Choices = nil
+
+			// Decode JSON data
+			data = data[6:]
+			if _, err := buf.WriteString(data); err != nil {
+				return err
+			} else if err := json.Unmarshal(buf.Bytes(), &stream); err != nil {
+				return err
+			}
+
+			// Check for sane data
+			if len(stream.Choices) == 0 {
+				return ErrUnexpectedResponse.With("no choices returned")
+			} else if stream.Choices[0].Index != 0 {
+				return ErrUnexpectedResponse.With("unexpected choice", stream.Choices[0].Index)
+			} else if stream.Choices[0].Delta == nil {
+				return ErrUnexpectedResponse.With("no delta returned")
+			}
+
+			// Append the choice
+			if len(m.Choices) == 0 {
+				message := schema.NewMessage(stream.Choices[0].Delta.Role, stream.Choices[0].Delta.Content)
+				m.Choices = append(m.Choices, schema.MessageChoice{
+					Index:        stream.Choices[0].Index,
+					Message:      message,
+					FinishReason: stream.Choices[0].FinishReason,
+				})
+			} else {
+				// Append text to the message
+				m.Choices[0].Message.Add(stream.Choices[0].Delta.Content)
+
+				// If the finish reason is set
+				if stream.Choices[0].FinishReason != "" {
+					m.Choices[0].FinishReason = stream.Choices[0].FinishReason
+				}
+			}
+
+			// Set the model and id
+			m.Id = stream.Id
+			m.Model = stream.Model
+
+			// Callback
+			if m.callback != nil {
+				m.callback(stream.Choices[0])
+			}
+
+			// Reset the buffer
+			buf.Reset()
+		default:
+			return ErrUnexpectedResponse.Withf("%q", data)
+		}
+	}
+
+	// Return any errors from the scanner
+	return scanner.Err()
+}

pkg/mistral/chat_test.go

@@ -1,6 +1,7 @@
 package mistral_test
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -16,7 +17,7 @@ func Test_chat_001(t *testing.T) {
 	client, err := mistral.New(GetApiKey(t), opts.OptTrace(os.Stderr, true))
 	assert.NoError(err)
 	assert.NotNil(client)
-	_, err = client.Chat([]schema.Message{
+	_, err = client.Chat(context.Background(), []*schema.Message{
 		{Role: "user", Content: "What is the weather"},
 	})
 	assert.NoError(err)

pkg/mistral/opts.go

@@ -1,6 +1,9 @@
 package mistral
 
 import (
+	// Packages
+	schema "github.com/mutablelogic/go-client/pkg/openai/schema"
+
 	// Namespace imports
 	. "github.com/djthorpe/go-errors"
 )
@@ -24,7 +27,8 @@ type options struct {
 // Opt is a function which can be used to set options on a request
 type Opt func(*options) error
 
-type Callback func()
+// Callback when new stream data is received
+type Callback func(schema.MessageChoice)
 
 ///////////////////////////////////////////////////////////////////////////////
 // OPTIONS

pkg/openai/schema/message.go

@@ -18,8 +18,8 @@ import (
 
 // A chat completion message
 type Message struct {
-	// user or assistant
-	Role string `json:"role"`
+	// user, system or assistant
+	Role string `json:"role,omitempty"`
 
 	// Message Id
 	Id string `json:"id,omitempty"`
@@ -30,20 +30,30 @@ type Message struct {
 	// Content can be a string, array of strings, content
 	// object or an array of content objects
 	Content any `json:"content,omitempty"`
+
+	// Time the message was created, in unix seconds
+	Created int64 `json:"created,omitempty"`
 }
 
 // One choice of chat completion messages
 type MessageChoice struct {
-	Message      `json:"message"`
-	Index        int    `json:"index"`
-	FinishReason string `json:"finish_reason"`
+	Message      *Message      `json:"message,omitempty"`
+	Delta        *MessageDelta `json:"delta,omitempty"`
+	Index        int           `json:"index"`
+	FinishReason string        `json:"finish_reason,omitempty"`
+}
+
+// Delta between messages (for streaming responses)
+type MessageDelta struct {
+	Role    string `json:"role,omitempty"`
+	Content string `json:"content,omitempty"`
 }
 
 // Message Content
 type Content struct {
 	Id     string         `json:"id,omitempty"`
-	Type   string         `json:"type,width:4"`
-	Text   string         `json:"text,omitempty,wrap,width:60"`
+	Type   string         `json:"type" writer:",width:4"`
+	Text   string         `json:"text,omitempty" writer:",width:60,wrap"`
 	Source *contentSource `json:"source,omitempty"`
 	toolUse
 
@@ -134,6 +144,11 @@ func (m Message) String() string {
 	return string(data)
 }
 
+func (m MessageChoice) String() string {
+	data, _ := json.MarshalIndent(m, "", "  ")
+	return string(data)
+}
+
 func (c Content) String() string {
 	data, _ := json.MarshalIndent(c, "", "  ")
 	return string(data)