diff --git a/core/bifrost.go b/core/bifrost.go index 81346351de..67ddcb09fa 100644 --- a/core/bifrost.go +++ b/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) } diff --git a/core/providers/anthropic.go b/core/providers/anthropic.go index 4afe105230..85e4ce7116 100644 --- a/core/providers/anthropic.go +++ b/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, + } } } } diff --git a/core/providers/azure.go b/core/providers/azure.go index b85f436720..1b74927e91 100644 --- a/core/providers/azure.go +++ b/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 diff --git a/core/providers/cohere.go b/core/providers/cohere.go index 6ea1414ab9..df14263cb9 100644 --- a/core/providers/cohere.go +++ b/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 diff --git a/core/providers/mistral.go b/core/providers/mistral.go new file mode 100644 index 0000000000..8c9d677974 --- /dev/null +++ b/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 +} diff --git a/core/providers/openai.go b/core/providers/openai.go index c7601753c8..105806d248 100644 --- a/core/providers/openai.go +++ b/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 diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index a594007c6f..661591ee7f 100644 --- a/core/schemas/bifrost.go +++ b/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"` diff --git a/core/schemas/meta/azure.go b/core/schemas/meta/azure.go index 15e0f27d8e..caaa53b945 100644 --- a/core/schemas/meta/azure.go +++ b/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 } diff --git a/core/schemas/meta/bedrock.go b/core/schemas/meta/bedrock.go index d645a4a2d8..bdff19e76a 100644 --- a/core/schemas/meta/bedrock.go +++ b/core/schemas/meta/bedrock.go @@ -43,27 +43,9 @@ func (c *BedrockMetaConfig) GetInferenceProfiles() map[string]string { return c.InferenceProfiles } -// This is not used for Bedrock. -func (c *BedrockMetaConfig) GetEndpoint() *string { - return nil -} - -// This is not used for Bedrock. -func (c *BedrockMetaConfig) GetDeployments() map[string]string { - return nil -} - -// This is not used for Bedrock. -func (c *BedrockMetaConfig) GetAPIVersion() *string { - return nil -} - -// This is not used for Bedrock. -func (c *BedrockMetaConfig) GetProjectID() *string { - return nil -} - -// This is not used for Bedrock. -func (c *BedrockMetaConfig) GetAuthCredentials() *string { - return nil -} +// These are not used for Bedrock. +func (c *BedrockMetaConfig) GetAPIVersion() *string { return nil } +func (c *BedrockMetaConfig) GetAuthCredentials() *string { return nil } +func (c *BedrockMetaConfig) GetDeployments() map[string]string { return nil } +func (c *BedrockMetaConfig) GetEndpoint() *string { return nil } +func (c *BedrockMetaConfig) GetProjectID() *string { return nil } diff --git a/core/schemas/meta/vertex.go b/core/schemas/meta/vertex.go index 1a9e71d78b..a82e46380d 100644 --- a/core/schemas/meta/vertex.go +++ b/core/schemas/meta/vertex.go @@ -11,47 +11,12 @@ type VertexMetaConfig struct { AuthCredentials string `json:"auth_credentials,omitempty"` } -// This is not used for Vertex. -func (c *VertexMetaConfig) GetSecretAccessKey() *string { - return nil -} - // GetRegion returns the Vertex region. // This is the region for the Vertex project. func (c *VertexMetaConfig) GetRegion() *string { return &c.Region } -// This is not used for Vertex. -func (c *VertexMetaConfig) GetSessionToken() *string { - return nil -} - -// This is not used for Vertex. -func (c *VertexMetaConfig) GetARN() *string { - return nil -} - -// This is not used for Vertex. -func (c *VertexMetaConfig) GetInferenceProfiles() map[string]string { - return nil -} - -// This is not used for Vertex. -func (c *VertexMetaConfig) GetEndpoint() *string { - return nil -} - -// This is not used for Vertex. -func (c *VertexMetaConfig) GetDeployments() map[string]string { - return nil -} - -// This is not used for Vertex. -func (c *VertexMetaConfig) GetAPIVersion() *string { - return nil -} - // GetProjectID returns the Vertex project ID. // This is the project ID for the Vertex project. func (c *VertexMetaConfig) GetProjectID() *string { @@ -63,3 +28,12 @@ func (c *VertexMetaConfig) GetProjectID() *string { func (c *VertexMetaConfig) GetAuthCredentials() *string { return &c.AuthCredentials } + +// These are not used for Vertex. +func (c *VertexMetaConfig) GetAPIVersion() *string { return nil } +func (c *VertexMetaConfig) GetARN() *string { return nil } +func (c *VertexMetaConfig) GetDeployments() map[string]string { return nil } +func (c *VertexMetaConfig) GetEndpoint() *string { return nil } +func (c *VertexMetaConfig) GetInferenceProfiles() map[string]string { return nil } +func (c *VertexMetaConfig) GetSecretAccessKey() *string { return nil } +func (c *VertexMetaConfig) GetSessionToken() *string { return nil } diff --git a/core/tests/account.go b/core/tests/account.go index d5652a3c14..a46dab51d5 100644 --- a/core/tests/account.go +++ b/core/tests/account.go @@ -27,7 +27,7 @@ type BaseAccount struct{} // - []schemas.SupportedModelProvider: A slice containing the supported provider identifiers // - error: Always returns nil as this implementation doesn't produce errors func (baseAccount *BaseAccount) GetConfiguredProviders() ([]schemas.ModelProvider, error) { - return []schemas.ModelProvider{schemas.OpenAI, schemas.Anthropic, schemas.Bedrock, schemas.Cohere, schemas.Azure, schemas.Vertex}, nil + return []schemas.ModelProvider{schemas.OpenAI, schemas.Anthropic, schemas.Bedrock, schemas.Cohere, schemas.Azure, schemas.Vertex, schemas.Mistral}, nil } // GetKeysForProvider returns the API keys and associated models for a given provider. @@ -89,6 +89,14 @@ func (baseAccount *BaseAccount) GetKeysForProvider(providerKey schemas.ModelProv Weight: 1.0, }, }, nil + case schemas.Mistral: + return []schemas.Key{ + { + Value: os.Getenv("MISTRAL_API_KEY"), + Models: []string{"mistral-large-2411", "ministral-3b-2410", "pixtral-12b-latest"}, + Weight: 1.0, + }, + }, nil default: return nil, fmt.Errorf("unsupported provider: %s", providerKey) } @@ -199,6 +207,11 @@ func (baseAccount *BaseAccount) GetConfigForProvider(providerKey schemas.ModelPr BufferSize: 10, }, }, nil + case schemas.Mistral: + return &schemas.ProviderConfig{ + NetworkConfig: schemas.DefaultNetworkConfig, + ConcurrencyAndBufferSize: schemas.DefaultConcurrencyAndBufferSize, + }, nil default: return nil, fmt.Errorf("unsupported provider: %s", providerKey) } diff --git a/core/tests/e2e_tool_test.go b/core/tests/e2e_tool_test.go index 902818b1c7..92ffd2b861 100644 --- a/core/tests/e2e_tool_test.go +++ b/core/tests/e2e_tool_test.go @@ -41,9 +41,11 @@ func TestToolCallingEndToEnd(t *testing.T) { toolParams := WeatherToolParams toolParams.ToolChoice = &schemas.ToolChoice{ - Type: schemas.ToolChoiceTypeFunction, - Function: schemas.ToolChoiceFunction{ - Name: "get_weather", + ToolChoiceStruct: &schemas.ToolChoiceStruct{ + Type: schemas.ToolChoiceTypeFunction, + Function: schemas.ToolChoiceFunction{ + Name: "get_weather", + }, }, } toolParams.MaxTokens = bifrost.Ptr(1000) diff --git a/core/tests/mistral_test.go b/core/tests/mistral_test.go new file mode 100644 index 0000000000..7e59d6fce7 --- /dev/null +++ b/core/tests/mistral_test.go @@ -0,0 +1,36 @@ +// Package tests provides test utilities and configurations for the Bifrost system. +// It includes test implementations of schemas, mock objects, and helper functions +// for testing the Bifrost functionality with various AI providers. +package tests + +import ( + "testing" + + schemas "github.com/maximhq/bifrost/core/schemas" +) + +func TestMistral(t *testing.T) { + bifrost, err := getBifrost() + if err != nil { + t.Fatalf("Error initializing bifrost: %v", err) + return + } + + config := TestConfig{ + Provider: schemas.Mistral, + TextModel: "ministral-3b-2410", + ChatModel: "pixtral-12b-latest", + SetupText: false, // Mistral does not support text completion + SetupToolCalls: true, + SetupImage: true, + SetupBaseImage: true, + Fallbacks: []schemas.Fallback{ + { + Provider: schemas.Anthropic, + Model: "claude-3-7-sonnet-20250219", + }, + }, + } + + SetupAllRequests(bifrost, config) +}