feat: add Mistral provider and enhance ToolChoice handling

core/bifrost.go

@@ -138,6 +138,8 @@ func (bifrost *Bifrost) createProviderFromProviderKey(providerKey schemas.ModelP
 		return providers.NewAzureProvider(config, bifrost.logger)
 	case schemas.Vertex:
 		return providers.NewVertexProvider(config, bifrost.logger)
+	case schemas.Mistral:
+		return providers.NewMistralProvider(config, bifrost.logger), nil
 	default:
 		return nil, fmt.Errorf("unsupported provider: %s", providerKey)
 	}

core/providers/anthropic.go

@@ -134,7 +134,7 @@ func NewAnthropicProvider(config *schemas.ProviderConfig, logger schemas.Logger)
 	client := &fasthttp.Client{
 		ReadTimeout:     time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
 		WriteTimeout:    time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
-		MaxConnsPerHost: config.ConcurrencyAndBufferSize.BufferSize,
+		MaxConnsPerHost: config.ConcurrencyAndBufferSize.Concurrency,
 	}
 
 	// Pre-warm response pools
@@ -515,17 +515,21 @@ func prepareAnthropicChatRequest(messages []schemas.BifrostMessage, params *sche
 
 	// Transform tool choice if present
 	if params != nil && params.ToolChoice != nil {
-		switch toolChoice := params.ToolChoice.Type; toolChoice {
-		case schemas.ToolChoiceTypeFunction:
-			fallthrough
-		case "tool":
-			preparedParams["tool_choice"] = map[string]interface{}{
-				"type": "tool",
-				"name": params.ToolChoice.Function.Name,
-			}
-		default:
-			preparedParams["tool_choice"] = map[string]interface{}{
-				"type": toolChoice,
+		if params.ToolChoice.ToolChoiceStr != nil {
+			preparedParams["tool_choice"] = *params.ToolChoice.ToolChoiceStr
+		} else if params.ToolChoice.ToolChoiceStruct != nil {
+			switch toolChoice := params.ToolChoice.ToolChoiceStruct.Type; toolChoice {
+			case schemas.ToolChoiceTypeFunction:
+				fallthrough
+			case "tool":
+				preparedParams["tool_choice"] = map[string]interface{}{
+					"type": "tool",
+					"name": params.ToolChoice.ToolChoiceStruct.Function.Name,
+				}
+			default:
+				preparedParams["tool_choice"] = map[string]interface{}{
+					"type": toolChoice,
+				}
 			}
 		}
 	}

core/providers/azure.go

@@ -114,7 +114,7 @@ func NewAzureProvider(config *schemas.ProviderConfig, logger schemas.Logger) (*A
 	client := &fasthttp.Client{
 		ReadTimeout:     time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
 		WriteTimeout:    time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
-		MaxConnsPerHost: config.ConcurrencyAndBufferSize.BufferSize,
+		MaxConnsPerHost: config.ConcurrencyAndBufferSize.Concurrency,
 	}
 
 	// Pre-warm response pools

core/providers/cohere.go

@@ -109,7 +109,7 @@ func NewCohereProvider(config *schemas.ProviderConfig, logger schemas.Logger) *C
 	client := &fasthttp.Client{
 		ReadTimeout:     time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
 		WriteTimeout:    time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
-		MaxConnsPerHost: config.ConcurrencyAndBufferSize.BufferSize,
+		MaxConnsPerHost: config.ConcurrencyAndBufferSize.Concurrency,
 	}
 
 	// Pre-warm response pools
@@ -325,7 +325,13 @@ func (provider *CohereProvider) ChatCompletion(ctx context.Context, model, key s
 	}
 	// Add tool choice if present
 	if params != nil && params.ToolChoice != nil {
-		requestBody["tool_choice"] = strings.ToUpper(string(params.ToolChoice.Type))
+		if params.ToolChoice.ToolChoiceStr != nil {
+			requestBody["tool_choice"] = *params.ToolChoice.ToolChoiceStr
+		} else if params.ToolChoice.ToolChoiceStruct != nil {
+			requestBody["tool_choice"] = map[string]interface{}{
+				"type": strings.ToUpper(string(params.ToolChoice.ToolChoiceStruct.Type)),
+			}
+		}
 	}
 
 	// Marshal request body

core/providers/mistral.go

@@ -0,0 +1,180 @@
+// Package providers implements various LLM providers and their utility functions.
+// This file contains the Mistral provider implementation.
+package providers
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/goccy/go-json"
+
+	schemas "github.com/maximhq/bifrost/core/schemas"
+	"github.com/valyala/fasthttp"
+)
+
+// MistralResponse represents the response structure from the Mistral API.
+type MistralResponse struct {
+	ID      string                          `json:"id"`
+	Object  string                          `json:"object"`
+	Choices []schemas.BifrostResponseChoice `json:"choices"`
+	Model   string                          `json:"model"`
+	Created int                             `json:"created"`
+	Usage   schemas.LLMUsage                `json:"usage"`
+}
+
+// mistralResponsePool provides a pool for Mistral response objects.
+var mistralResponsePool = sync.Pool{
+	New: func() interface{} {
+		return &MistralResponse{}
+	},
+}
+
+// acquireMistralResponse gets a Mistral response from the pool and resets it.
+func acquireMistralResponse() *MistralResponse {
+	resp := mistralResponsePool.Get().(*MistralResponse)
+	*resp = MistralResponse{} // Reset the struct
+	return resp
+}
+
+// releaseMistralResponse returns a Mistral response to the pool.
+func releaseMistralResponse(resp *MistralResponse) {
+	if resp != nil {
+		mistralResponsePool.Put(resp)
+	}
+}
+
+// MistralProvider implements the Provider interface for Mistral's API.
+type MistralProvider struct {
+	logger  schemas.Logger   // Logger for provider operations
+	client  *fasthttp.Client // HTTP client for API requests
+	baseURL string           // Base URL for the provider
+}
+
+// NewMistralProvider creates a new Mistral provider instance.
+// It initializes the HTTP client with the provided configuration and sets up response pools.
+// The client is configured with timeouts, concurrency limits, and optional proxy settings.
+func NewMistralProvider(config *schemas.ProviderConfig, logger schemas.Logger) *MistralProvider {
+	config.CheckAndSetDefaults()
+
+	client := &fasthttp.Client{
+		ReadTimeout:     time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
+		WriteTimeout:    time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
+		MaxConnsPerHost: config.ConcurrencyAndBufferSize.Concurrency,
+	}
+
+	// Pre-warm response pools
+	for range config.ConcurrencyAndBufferSize.Concurrency {
+		mistralResponsePool.Put(&MistralResponse{})
+		bifrostResponsePool.Put(&schemas.BifrostResponse{})
+	}
+
+	// Configure proxy if provided
+	client = configureProxy(client, config.ProxyConfig, logger)
+
+	baseURL := strings.TrimRight(config.NetworkConfig.BaseURL, "/")
+	if baseURL == "" {
+		baseURL = "https://api.mistral.ai"
+	}
+
+	return &MistralProvider{
+		logger:  logger,
+		client:  client,
+		baseURL: baseURL,
+	}
+}
+
+// GetProviderKey returns the provider identifier for Mistral.
+func (provider *MistralProvider) GetProviderKey() schemas.ModelProvider {
+	return schemas.Mistral
+}
+
+// TextCompletion is not supported by the Mistral provider.
+func (provider *MistralProvider) TextCompletion(ctx context.Context, model, key, text string, params *schemas.ModelParameters) (*schemas.BifrostResponse, *schemas.BifrostError) {
+	return nil, &schemas.BifrostError{
+		IsBifrostError: false,
+		Error: schemas.ErrorField{
+			Message: "text completion is not supported by mistral provider",
+		},
+	}
+}
+
+// ChatCompletion performs a chat completion request to the Mistral API.
+func (provider *MistralProvider) ChatCompletion(ctx context.Context, model, key string, messages []schemas.BifrostMessage, params *schemas.ModelParameters) (*schemas.BifrostResponse, *schemas.BifrostError) {
+	formattedMessages, preparedParams := prepareOpenAIChatRequest(messages, params)
+
+	requestBody := mergeConfig(map[string]interface{}{
+		"model":    model,
+		"messages": formattedMessages,
+	}, preparedParams)
+
+	jsonBody, err := json.Marshal(requestBody)
+	if err != nil {
+		return nil, &schemas.BifrostError{
+			IsBifrostError: true,
+			Error: schemas.ErrorField{
+				Message: schemas.ErrProviderJSONMarshaling,
+				Error:   err,
+			},
+		}
+	}
+
+	// Create request
+	req := fasthttp.AcquireRequest()
+	resp := fasthttp.AcquireResponse()
+	defer fasthttp.ReleaseRequest(req)
+	defer fasthttp.ReleaseResponse(resp)
+
+	req.SetRequestURI(provider.baseURL + "/v1/chat/completions")
+	req.Header.SetMethod("POST")
+	req.Header.SetContentType("application/json")
+	req.Header.Set("Authorization", "Bearer "+key)
+	req.SetBody(jsonBody)
+
+	// Make request
+	bifrostErr := makeRequestWithContext(ctx, provider.client, req, resp)
+	if bifrostErr != nil {
+		return nil, bifrostErr
+	}
+
+	// Handle error response
+	if resp.StatusCode() != fasthttp.StatusOK {
+		provider.logger.Debug(fmt.Sprintf("error from mistral provider: %s", string(resp.Body())))
+
+		var errorResp map[string]interface{}
+		bifrostErr := handleProviderAPIError(resp, &errorResp)
+		bifrostErr.Error.Message = fmt.Sprintf("Mistral error: %v", errorResp)
+		return nil, bifrostErr
+	}
+
+	responseBody := resp.Body()
+
+	// Pre-allocate response structs from pools
+	response := acquireMistralResponse()
+	defer releaseMistralResponse(response)
+
+	result := acquireBifrostResponse()
+	defer releaseBifrostResponse(result)
+
+	// Use enhanced response handler with pre-allocated response
+	rawResponse, bifrostErr := handleProviderResponse(responseBody, response)
+	if bifrostErr != nil {
+		return nil, bifrostErr
+	}
+
+	// Populate result from response
+	result.ID = response.ID
+	result.Choices = response.Choices
+	result.Object = response.Object
+	result.Usage = response.Usage
+	result.Model = response.Model
+	result.Created = response.Created
+	result.ExtraFields = schemas.BifrostResponseExtraFields{
+		Provider:    schemas.Mistral,
+		RawResponse: rawResponse,
+	}
+
+	return result, nil
+}

core/providers/openai.go

@@ -79,7 +79,7 @@ func NewOpenAIProvider(config *schemas.ProviderConfig, logger schemas.Logger) *O
 	client := &fasthttp.Client{
 		ReadTimeout:     time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
 		WriteTimeout:    time.Second * time.Duration(config.NetworkConfig.DefaultRequestTimeoutInSeconds),
-		MaxConnsPerHost: config.ConcurrencyAndBufferSize.BufferSize,
+		MaxConnsPerHost: config.ConcurrencyAndBufferSize.Concurrency,
 	}
 
 	// Pre-warm response pools

core/schemas/bifrost.go

@@ -42,6 +42,7 @@ const (
 	Bedrock   ModelProvider = "bedrock"
 	Cohere    ModelProvider = "cohere"
 	Vertex    ModelProvider = "vertex"
+	Mistral   ModelProvider = "mistral"
 )
 
 //* Request Structs
@@ -138,12 +139,62 @@ type ToolChoiceFunction struct {
 	Name string `json:"name"` // Name of the function to call
 }
 
-// ToolChoice represents how a tool should be chosen for a request.
-type ToolChoice struct {
+// ToolChoiceStruct represents a specific tool choice.
+type ToolChoiceStruct struct {
 	Type     ToolChoiceType     `json:"type"`               // Type of tool choice
 	Function ToolChoiceFunction `json:"function,omitempty"` // Function to call if type is ToolChoiceTypeFunction
 }
 
+// ToolChoice represents how a tool should be chosen for a request. (either a string or a struct)
+type ToolChoice struct {
+	ToolChoiceStr    *string
+	ToolChoiceStruct *ToolChoiceStruct
+}
+
+// MarshalJSON implements custom JSON marshalling for ToolChoice.
+// It marshals either ToolChoiceStr or ToolChoiceStruct directly without wrapping.
+func (tc ToolChoice) MarshalJSON() ([]byte, error) {
+	// Validation: ensure only one field is set at a time
+	if tc.ToolChoiceStr != nil && tc.ToolChoiceStruct != nil {
+		return nil, fmt.Errorf("both ToolChoiceStr and ToolChoiceStruct are set; only one should be non-nil")
+	}
+
+	if tc.ToolChoiceStr != nil {
+		return json.Marshal(*tc.ToolChoiceStr)
+	}
+	if tc.ToolChoiceStruct != nil {
+		return json.Marshal(*tc.ToolChoiceStruct)
+	}
+	// If both are nil, return null
+	return json.Marshal(nil)
+}
+
+// UnmarshalJSON implements custom JSON unmarshalling for ToolChoice.
+// It determines whether "tool_choice" is a string or struct and assigns to the appropriate field.
+// It also handles direct string/array content without a wrapper object.
+func (tc *ToolChoice) UnmarshalJSON(data []byte) error {
+	// First, try to unmarshal as a direct string
+	var stringContent string
+	if err := json.Unmarshal(data, &stringContent); err == nil {
+		tc.ToolChoiceStr = &stringContent
+		return nil
+	}
+
+	// Try to unmarshal as a direct struct of ToolChoiceStruct
+	var toolChoiceStruct ToolChoiceStruct
+	if err := json.Unmarshal(data, &toolChoiceStruct); err == nil {
+		// Validate the Type field is not empty and is a valid value
+		if toolChoiceStruct.Type == "" {
+			return fmt.Errorf("tool_choice struct has empty type field")
+		}
+
+		tc.ToolChoiceStruct = &toolChoiceStruct
+		return nil
+	}
+
+	return fmt.Errorf("tool_choice field is neither a string nor a struct")
+}
+
 // BifrostMessage represents a message in a chat conversation.
 type BifrostMessage struct {
 	Role    ModelChatMessageRole `json:"role"`

core/schemas/meta/azure.go

@@ -11,31 +11,6 @@ type AzureMetaConfig struct {
 	APIVersion  *string           `json:"api_version,omitempty"` // Azure API version to use; defaults to "2024-02-01"
 }
 
-// This is not used for Azure.
-func (c *AzureMetaConfig) GetSecretAccessKey() *string {
-	return nil
-}
-
-// This is not used for Azure.
-func (c *AzureMetaConfig) GetRegion() *string {
-	return nil
-}
-
-// This is not used for Azure.
-func (c *AzureMetaConfig) GetSessionToken() *string {
-	return nil
-}
-
-// This is not used for Azure.
-func (c *AzureMetaConfig) GetARN() *string {
-	return nil
-}
-
-// This is not used for Azure.
-func (c *AzureMetaConfig) GetInferenceProfiles() map[string]string {
-	return nil
-}
-
 // GetEndpoint returns the Azure service endpoint.
 // This specifies the base URL for Azure API requests.
 func (c *AzureMetaConfig) GetEndpoint() *string {
@@ -55,12 +30,11 @@ func (c *AzureMetaConfig) GetAPIVersion() *string {
 	return c.APIVersion
 }
 
-// This is not used for Azure.
-func (c *AzureMetaConfig) GetProjectID() *string {
-	return nil
-}
-
-// This is not used for Azure.
-func (c *AzureMetaConfig) GetAuthCredentials() *string {
-	return nil
-}
+// These are not used for Azure.
+func (c *AzureMetaConfig) GetARN() *string                         { return nil }
+func (c *AzureMetaConfig) GetAuthCredentials() *string             { return nil }
+func (c *AzureMetaConfig) GetInferenceProfiles() map[string]string { return nil }
+func (c *AzureMetaConfig) GetProjectID() *string                   { return nil }
+func (c *AzureMetaConfig) GetRegion() *string                      { return nil }
+func (c *AzureMetaConfig) GetSecretAccessKey() *string             { return nil }
+func (c *AzureMetaConfig) GetSessionToken() *string                { return nil }