Updated mistral

Author: djthorpe

pkg/mistral/chat_completion.go

@@ -78,11 +78,8 @@ func (mistral *Client) ChatCompletion(ctx context.Context, context llm.Context,
 	}
 
 	// Always append the first message of each completion
-	for _, completion := range context.(*session).seq {
-		if completion.Num() == 0 {
-			continue
-		}
-		messages = append(messages, completion.Message(0))
+	for _, message := range context.(*session).seq {
+		messages = append(messages, message)
 	}
 
 	// Request
@@ -179,17 +176,19 @@ func appendCompletion(response *Response, c *Completion) {
 		response.Completions = append(response.Completions, Completion{
 			Index: c.Index,
 			Message: &Message{
-				Role:    c.Delta.Role,
-				Content: "",
+				RoleContent: RoleContent{
+					Role:    c.Delta.Role(),
+					Content: "",
+				},
 			},
 		})
 	}
 	// Add the completion delta
 	if c.Reason != "" {
 		response.Completions[c.Index].Reason = c.Reason
 	}
-	if c.Delta.Role != "" {
-		response.Completions[c.Index].Message.Role = c.Delta.Role
+	if role := c.Delta.Role(); role != "" {
+		response.Completions[c.Index].Message.RoleContent.Role = role
 	}
 
 	// TODO: We only allow deltas which are strings at the moment...

pkg/mistral/chat_completion_test.go

@@ -2,7 +2,7 @@ package mistral_test
 
 import (
 	"context"
-	"errors"
+	"fmt"
 	"os"
 	"strings"
 	"testing"
@@ -243,6 +243,6 @@ func (weather) Description() string {
 	return "Get the weather for a city"
 }
 
-func (weather) Run(ctx context.Context) (any, error) {
-	return nil, errors.New("I couldn't retrieve the weather for that city")
+func (w weather) Run(ctx context.Context) (any, error) {
+	return fmt.Sprintf("The weather in %q is sunny and warm", w.City), nil
 }

pkg/mistral/message.go

@@ -1,7 +1,10 @@
 package mistral
 
 import (
+	"encoding/json"
+
 	"github.com/mutablelogic/go-llm"
+	"github.com/mutablelogic/go-llm/pkg/tool"
 )
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -12,6 +15,21 @@ type Completions []Completion
 
 var _ llm.Completion = Completions{}
 
+// Message with text or object content
+type Message struct {
+	RoleContent
+	ToolCallArray `json:"tool_calls,omitempty"`
+}
+
+type RoleContent struct {
+	Role    string `json:"role,omitempty"`         // assistant, user, tool, system
+	Id      string `json:"tool_call_id,omitempty"` // tool call - when role is tool
+	Name    string `json:"name,omitempty"`         // function name - when role is tool
+	Content any    `json:"content,omitempty"`      // string or array of text, reference, image_url
+}
+
+var _ llm.Completion = (*Message)(nil)
+
 // Completion Variation
 type Completion struct {
 	Index   uint64   `json:"index"`
@@ -20,23 +38,15 @@ type Completion struct {
 	Reason  string   `json:"finish_reason,omitempty"`
 }
 
-// Message with text or object content
-type Message struct {
-	Role      string `json:"role,omitempty"` // assistant, user, tool, system
-	Prefix    bool   `json:"prefix,omitempty"`
-	Content   any    `json:"content,omitempty"`
-	ToolCalls `json:"tool_calls,omitempty"`
-}
-
 type Content struct {
-	Type        string                       `json:"type"` // text, reference, image_url
+	Type        string                       `json:"type,omitempty"` // text, reference, image_url
 	*Text       `json:"text,omitempty"`      // text content
 	*Prediction `json:"content,omitempty"`   // prediction
 	*Image      `json:"image_url,omitempty"` // image_url
 }
 
 // A set of tool calls
-type ToolCalls []ToolCall
+type ToolCallArray []ToolCall
 
 // text content
 type Text string
@@ -78,62 +88,87 @@ func NewImageAttachment(a *llm.Attachment) *Content {
 }
 
 ///////////////////////////////////////////////////////////////////////////////
-// PUBLIC METHODS
+// PUBLIC METHODS - MESSAGE
 
-// Return the number of completions
-func (c Completions) Num() int {
-	return len(c)
+func (m Message) Num() int {
+	return 1
 }
 
-// Return the role of the completion
-func (c Completions) Role() string {
-	// The role should be the same for all completions, let's use the first one
-	if len(c) == 0 {
-		return ""
-	}
-	return c[0].Message.Role
+func (m Message) Role() string {
+	return m.RoleContent.Role
 }
 
-// Return the text content for a specific completion
-func (c Completions) Text(index int) string {
-	if index < 0 || index >= len(c) {
+func (m Message) Text(index int) string {
+	if index != 0 {
 		return ""
 	}
-	completion := c[index].Message
-	if text, ok := completion.Content.(string); ok {
+	// If content is text, return it
+	if text, ok := m.Content.(string); ok {
 		return text
 	}
-	// TODO: Will the text be in other forms?
+	// For other kinds, return empty string for the moment
 	return ""
 }
 
-// Return the current session tool calls given the completion index.
-// Will return nil if no tool calls were returned.
-func (c Completions) ToolCalls(index int) []llm.ToolCall {
-	if index < 0 || index >= len(c) {
-		return nil
-	}
-
-	// Get the completion
-	completion := c[index].Message
-	if completion == nil {
+func (m Message) ToolCalls(index int) []llm.ToolCall {
+	if index != 0 {
 		return nil
 	}
 
 	// Make the tool calls
-	calls := make([]llm.ToolCall, 0, len(completion.ToolCalls))
-	for _, call := range completion.ToolCalls {
-		calls = append(calls, &toolcall{call})
+	calls := make([]llm.ToolCall, 0, len(m.ToolCallArray))
+	for _, call := range m.ToolCallArray {
+		var args map[string]any
+		if call.Function.Arguments != "" {
+			if err := json.Unmarshal([]byte(call.Function.Arguments), &args); err != nil {
+				return nil
+			}
+		}
+		calls = append(calls, tool.NewCall(call.Id, call.Function.Name, args))
 	}
 
 	// Return success
 	return calls
 }
 
+///////////////////////////////////////////////////////////////////////////////
+// PUBLIC METHODS - COMPLETIONS
+
+// Return the number of completions
+func (c Completions) Num() int {
+	return len(c)
+}
+
 // Return message for a specific completion
 func (c Completions) Message(index int) *Message {
 	if index < 0 || index >= len(c) {
 		return nil
 	}
 	return c[index].Message
 }
+
+// Return the role of the completion
+func (c Completions) Role() string {
+	// The role should be the same for all completions, let's use the first one
+	if len(c) == 0 {
+		return ""
+	}
+	return c[0].Message.Role()
+}
+
+// Return the text content for a specific completion
+func (c Completions) Text(index int) string {
+	if index < 0 || index >= len(c) {
+		return ""
+	}
+	return c[index].Message.Text(0)
+}
+
+// Return the current session tool calls given the completion index.
+// Will return nil if no tool calls were returned.
+func (c Completions) ToolCalls(index int) []llm.ToolCall {
+	if index < 0 || index >= len(c) {
+		return nil
+	}
+	return c[index].Message.ToolCalls(0)
+}

pkg/mistral/model.go

@@ -12,6 +12,7 @@ import (
 // TYPES
 
 type model struct {
+	*Client
 	meta Model
 }
 
@@ -65,7 +66,7 @@ func (c *Client) ListModels(ctx context.Context) ([]llm.Model, error) {
 	//  Make models
 	result := make([]llm.Model, 0, len(response.Data))
 	for _, meta := range response.Data {
-		result = append(result, &model{meta: meta})
+		result = append(result, &model{c, meta})
 	}
 
 	// Return models