feat: core extended

Author: Pratham-Mishra04

core/bifrost.go

@@ -46,6 +46,7 @@ type Bifrost struct {
 	logger              schemas.Logger                   // logger instance, default logger is used if not provided
 	mcpManager          *MCPManager                      // MCP integration manager (nil if MCP not configured)
 	dropExcessRequests  atomic.Bool                      // If true, in cases where the queue is full, requests will not wait for the queue to be empty and will be dropped instead.
+	keySelector         schemas.KeySelector              // Custom key selector function
 }
 
 // PluginPipeline encapsulates the execution of plugin PreHooks and PostHooks, tracks how many plugins ran, and manages short-circuiting and error aggregation.
@@ -86,10 +87,15 @@ func Init(ctx context.Context, config schemas.BifrostConfig) (*Bifrost, error) {
 		plugins:       atomic.Pointer[[]schemas.Plugin]{},
 		requestQueues: sync.Map{},
 		waitGroups:    sync.Map{},
+		keySelector:   config.KeySelector,
 	}
 	bifrost.plugins.Store(&config.Plugins)
 	bifrost.dropExcessRequests.Store(config.DropExcessRequests)
 
+	if bifrost.keySelector == nil {
+		bifrost.keySelector = WeightedRandomKeySelector
+	}
+
 	// Initialize object pools
 	bifrost.channelMessagePool = sync.Pool{
 		New: func() interface{} {
@@ -626,12 +632,12 @@ func (bifrost *Bifrost) UpdateProviderConcurrency(providerKey schemas.ModelProvi
 		return bifrost.prepareProvider(providerKey, providerConfig)
 	}
 
-	oldQueue := oldQueueValue.(chan ChannelMessage)
+	oldQueue := oldQueueValue.(chan *ChannelMessage)
 
 	bifrost.logger.Debug("gracefully stopping existing workers for provider %s", providerKey)
 
 	// Step 1: Create new queue with updated buffer size
-	newQueue := make(chan ChannelMessage, providerConfig.ConcurrencyAndBufferSize.BufferSize)
+	newQueue := make(chan *ChannelMessage, providerConfig.ConcurrencyAndBufferSize.BufferSize)
 
 	// Step 2: Transfer any buffered requests from old queue to new queue
 	// This prevents request loss during the transition
@@ -647,7 +653,7 @@ func (bifrost *Bifrost) UpdateProviderConcurrency(providerKey schemas.ModelProvi
 				// New queue is full, handle this request in a goroutine
 				// This is unlikely with proper buffer sizing but provides safety
 				transferWaitGroup.Add(1)
-				go func(m ChannelMessage) {
+				go func(m *ChannelMessage) {
 					defer transferWaitGroup.Done()
 					select {
 					case newQueue <- m:
@@ -1011,7 +1017,7 @@ func (bifrost *Bifrost) prepareProvider(providerKey schemas.ModelProvider, confi
 		return fmt.Errorf("failed to get config for provider: %v", err)
 	}
 
-	queue := make(chan ChannelMessage, providerConfig.ConcurrencyAndBufferSize.BufferSize) // Buffered channel per provider
+	queue := make(chan *ChannelMessage, providerConfig.ConcurrencyAndBufferSize.BufferSize) // Buffered channel per provider
 
 	bifrost.requestQueues.Store(providerKey, queue)
 
@@ -1038,13 +1044,13 @@ func (bifrost *Bifrost) prepareProvider(providerKey schemas.ModelProvider, confi
 // If the queue doesn't exist, it creates one at runtime and initializes the provider,
 // given the provider config is provided in the account interface implementation.
 // This function uses read locks to prevent race conditions during provider updates.
-func (bifrost *Bifrost) getProviderQueue(providerKey schemas.ModelProvider) (chan ChannelMessage, error) {
+func (bifrost *Bifrost) getProviderQueue(providerKey schemas.ModelProvider) (chan *ChannelMessage, error) {
 	// Use read lock to allow concurrent reads but prevent concurrent updates
 	providerMutex := bifrost.getProviderMutex(providerKey)
 	providerMutex.RLock()
 
 	if queueValue, exists := bifrost.requestQueues.Load(providerKey); exists {
-		queue := queueValue.(chan ChannelMessage)
+		queue := queueValue.(chan *ChannelMessage)
 		providerMutex.RUnlock()
 		return queue, nil
 	}
@@ -1057,7 +1063,7 @@ func (bifrost *Bifrost) getProviderQueue(providerKey schemas.ModelProvider) (cha
 
 	// Double-check after acquiring write lock (another goroutine might have created it)
 	if queueValue, exists := bifrost.requestQueues.Load(providerKey); exists {
-		queue := queueValue.(chan ChannelMessage)
+		queue := queueValue.(chan *ChannelMessage)
 		return queue, nil
 	}
 
@@ -1073,7 +1079,7 @@ func (bifrost *Bifrost) getProviderQueue(providerKey schemas.ModelProvider) (cha
 	}
 
 	queueValue, _ := bifrost.requestQueues.Load(providerKey)
-	queue := queueValue.(chan ChannelMessage)
+	queue := queueValue.(chan *ChannelMessage)
 
 	return queue, nil
 }
@@ -1335,9 +1341,8 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont
 
 	msg := bifrost.getChannelMessage(*preReq)
 	msg.Context = ctx
-	startTime := time.Now()
 	select {
-	case queue <- *msg:
+	case queue <- msg:
 		// Message was sent successfully
 	case <-ctx.Done():
 		bifrost.releaseChannelMessage(msg)
@@ -1349,7 +1354,7 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont
 			return nil, newBifrostErrorFromMsg("request dropped: queue is full")
 		}
 		select {
-		case queue <- *msg:
+		case queue <- msg:
 			// Message was sent successfully
 		case <-ctx.Done():
 			bifrost.releaseChannelMessage(msg)
@@ -1362,11 +1367,7 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont
 	pluginCount := len(*bifrost.plugins.Load())
 	select {
 	case result = <-msg.Response:
-		latency := time.Since(startTime).Milliseconds()
-		if result.ExtraFields.Latency == nil {
-			result.ExtraFields.Latency = Ptr(float64(latency))
-		}
-		resp, bifrostErr := pipeline.RunPostHooks(&ctx, result, nil, pluginCount)
+		resp, bifrostErr := pipeline.RunPostHooks(&msg.Context, result, nil, pluginCount)
 		if bifrostErr != nil {
 			bifrost.releaseChannelMessage(msg)
 			return nil, bifrostErr
@@ -1375,7 +1376,7 @@ func (bifrost *Bifrost) tryRequest(req *schemas.BifrostRequest, ctx context.Cont
 		return resp, nil
 	case bifrostErrVal := <-msg.Err:
 		bifrostErrPtr := &bifrostErrVal
-		resp, bifrostErrPtr = pipeline.RunPostHooks(&ctx, nil, bifrostErrPtr, pluginCount)
+		resp, bifrostErrPtr = pipeline.RunPostHooks(&msg.Context, nil, bifrostErrPtr, pluginCount)
 		bifrost.releaseChannelMessage(msg)
 		if bifrostErrPtr != nil {
 			return nil, bifrostErrPtr
@@ -1457,7 +1458,7 @@ func (bifrost *Bifrost) tryStreamRequest(req *schemas.BifrostRequest, ctx contex
 	msg.Context = ctx
 
 	select {
-	case queue <- *msg:
+	case queue <- msg:
 		// Message was sent successfully
 	case <-ctx.Done():
 		bifrost.releaseChannelMessage(msg)
@@ -1469,7 +1470,7 @@ func (bifrost *Bifrost) tryStreamRequest(req *schemas.BifrostRequest, ctx contex
 			return nil, newBifrostErrorFromMsg("request dropped: queue is full")
 		}
 		select {
-		case queue <- *msg:
+		case queue <- msg:
 			// Message was sent successfully
 		case <-ctx.Done():
 			bifrost.releaseChannelMessage(msg)
@@ -1500,7 +1501,7 @@ func (bifrost *Bifrost) tryStreamRequest(req *schemas.BifrostRequest, ctx contex
 
 // requestWorker handles incoming requests from the queue for a specific provider.
 // It manages retries, error handling, and response processing.
-func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas.ProviderConfig, queue chan ChannelMessage) {
+func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas.ProviderConfig, queue chan *ChannelMessage) {
 	defer func() {
 		if waitGroupValue, ok := bifrost.waitGroups.Load(provider.GetProviderKey()); ok {
 			waitGroup := waitGroupValue.(*sync.WaitGroup)
@@ -1535,6 +1536,7 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
 				}
 				continue
 			}
+			req.Context = context.WithValue(req.Context, schemas.BifrostContextKeySelectedKey, key.ID)
 		}
 
 		// Track attempts
@@ -1570,12 +1572,12 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
 
 			// Attempt the request
 			if IsStreamRequestType(req.RequestType) {
-				stream, bifrostError = handleProviderStreamRequest(provider, &req, key, postHookRunner)
+				stream, bifrostError = handleProviderStreamRequest(provider, req, key, postHookRunner)
 				if bifrostError != nil && !bifrostError.IsBifrostError {
 					break // Don't retry client errors
 				}
 			} else {
-				result, bifrostError = handleProviderRequest(provider, &req, key)
+				result, bifrostError = handleProviderRequest(provider, req, key)
 				if bifrostError != nil {
 					break // Don't retry client errors
 				}
@@ -1924,9 +1926,19 @@ func (bifrost *Bifrost) selectKeyFromProviderForModel(ctx *context.Context, prov
 		return supportedKeys[0], nil
 	}
 
+	selectedKey, err := bifrost.keySelector(ctx, supportedKeys, providerKey, model)
+	if err != nil {
+		return schemas.Key{}, err
+	}
+
+	return selectedKey, nil
+
+}
+
+func WeightedRandomKeySelector(ctx *context.Context, keys []schemas.Key, providerKey schemas.ModelProvider, model string) (schemas.Key, error) {
 	// Use a weighted random selection based on key weights
 	totalWeight := 0
-	for _, key := range supportedKeys {
+	for _, key := range keys {
 		totalWeight += int(key.Weight * 100) // Convert float to int for better performance
 	}
 
@@ -1936,15 +1948,15 @@ func (bifrost *Bifrost) selectKeyFromProviderForModel(ctx *context.Context, prov
 
 	// Select key based on weight
 	currentWeight := 0
-	for _, key := range supportedKeys {
+	for _, key := range keys {
 		currentWeight += int(key.Weight * 100)
 		if randomValue < currentWeight {
 			return key, nil
 		}
 	}
 
 	// Fallback to first key if something goes wrong
-	return supportedKeys[0], nil
+	return keys[0], nil
 }
 
 // Shutdown gracefully stops all workers when triggered.
@@ -1954,7 +1966,7 @@ func (bifrost *Bifrost) Shutdown() {
 
 	// Close all provider queues to signal workers to stop
 	bifrost.requestQueues.Range(func(key, value interface{}) bool {
-		close(value.(chan ChannelMessage))
+		close(value.(chan *ChannelMessage))
 		return true
 	})
 

core/providers/anthropic.go

@@ -124,11 +124,11 @@ func (provider *AnthropicProvider) GetProviderKey() schemas.ModelProvider {
 // completeRequest sends a request to Anthropic's API and handles the response.
 // It constructs the API URL, sets up authentication, and processes the response.
 // Returns the response body or an error if the request fails.
-func (provider *AnthropicProvider) completeRequest(ctx context.Context, requestBody interface{}, url string, key string) ([]byte, *schemas.BifrostError) {
+func (provider *AnthropicProvider) completeRequest(ctx context.Context, requestBody interface{}, url string, key string) ([]byte, time.Duration, *schemas.BifrostError) {
 	// Marshal the request body
 	jsonData, err := sonic.Marshal(requestBody)
 	if err != nil {
-		return nil, newBifrostOperationError(schemas.ErrProviderJSONMarshaling, err, provider.GetProviderKey())
+		return nil, 0, newBifrostOperationError(schemas.ErrProviderJSONMarshaling, err, provider.GetProviderKey())
 	}
 
 	// Create the request with the JSON body
@@ -149,9 +149,9 @@ func (provider *AnthropicProvider) completeRequest(ctx context.Context, requestB
 	req.SetBody(jsonData)
 
 	// Send the request
-	bifrostErr := makeRequestWithContext(ctx, provider.client, req, resp)
+	latency, bifrostErr := makeRequestWithContext(ctx, provider.client, req, resp)
 	if bifrostErr != nil {
-		return nil, bifrostErr
+		return nil, latency, bifrostErr
 	}
 
 	// Handle error response
@@ -164,13 +164,13 @@ func (provider *AnthropicProvider) completeRequest(ctx context.Context, requestB
 		bifrostErr.Error.Type = &errorResp.Error.Type
 		bifrostErr.Error.Message = errorResp.Error.Message
 
-		return nil, bifrostErr
+		return nil, latency, bifrostErr
 	}
 
 	// Read the response body
 	body := resp.Body()
 
-	return body, nil
+	return body, latency, nil
 }
 
 // TextCompletion performs a text completion request to Anthropic's API.
@@ -188,7 +188,7 @@ func (provider *AnthropicProvider) TextCompletion(ctx context.Context, key schem
 	}
 
 	// Use struct directly for JSON marshaling
-	responseBody, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+"/v1/complete", key.Value)
+	responseBody, latency, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+"/v1/complete", key.Value)
 	if err != nil {
 		return nil, err
 	}
@@ -208,6 +208,7 @@ func (provider *AnthropicProvider) TextCompletion(ctx context.Context, key schem
 	bifrostResponse.ExtraFields.Provider = provider.GetProviderKey()
 	bifrostResponse.ExtraFields.ModelRequested = request.Model
 	bifrostResponse.ExtraFields.RequestType = schemas.TextCompletionRequest
+	bifrostResponse.ExtraFields.Latency = latency.Milliseconds()
 
 	// Set raw response if enabled
 	if provider.sendBackRawResponse {
@@ -239,7 +240,7 @@ func (provider *AnthropicProvider) ChatCompletion(ctx context.Context, key schem
 	}
 
 	// Use struct directly for JSON marshaling
-	responseBody, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+"/v1/messages", key.Value)
+	responseBody, latency, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+"/v1/messages", key.Value)
 	if err != nil {
 		return nil, err
 	}
@@ -260,6 +261,7 @@ func (provider *AnthropicProvider) ChatCompletion(ctx context.Context, key schem
 	bifrostResponse.ExtraFields.Provider = provider.GetProviderKey()
 	bifrostResponse.ExtraFields.ModelRequested = request.Model
 	bifrostResponse.ExtraFields.RequestType = schemas.ChatCompletionRequest
+	bifrostResponse.ExtraFields.Latency = latency.Milliseconds()
 
 	// Set raw response if enabled
 	if provider.sendBackRawResponse {
@@ -284,7 +286,7 @@ func (provider *AnthropicProvider) Responses(ctx context.Context, key schemas.Ke
 	}
 
 	// Use struct directly for JSON marshaling
-	responseBody, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+"/v1/messages", key.Value)
+	responseBody, latency, err := provider.completeRequest(ctx, reqBody, provider.networkConfig.BaseURL+"/v1/messages", key.Value)
 	if err != nil {
 		return nil, err
 	}
@@ -305,6 +307,7 @@ func (provider *AnthropicProvider) Responses(ctx context.Context, key schemas.Ke
 	bifrostResponse.ExtraFields.Provider = provider.GetProviderKey()
 	bifrostResponse.ExtraFields.ModelRequested = request.Model
 	bifrostResponse.ExtraFields.RequestType = schemas.ResponsesRequest
+	bifrostResponse.ExtraFields.Latency = latency.Milliseconds()
 
 	// Set raw response if enabled
 	if provider.sendBackRawResponse {