Added support for function calling, json mode & new models

Author: sam-ulrich1

.gitignore

@@ -21,3 +21,4 @@
 go.work
 
 .vscode/
+.idea/

chat.go

@@ -9,26 +9,16 @@ import (
 	"net/http"
 )
 
-const (
-	RoleUser      = "user"
-	RoleAssistant = "assistant"
-	RoleSystem    = "system"
-)
-
-type FinishReason string
-
-const (
-	FinishReasonStop   FinishReason = "stop"
-	FinishReasonLength FinishReason = "length"
-)
-
 // ChatRequestParams represents the parameters for the Chat/ChatStream method of MistralClient.
 type ChatRequestParams struct {
-	Temperature float64 `json:"temperature"` // The temperature to use for sampling. Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend altering this or TopP but not both.
-	TopP        float64 `json:"top_p"`       // An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered. We generally recommend altering this or Temperature but not both.
-	RandomSeed  int     `json:"random_seed"`
-	MaxTokens   int     `json:"max_tokens"`
-	SafePrompt  bool    `json:"safe_prompt"` // Adds a Mistral defined safety message to the system prompt to enforce guardrailing
+	Temperature    float64        `json:"temperature"` // The temperature to use for sampling. Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend altering this or TopP but not both.
+	TopP           float64        `json:"top_p"`       // An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered. We generally recommend altering this or Temperature but not both.
+	RandomSeed     int            `json:"random_seed"`
+	MaxTokens      int            `json:"max_tokens"`
+	SafePrompt     bool           `json:"safe_prompt"` // Adds a Mistral defined safety message to the system prompt to enforce guardrailing
+	Tools          []Tool         `json:"tools"`
+	ToolChoice     string         `json:"tool_choice"`
+	ResponseFormat ResponseFormat `json:"response_format"`
 }
 
 var DefaultChatRequestParams = ChatRequestParams{
@@ -39,12 +29,6 @@ var DefaultChatRequestParams = ChatRequestParams{
 	SafePrompt:  false,
 }
 
-// ChatMessage represents a single message in a chat.
-type ChatMessage struct {
-	Role    string `json:"role"`
-	Content string `json:"content"`
-}
-
 // ChatCompletionResponseChoice represents a choice in the chat completion response.
 type ChatCompletionResponseChoice struct {
 	Index        int          `json:"index"`
@@ -55,7 +39,7 @@ type ChatCompletionResponseChoice struct {
 // ChatCompletionResponseChoice represents a choice in the chat completion response.
 type ChatCompletionResponseChoiceStream struct {
 	Index        int          `json:"index"`
-	Delta        ChatMessage  `json:"delta"`
+	Delta        DeltaMessage `json:"delta"`
 	FinishReason FinishReason `json:"finish_reason,omitempty"`
 }
 
@@ -102,6 +86,16 @@ func (c *MistralClient) Chat(model string, messages []ChatMessage, params *ChatR
 		"safe_prompt": params.SafePrompt,
 	}
 
+	if params.Tools != nil {
+		requestData["tools"] = params.Tools
+	}
+	if params.ToolChoice != "" {
+		requestData["tool_choice"] = params.ToolChoice
+	}
+	if params.ResponseFormat != "" {
+		requestData["response_format"] = map[string]any{"type": params.ResponseFormat}
+	}
+
 	response, err := c.request(http.MethodPost, requestData, "v1/chat/completions", false, nil)
 	if err != nil {
 		return nil, err
@@ -140,6 +134,16 @@ func (c *MistralClient) ChatStream(model string, messages []ChatMessage, params
 		"stream":      true,
 	}
 
+	if params.Tools != nil {
+		requestData["tools"] = params.Tools
+	}
+	if params.ToolChoice != "" {
+		requestData["tool_choice"] = params.ToolChoice
+	}
+	if params.ResponseFormat != "" {
+		requestData["response_format"] = map[string]any{"type": params.ResponseFormat}
+	}
+
 	response, err := c.request(http.MethodPost, requestData, "v1/chat/completions", true, nil)
 	if err != nil {
 		return nil, err

chat_test.go

@@ -12,7 +12,7 @@ func TestChat(t *testing.T) {
 	params.MaxTokens = 10
 	params.Temperature = 0
 	res, err := client.Chat(
-		"mistral-tiny",
+		ModelMistralTiny,
 		[]ChatMessage{
 			{
 				Role:    RoleUser,
@@ -30,13 +30,99 @@ func TestChat(t *testing.T) {
 	assert.Equal(t, res.Choices[0].Message.Content, "Test Succeeded")
 }
 
+func TestChatFunctionCall(t *testing.T) {
+	client := NewMistralClientDefault("")
+	params := DefaultChatRequestParams
+	params.Temperature = 0
+	params.Tools = []Tool{
+		{
+			Type: ToolTypeFunction,
+			Function: Function{
+				Name:        "get_weather",
+				Description: "Retrieve the weather for a city in the US",
+				Parameters: map[string]interface{}{
+					"type":     "object",
+					"required": []string{"city", "state"},
+					"properties": map[string]interface{}{
+						"city":  map[string]interface{}{"type": "string", "description": "Name of the city for the weather"},
+						"state": map[string]interface{}{"type": "string", "description": "Name of the state for the weather"},
+					},
+				},
+			},
+		},
+		{
+			Type: ToolTypeFunction,
+			Function: Function{
+				Name:        "send_text",
+				Description: "Send text message using SMS service",
+				Parameters: map[string]interface{}{
+					"type":     "object",
+					"required": []string{"contact_name", "message"},
+					"properties": map[string]interface{}{
+						"contact_name": map[string]interface{}{"type": "string", "description": "Name of the contact that will receive the message"},
+						"message":      map[string]interface{}{"type": "string", "description": "Content of the message that will be sent"},
+					},
+				},
+			},
+		},
+	}
+	params.ToolChoice = ToolChoiceAuto
+	res, err := client.Chat(
+		ModelMistralSmallLatest,
+		[]ChatMessage{
+			{
+				Role:    RoleUser,
+				Content: "What's the weather like in Dallas, TX?",
+			},
+		},
+		&params,
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, res)
+
+	assert.Greater(t, len(res.Choices), 0)
+	assert.Greater(t, len(res.Choices[0].Message.ToolCalls), 0)
+	assert.Equal(t, res.Choices[0].Message.Role, RoleAssistant)
+	assert.Equal(t, res.Choices[0].Message.ToolCalls[0].Function.Name, "get_weather")
+	assert.Equal(t, res.Choices[0].Message.ToolCalls[0].Function.Arguments, "{\"city\": \"Dallas\", \"state\": \"TX\"}")
+}
+
+func TestChatJsonMode(t *testing.T) {
+	client := NewMistralClientDefault("")
+	params := DefaultChatRequestParams
+	params.Temperature = 0
+	params.ResponseFormat = ResponseFormatJsonObject
+	res, err := client.Chat(
+		ModelMistralSmallLatest,
+		[]ChatMessage{
+			{
+				Role: RoleUser,
+				Content: "Extract all of the code symbols in this text chunk and return them in the following JSON: " +
+					"{\"symbols\":[\"SymbolOne\",\"SymbolTwo\"]}\n```\nI'm working on updating the Go client for the " +
+					"new release, is it expected that the function call will be passed back into the model or just " +
+					"the tool response?\nI ask because ChatMessage can handle the tool response but the messages list " +
+					"has an Any option that I assume would be for the FunctionCall/ToolCall type\nAdditionally the " +
+					"example in the docs only shows the tool response appended to the messages\n```",
+			},
+		},
+		&params,
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, res)
+
+	assert.Greater(t, len(res.Choices), 0)
+	assert.Greater(t, len(res.Choices[0].Message.Content), 0)
+	assert.Equal(t, res.Choices[0].Message.Role, RoleAssistant)
+	assert.Equal(t, res.Choices[0].Message.Content, "{\"symbols\": [\"Go\", \"ChatMessage\", \"FunctionCall\", \"ToolCall\"]}")
+}
+
 func TestChatStream(t *testing.T) {
 	client := NewMistralClientDefault("")
 	params := DefaultChatRequestParams
 	params.MaxTokens = 50
 	params.Temperature = 0
 	resChan, err := client.ChatStream(
-		"mistral-tiny",
+		ModelMistralTiny,
 		[]ChatMessage{
 			{
 				Role:    RoleUser,
@@ -53,16 +139,141 @@ func TestChatStream(t *testing.T) {
 	for res := range resChan {
 		assert.NoError(t, res.Error)
 
+		assert.Greater(t, len(res.Choices), 0)
+		if idx == 0 {
+			assert.Equal(t, res.Choices[0].Delta.Role, RoleAssistant)
+		}
+		totalOutput += res.Choices[0].Delta.Content
+		idx++
+
 		if res.Choices[0].FinishReason == FinishReasonStop {
 			break
 		}
+	}
+	assert.Equal(t, totalOutput, "Test Succeeded, Test Succeeded, Test Succeeded, Test Succeeded, Test Succeeded, Test Succeeded")
+}
+
+func TestChatStreamFunctionCall(t *testing.T) {
+	client := NewMistralClientDefault("")
+	params := DefaultChatRequestParams
+	params.Temperature = 0
+	params.Tools = []Tool{
+		{
+			Type: ToolTypeFunction,
+			Function: Function{
+				Name:        "get_weather",
+				Description: "Retrieve the weather for a city in the US",
+				Parameters: map[string]interface{}{
+					"type":     "object",
+					"required": []string{"city", "state"},
+					"properties": map[string]interface{}{
+						"city":  map[string]interface{}{"type": "string", "description": "Name of the city for the weather"},
+						"state": map[string]interface{}{"type": "string", "description": "Name of the state for the weather"},
+					},
+				},
+			},
+		},
+		{
+			Type: ToolTypeFunction,
+			Function: Function{
+				Name:        "send_text",
+				Description: "Send text message using SMS service",
+				Parameters: map[string]interface{}{
+					"type":     "object",
+					"required": []string{"contact_name", "message"},
+					"properties": map[string]interface{}{
+						"contact_name": map[string]interface{}{"type": "string", "description": "Name of the contact that will receive the message"},
+						"message":      map[string]interface{}{"type": "string", "description": "Content of the message that will be sent"},
+					},
+				},
+			},
+		},
+	}
+	params.ToolChoice = ToolChoiceAuto
+	resChan, err := client.ChatStream(
+		ModelMistralSmallLatest,
+		[]ChatMessage{
+			{
+				Role:    RoleUser,
+				Content: "What's the weather like in Dallas, TX?",
+			},
+		},
+		&params,
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, resChan)
+
+	totalOutput := ""
+	var functionCall *ToolCall
+	idx := 0
+	for res := range resChan {
+		assert.NoError(t, res.Error)
 
 		assert.Greater(t, len(res.Choices), 0)
 		if idx == 0 {
 			assert.Equal(t, res.Choices[0].Delta.Role, RoleAssistant)
 		}
 		totalOutput += res.Choices[0].Delta.Content
+		if len(res.Choices[0].Delta.ToolCalls) > 0 {
+			functionCall = &res.Choices[0].Delta.ToolCalls[0]
+		}
 		idx++
+
+		if res.Choices[0].FinishReason == FinishReasonStop {
+			break
+		}
 	}
-	assert.Equal(t, totalOutput, "Test Succeeded, Test Succeeded, Test Succeeded, Test Succeeded, Test Succeeded, Test Succeeded")
+
+	assert.Equal(t, totalOutput, "")
+	assert.NotNil(t, functionCall)
+	assert.Equal(t, functionCall.Function.Name, "get_weather")
+	assert.Equal(t, functionCall.Function.Arguments, "{\"city\": \"Dallas\", \"state\": \"TX\"}")
+}
+
+func TestChatStreamJsonMode(t *testing.T) {
+	client := NewMistralClientDefault("")
+	params := DefaultChatRequestParams
+	params.Temperature = 0
+	params.ResponseFormat = ResponseFormatJsonObject
+	resChan, err := client.ChatStream(
+		ModelMistralSmallLatest,
+		[]ChatMessage{
+			{
+				Role: RoleUser,
+				Content: "Extract all of the code symbols in this text chunk and return them in the following JSON: " +
+					"{\"symbols\":[\"SymbolOne\",\"SymbolTwo\"]}\n```\nI'm working on updating the Go client for the " +
+					"new release, is it expected that the function call will be passed back into the model or just " +
+					"the tool response?\nI ask because ChatMessage can handle the tool response but the messages list " +
+					"has an Any option that I assume would be for the FunctionCall/ToolCall type\nAdditionally the " +
+					"example in the docs only shows the tool response appended to the messages\n```",
+			},
+		},
+		&params,
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, resChan)
+
+	totalOutput := ""
+	var functionCall *ToolCall
+	idx := 0
+	for res := range resChan {
+		assert.NoError(t, res.Error)
+
+		assert.Greater(t, len(res.Choices), 0)
+		if idx == 0 {
+			assert.Equal(t, res.Choices[0].Delta.Role, RoleAssistant)
+		}
+		totalOutput += res.Choices[0].Delta.Content
+		if len(res.Choices[0].Delta.ToolCalls) > 0 {
+			functionCall = &res.Choices[0].Delta.ToolCalls[0]
+		}
+		idx++
+
+		if res.Choices[0].FinishReason == FinishReasonStop {
+			break
+		}
+	}
+
+	assert.Equal(t, totalOutput, "{\"symbols\": [\"Go\", \"ChatMessage\", \"FunctionCall\", \"ToolCall\"]}")
+	assert.Nil(t, functionCall)
 }