Added tool calls to mistral

Author: djthorpe

pkg/mistral/chat.go

@@ -1,13 +1,10 @@
 package mistral
 
 import (
-	"bufio"
-	"bytes"
 	"context"
 	"encoding/json"
 	"io"
 	"reflect"
-	"strings"
 
 	// Packages
 	client "github.com/mutablelogic/go-client"
@@ -20,37 +17,46 @@ import (
 ///////////////////////////////////////////////////////////////////////////////
 // TYPES
 
+// A request for a chat completion
 type reqChat struct {
 	options
+	Tools    []reqChatTools    `json:"tools,omitempty"`
 	Messages []*schema.Message `json:"messages,omitempty"`
 }
 
-type respChat struct {
-	Id      string                 `json:"id"`
-	Created int64                  `json:"created"`
-	Model   string                 `json:"model"`
-	Choices []schema.MessageChoice `json:"choices,omitempty"`
-	Usage   *respUsage             `json:"usage,omitempty"`
-
-	// Private fields
-	callback Callback `json:"-"`
+type reqChatTools struct {
+	Type     string       `json:"type"`
+	Function *schema.Tool `json:"function"`
 }
 
-type respUsage struct {
-	PromptTokens     int `json:"prompt_tokens"`
-	CompletionTokens int `json:"completion_tokens"`
-	TotalTokens      int `json:"total_tokens"`
+// A chat completion object
+type respChat struct {
+	Id         string                  `json:"id"`
+	Created    int64                   `json:"created"`
+	Model      string                  `json:"model"`
+	Choices    []*schema.MessageChoice `json:"choices,omitempty"`
+	TokenUsage schema.TokenUsage       `json:"usage,omitempty"`
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // GLOBALS
 
 const (
 	defaultChatCompletionModel = "mistral-small-latest"
-	contentTypeTextStream      = "text/event-stream"
-	endOfStream                = "[DONE]"
+	endOfStreamToken           = "[DONE]"
 )
 
+///////////////////////////////////////////////////////////////////////////////
+// STRINGIFY
+
+func (v respChat) String() string {
+	if data, err := json.MarshalIndent(v, "", "  "); err != nil {
+		return err.Error()
+	} else {
+		return string(data)
+	}
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // API CALLS
 
@@ -59,11 +65,6 @@ func (c *Client) Chat(ctx context.Context, messages []*schema.Message, opts ...O
 	var request reqChat
 	var response respChat
 
-	// Check messages
-	if len(messages) == 0 {
-		return nil, ErrBadParameter.With("missing messages")
-	}
-
 	// Process options
 	request.Model = defaultChatCompletionModel
 	request.Messages = messages
@@ -73,13 +74,28 @@ func (c *Client) Chat(ctx context.Context, messages []*schema.Message, opts ...O
 		}
 	}
 
-	// Set the callback
-	response.callback = request.callback
+	// Append tools
+	for _, tool := range request.options.Tools {
+		request.Tools = append(request.Tools, reqChatTools{
+			Type:     "function",
+			Function: tool,
+		})
+	}
+
+	// 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
 	} else if len(response.Choices) == 0 {
 		return nil, ErrUnexpectedResponse.With("no choices returned")
@@ -91,6 +107,9 @@ func (c *Client) Chat(ctx context.Context, messages []*schema.Message, opts ...O
 		if choice.Message == nil || choice.Message.Content == nil {
 			continue
 		}
+		for _, tool := range choice.Message.ToolCalls {
+			result = append(result, schema.ToolUse(tool))
+		}
 		switch v := choice.Message.Content.(type) {
 		case []string:
 			for _, v := range v {
@@ -108,97 +127,76 @@ func (c *Client) Chat(ctx context.Context, messages []*schema.Message, opts ...O
 }
 
 ///////////////////////////////////////////////////////////////////////////////
-// STRINGIFY
+// PRIVATE METHODS
 
-func (s respChat) String() string {
-	data, _ := json.MarshalIndent(s, "", "  ")
-	return string(data)
-}
+func (response *respChat) streamCallback(v client.TextStreamEvent, fn Callback) error {
+	var delta schema.MessageChunk
 
-///////////////////////////////////////////////////////////////////////////////
-// 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)
+	// [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
 	}
-}
 
-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
-			}
+	// 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.TokenUsage != nil {
+		response.TokenUsage = *delta.TokenUsage
+	}
 
-			// 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")
-			}
+	// With no choices, return success
+	if len(delta.Choices) == 0 {
+		return nil
+	}
 
-			// 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
-				}
+	// 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
 			}
-
-			// Set the model and id
-			m.Id = stream.Id
-			m.Model = stream.Model
-
-			// Callback
-			if m.callback != nil {
-				m.callback(stream.Choices[0])
+			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)
+			}
+		}
 
-			// Reset the buffer
-			buf.Reset()
-		default:
-			return ErrUnexpectedResponse.Withf("%q", data)
+		// Callback to the client
+		if fn != nil {
+			fn(choice)
 		}
 	}
 
-	// Return any errors from the scanner
-	return scanner.Err()
+	// Return success
+	return nil
 }

pkg/mistral/chat_test.go

@@ -3,6 +3,7 @@ package mistral_test
 import (
 	"context"
 	"os"
+	"reflect"
 	"testing"
 
 	// Packages
@@ -22,3 +23,31 @@ func Test_chat_001(t *testing.T) {
 	})
 	assert.NoError(err)
 }
+
+func Test_chat_002(t *testing.T) {
+	assert := assert.New(t)
+	client, err := mistral.New(GetApiKey(t), opts.OptTrace(os.Stderr, true))
+	assert.NoError(err)
+	assert.NotNil(client)
+	_, err = client.Chat(context.Background(), []*schema.Message{
+		{Role: "user", Content: "What is the weather"},
+	}, mistral.OptStream(func(message schema.MessageChoice) {
+		t.Log(message)
+	}))
+	assert.NoError(err)
+}
+
+func Test_chat_003(t *testing.T) {
+	assert := assert.New(t)
+	client, err := mistral.New(GetApiKey(t), opts.OptTrace(os.Stderr, true))
+	assert.NoError(err)
+	assert.NotNil(client)
+
+	tool := schema.NewTool("weather", "get weather in a specific city")
+	tool.Add("city", "name of the city, if known", false, reflect.TypeOf(""))
+
+	_, err = client.Chat(context.Background(), []*schema.Message{
+		{Role: "user", Content: "What is the weather in Berlin"},
+	}, mistral.OptTool(tool))
+	assert.NoError(err)
+}

pkg/mistral/opts.go

@@ -12,16 +12,18 @@ import (
 // TYPES
 
 type options struct {
+	// Common options
 	Model          string   `json:"model,omitempty"`
 	EncodingFormat string   `json:"encoding_format,omitempty"`
-	Temperature    *float64 `json:"temperature,omitempty"`
+	Temperature    *float32 `json:"temperature,omitempty"`
 	MaxTokens      int      `json:"max_tokens,omitempty"`
-	Stream         bool     `json:"stream,omitempty"`
 	SafePrompt     bool     `json:"safe_prompt,omitempty"`
 	Seed           int      `json:"random_seed,omitempty"`
 
-	// Private methods
-	callback Callback `json:"-"`
+	// Options for chat
+	Stream         bool           `json:"stream,omitempty"`
+	StreamCallback Callback       `json:"-"`
+	Tools          []*schema.Tool `json:"-"`
 }
 
 // Opt is a function which can be used to set options on a request
@@ -61,7 +63,7 @@ func OptMaxTokens(v int) Opt {
 func OptStream(fn Callback) Opt {
 	return func(o *options) error {
 		o.Stream = true
-		o.callback = fn
+		o.StreamCallback = fn
 		return nil
 	}
 }
@@ -83,7 +85,7 @@ func OptSeed(v int) Opt {
 }
 
 // Amount of randomness injected into the response.
-func OptTemperature(v float64) Opt {
+func OptTemperature(v float32) Opt {
 	return func(o *options) error {
 		if v < 0.0 || v > 1.0 {
 			return ErrBadParameter.With("OptTemperature")
@@ -92,3 +94,21 @@ func OptTemperature(v float64) Opt {
 		return nil
 	}
 }
+
+// A list of tools the model may call.
+func OptTool(value ...*schema.Tool) Opt {
+	return func(o *options) error {
+		// Check tools
+		for _, tool := range value {
+			if tool == nil {
+				return ErrBadParameter.With("OptTool")
+			}
+		}
+
+		// Append tools
+		o.Tools = append(o.Tools, value...)
+
+		// Return success
+		return nil
+	}
+}