key_manager.py 883 B

1234567891011121314151617181920212223242526272829
  1. import random
  2. def Singleton(cls):
  3. _instance = {}
  4. def _singleton(*args, **kargs):
  5. if cls not in _instance:
  6. _instance[cls] = cls(*args, **kargs)
  7. return _instance[cls]
  8. return _singleton
  9. @Singleton
  10. class OpenAI_ApiKeyManager():
  11. def __init__(self, mode='blacklist') -> None:
  12. # self.key_avail_list = []
  13. self.key_black_list = []
  14. def add_key_to_blacklist(self, key):
  15. self.key_black_list.append(key)
  16. def select_avail_key(self, key_list):
  17. # select key from key_list, but avoid keys also in self.key_black_list, raise error if no key can be found
  18. available_keys = [key for key in key_list if key not in self.key_black_list]
  19. if not available_keys:
  20. raise KeyError("No available key found.")
  21. selected_key = random.choice(available_keys)
  22. return selected_key