Added mistral options

Author: djthorpe

cmd/api/main.go

@@ -24,6 +24,7 @@ func main() {
 	elRegister(flags)
 	haRegister(flags)
 	ipifyRegister(flags)
+	mistralRegister(flags)
 	newsapiRegister(flags)
 	samRegister(flags)
 	weatherapiRegister(flags)

cmd/api/mistral.go

@@ -0,0 +1,169 @@
+package main
+
+import (
+	"context"
+	"fmt"
+
+	// Packages
+
+	"github.com/djthorpe/go-tablewriter"
+	"github.com/mutablelogic/go-client"
+	"github.com/mutablelogic/go-client/pkg/mistral"
+	"github.com/mutablelogic/go-client/pkg/openai/schema"
+)
+
+///////////////////////////////////////////////////////////////////////////////
+// GLOBALS
+
+var (
+	mistralName           = "mistral"
+	mistralClient         *mistral.Client
+	mistralModel          string
+	mistralEncodingFormat string
+	mistralTemperature    *float64
+	mistralMaxTokens      *uint64
+	mistralStream         *bool
+	mistralSafePrompt     bool
+	mistralSeed           *uint64
+	mistralSystemPrompt   string
+)
+
+///////////////////////////////////////////////////////////////////////////////
+// LIFECYCLE
+
+func mistralRegister(flags *Flags) {
+	// Register flags required
+	flags.String(mistralName, "mistral-api-key", "${MISTRAL_API_KEY}", "API Key")
+	flags.String(mistralName, "model", "", "Model to use")
+	flags.String(mistralName, "encoding-format", "", "The format of the output data")
+	flags.String(mistralName, "system", "", "Provide a system prompt to the model")
+	flags.Float(mistralName, "temperature", 0, "Sampling temperature to use, between 0.0 and 1.0")
+	flags.Unsigned(mistralName, "max-tokens", 0, "Maximum number of tokens to generate")
+	flags.Bool(mistralName, "stream", false, "Stream output")
+	flags.Bool(mistralName, "safe-prompt", false, "Inject a safety prompt before all conversations.")
+	flags.Unsigned(mistralName, "seed", 0, "Set random seed")
+
+	flags.Register(Cmd{
+		Name:        mistralName,
+		Description: "Interact with Mistral models, from https://docs.mistral.ai/api/",
+		Parse:       mistralParse,
+		Fn: []Fn{
+			{Name: "models", Call: mistralModels, Description: "Gets a list of available models"},
+			{Name: "embeddings", Call: mistralEmbeddings, Description: "Create embeddings from text", MinArgs: 1, Syntax: "<text>..."},
+			{Name: "chat", Call: mistralChat, Description: "Create a chat completion", MinArgs: 1, Syntax: "<text>..."},
+		},
+	})
+}
+
+func mistralParse(flags *Flags, opts ...client.ClientOpt) error {
+	apiKey := flags.GetString("mistral-api-key")
+	if apiKey == "" {
+		return fmt.Errorf("missing -mistral-api-key flag")
+	} else if client, err := mistral.New(apiKey, opts...); err != nil {
+		return err
+	} else {
+		mistralClient = client
+	}
+
+	// Get the command-line parameters
+	mistralModel = flags.GetString("model")
+	mistralEncodingFormat = flags.GetString("encoding-format")
+	mistralSafePrompt = flags.GetBool("safe-prompt")
+	mistralSystemPrompt = flags.GetString("system")
+	if temp, err := flags.GetValue("temperature"); err == nil {
+		t := temp.(float64)
+		mistralTemperature = &t
+	}
+	if maxtokens, err := flags.GetValue("max-tokens"); err == nil {
+		t := maxtokens.(uint64)
+		mistralMaxTokens = &t
+	}
+	if stream, err := flags.GetValue("stream"); err == nil {
+		t := stream.(bool)
+		mistralStream = &t
+	}
+	if seed, err := flags.GetValue("seed"); err == nil {
+		t := seed.(uint64)
+		mistralSeed = &t
+	}
+
+	// Return success
+	return nil
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// METHODS
+
+func mistralModels(ctx context.Context, writer *tablewriter.Writer, args []string) error {
+	// Get models
+	models, err := mistralClient.ListModels()
+	if err != nil {
+		return err
+	}
+
+	return writer.Write(models)
+}
+
+func mistralEmbeddings(ctx context.Context, writer *tablewriter.Writer, args []string) error {
+	// Set options
+	opts := []mistral.Opt{}
+	if mistralModel != "" {
+		opts = append(opts, mistral.OptModel(mistralModel))
+	}
+	if mistralEncodingFormat != "" {
+		opts = append(opts, mistral.OptEncodingFormat(mistralEncodingFormat))
+	}
+
+	// Get embeddings
+	embeddings, err := mistralClient.CreateEmbedding(args, opts...)
+	if err != nil {
+		return err
+	}
+	return writer.Write(embeddings)
+}
+
+func mistralChat(ctx context.Context, w *tablewriter.Writer, args []string) error {
+	var messages []*schema.Message
+
+	// Set options
+	opts := []mistral.Opt{}
+	if mistralModel != "" {
+		opts = append(opts, mistral.OptModel(mistralModel))
+	}
+	if mistralTemperature != nil {
+		opts = append(opts, mistral.OptTemperature(*mistralTemperature))
+	}
+	if mistralMaxTokens != nil {
+		opts = append(opts, mistral.OptMaxTokens(int(*mistralMaxTokens)))
+	}
+	if mistralStream != nil {
+		opts = append(opts, mistral.OptStream(func() {
+			fmt.Println("STREAM")
+		}))
+	}
+	if mistralSafePrompt {
+		opts = append(opts, mistral.OptSafePrompt())
+	}
+	if mistralSeed != nil {
+		opts = append(opts, mistral.OptSeed(int(*mistralSeed)))
+	}
+	if mistralSystemPrompt != "" {
+		messages = append(messages, schema.NewMessage("system").Add(schema.Text(mistralSystemPrompt)))
+	}
+
+	// Append user message
+	message := schema.NewMessage("user")
+	for _, arg := range args {
+		message.Add(schema.Text(arg))
+	}
+	messages = append(messages, message)
+
+	// Request -> Response
+	responses, err := mistralClient.Chat(ctx, messages, opts...)
+	if err != nil {
+		return err
+	}
+
+	// Write table
+	return w.Write(responses)
+}

pkg/anthropic/message.go renamed to pkg/anthropic/messages.go

@@ -1,9 +1,11 @@
 package anthropic
 
 import (
+	// Package imports
+	schema "github.com/mutablelogic/go-client/pkg/openai/schema"
+
 	// Namespace imports
 	. "github.com/djthorpe/go-errors"
-	"github.com/mutablelogic/go-client/pkg/openai/schema"
 )
 
 ///////////////////////////////////////////////////////////////////////////////

pkg/anthropic/message_test.go renamed to pkg/anthropic/messages_test.go

@@ -2,6 +2,9 @@ package mistral
 
 import (
 	// Packages
+	"context"
+	"reflect"
+
 	client "github.com/mutablelogic/go-client"
 	schema "github.com/mutablelogic/go-client/pkg/openai/schema"
 
@@ -13,14 +16,8 @@ import (
 // TYPES
 
 type reqChat struct {
-	Model       string           `json:"model"`
-	Messages    []schema.Message `json:"messages,omitempty"`
-	Temperature float64          `json:"temperature,omitempty"`
-	TopP        float64          `json:"top_p,omitempty"`
-	MaxTokens   int              `json:"max_tokens,omitempty"`
-	Stream      bool             `json:"stream,omitempty"`
-	SafePrompt  bool             `json:"safe_prompt,omitempty"`
-	Seed        int              `json:"random_seed,omitempty"`
+	options
+	Messages []*schema.Message `json:"messages,omitempty"`
 }
 
 type respChat struct {
@@ -46,21 +43,43 @@ const (
 // API CALLS
 
 // Chat creates a model response for the given chat conversation.
-func (c *Client) Chat(messages []schema.Message) (schema.Message, error) {
+func (c *Client) Chat(ctx context.Context, messages []*schema.Message, opts ...Opt) ([]*schema.Content, error) {
 	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
+	for _, opt := range opts {
+		if err := opt(&request.options); err != nil {
+			return nil, err
+		}
+	}
 
-	// Return the response
+	// Request->Response
 	if payload, err := client.NewJSONRequest(request); err != nil {
-		return schema.Message{}, err
-	} else if err := c.Do(payload, &response, client.OptPath("chat/completions")); err != nil {
-		return schema.Message{}, err
+		return nil, err
+	} else if err := c.DoWithContext(ctx, payload, &response, client.OptPath("chat/completions")); err != nil {
+		return nil, err
 	} else if len(response.Choices) == 0 {
-		return schema.Message{}, ErrNotFound
-	} else {
-		return response.Choices[0].Message, nil
+		return nil, ErrUnexpectedResponse.With("no choices returned")
 	}
+
+	// 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))
+		}
+	}
+
+	// Return success
+	return result, nil
 }

pkg/anthropic/opts.go

@@ -14,9 +14,8 @@ import (
 
 // A request to create embeddings
 type reqCreateEmbedding struct {
-	Input          []string `json:"input"`
-	Model          string   `json:"model"`
-	EncodingFormat string   `json:"encoding_format,omitempty"`
+	Input []string `json:"input"`
+	options
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -30,12 +29,17 @@ const (
 // API CALLS
 
 // CreateEmbedding creates an embedding from a string or array of strings
-func (c *Client) CreateEmbedding(content any) (schema.Embeddings, error) {
+func (c *Client) CreateEmbedding(content any, opts ...Opt) (schema.Embeddings, error) {
 	var request reqCreateEmbedding
 	var response schema.Embeddings
 
-	// Set default model
+	// Set options
 	request.Model = defaultEmbeddingModel
+	for _, opt := range opts {
+		if err := opt(&request.options); err != nil {
+			return response, err
+		}
+	}
 
 	// Set the input, which is either a string or array of strings
 	switch v := content.(type) {
@@ -44,7 +48,7 @@ func (c *Client) CreateEmbedding(content any) (schema.Embeddings, error) {
 	case []string:
 		request.Input = v
 	default:
-		return response, ErrBadParameter
+		return response, ErrBadParameter.With("CreateEmbedding")
 	}
 
 	// Request->Response