Rui
51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
import redis
|
|
import json
|
|
import random
|
|
|
|
class KeyManager:
|
|
def __init__(self):
|
|
self.r = redis.Redis(
|
|
host='10.99.0.3',
|
|
port=6379,
|
|
password='clube67_dragonfly_pass_9903',
|
|
db=0,
|
|
decode_responses=True
|
|
)
|
|
self.pool_key = 'swarm:keys:pool'
|
|
|
|
def get_key(self, provider):
|
|
"""Busca uma chave ativa para o provedor (google ou anthropic)"""
|
|
pool_data = self.r.get(self.pool_key)
|
|
if not pool_data:
|
|
return None
|
|
|
|
pool = json.loads(pool_data)
|
|
keys = pool.get(provider, [])
|
|
|
|
# Filtra apenas as chaves ativas
|
|
active_keys = [k for k in keys if k.get('status') == 'active']
|
|
|
|
if not active_keys:
|
|
print(f" [!] Erro: Nenhuma chave ativa para {provider}")
|
|
return None
|
|
|
|
# Seleciona uma chave aleatória do pool ativo (Load Balancing)
|
|
selected = random.choice(active_keys)
|
|
return selected['key']
|
|
|
|
def report_failure(self, provider, failed_key):
|
|
"""Marca uma chave como inativa caso falhe (ex: sem saldo)"""
|
|
pool_data = self.r.get(self.pool_key)
|
|
pool = json.loads(pool_data)
|
|
|
|
for k in pool.get(provider, []):
|
|
if k['key'] == failed_key:
|
|
k['status'] = 'inactive'
|
|
print(f" [W] Chave {failed_key[:10]}... marcada como INATIVA.")
|
|
|
|
self.r.set(self.pool_key, json.dumps(pool))
|
|
|
|
# Exemplo de uso:
|
|
# km = KeyManager()
|
|
# key = km.get_key('google')
|