36 lines
933 B
Go
36 lines
933 B
Go
package keyselectors
|
|
|
|
import (
|
|
"math/rand"
|
|
|
|
"github.com/maximhq/bifrost/core/schemas"
|
|
)
|
|
|
|
func WeightedRandom(ctx *schemas.BifrostContext, 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 keys {
|
|
totalWeight += int(key.Weight * 100) // Convert float to int for better performance
|
|
}
|
|
|
|
// If all keys have zero weight, fall back to uniform random selection
|
|
if totalWeight == 0 {
|
|
return keys[rand.Intn(len(keys))], nil
|
|
}
|
|
|
|
// Use global thread-safe random (Go 1.20+) - no allocation, no syscall
|
|
randomValue := rand.Intn(totalWeight)
|
|
|
|
// Select key based on weight
|
|
currentWeight := 0
|
|
for _, key := range keys {
|
|
currentWeight += int(key.Weight * 100)
|
|
if randomValue < currentWeight {
|
|
return key, nil
|
|
}
|
|
}
|
|
|
|
// Fallback to first key if something goes wrong
|
|
return keys[0], nil
|
|
}
|