137 lines
3.1 KiB
Go
137 lines
3.1 KiB
Go
package store
|
|
|
|
import (
|
|
"errors"
|
|
"log"
|
|
"math/rand"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/Sakurasan/to"
|
|
"github.com/patrickmn/go-cache"
|
|
)
|
|
|
|
var (
|
|
KeysCache *cache.Cache
|
|
AuthCache *cache.Cache
|
|
)
|
|
|
|
func init() {
|
|
KeysCache = cache.New(cache.NoExpiration, cache.NoExpiration)
|
|
AuthCache = cache.New(cache.NoExpiration, cache.NoExpiration)
|
|
}
|
|
|
|
func LoadKeysCache() {
|
|
KeysCache = cache.New(cache.NoExpiration, cache.NoExpiration)
|
|
keys, err := GetAllKeys()
|
|
if err != nil {
|
|
log.Println(err)
|
|
return
|
|
}
|
|
for idx, key := range keys {
|
|
KeysCache.Set(to.String(idx), key, cache.NoExpiration)
|
|
}
|
|
}
|
|
|
|
func FromKeyCacheRandomItemKey() Key {
|
|
items := KeysCache.Items()
|
|
if len(items) == 1 {
|
|
return items[to.String(0)].Object.(Key)
|
|
}
|
|
idx := rand.Intn(len(items))
|
|
item := items[to.String(idx)]
|
|
return item.Object.(Key)
|
|
}
|
|
|
|
func SelectKeyCache(apitype string) (Key, error) {
|
|
var keys []Key
|
|
items := KeysCache.Items()
|
|
for _, item := range items {
|
|
if item.Object.(Key).ApiType == apitype {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
if apitype == "openai" {
|
|
if item.Object.(Key).ApiType == "azure" {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
if item.Object.(Key).ApiType == "github" {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
}
|
|
if apitype == "claude" {
|
|
if item.Object.(Key).ApiType == "vertex" {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
}
|
|
}
|
|
if len(keys) == 0 {
|
|
return Key{}, errors.New("No key found")
|
|
} else if len(keys) == 1 {
|
|
return keys[0], nil
|
|
}
|
|
rand.Seed(time.Now().UnixNano())
|
|
idx := rand.Intn(len(keys))
|
|
return keys[idx], nil
|
|
}
|
|
|
|
func SelectKeyCacheByModel(model string) (Key, error) {
|
|
var keys []Key
|
|
if os.Getenv("OPENAI_API_KEY") != "" {
|
|
keys = append(keys, Key{ApiType: "openai", Key: os.Getenv("OPENAI_API_KEY")})
|
|
}
|
|
items := KeysCache.Items()
|
|
for _, item := range items {
|
|
if strings.Contains(model, "realtime") || strings.HasPrefix(model, "o1-") {
|
|
if item.Object.(Key).ApiType == "openai" {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
if item.Object.(Key).ApiType == "azure" {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
}
|
|
if strings.HasPrefix(model, "gpt-") {
|
|
if item.Object.(Key).ApiType == "openai" {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
if item.Object.(Key).ApiType == "azure" {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
if item.Object.(Key).ApiType == "github" {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
}
|
|
if strings.HasPrefix(model, "chatgpt-") {
|
|
if item.Object.(Key).ApiType == "openai" {
|
|
keys = append(keys, item.Object.(Key))
|
|
}
|
|
}
|
|
}
|
|
if len(keys) == 0 {
|
|
return Key{}, errors.New("No key found")
|
|
} else if len(keys) == 1 {
|
|
return keys[0], nil
|
|
}
|
|
rand.Seed(time.Now().UnixNano())
|
|
idx := rand.Intn(len(keys))
|
|
return keys[idx], nil
|
|
}
|
|
|
|
func LoadAuthCache() {
|
|
AuthCache = cache.New(cache.NoExpiration, cache.NoExpiration)
|
|
users, err := GetAllUsers()
|
|
if err != nil {
|
|
log.Println(err)
|
|
return
|
|
}
|
|
for _, user := range users {
|
|
AuthCache.Set(user.Token, true, cache.NoExpiration)
|
|
}
|
|
}
|
|
|
|
func IsExistAuthCache(auth string) bool {
|
|
items := AuthCache.Items()
|
|
_, ok := items[auth]
|
|
return ok
|
|
}
|