A thread-safe OpenAI api invoker which support api_key pool

openai_invoker.py

import openai
import threading
import time
from collections import OrderedDict


class OpenAIInvoker:
    __unblock_time = 60 * 1000

    Completion = openai.Completion()
    Customer = openai.Customer()
    Edit = openai.Edit()
    Embedding = openai.Embedding()

    def __init__(self, keys=[]):
        self.__keys = OrderedDict()
        self.__lock = threading.Lock()
        for key in keys:
            self.__keys[key] = 0

    def append_key(self, key):
        with self.__lock:
            if key not in self.__keys:
                self.__keys[key] = 0

    def remove_key(self, key):
        with self.__lock:
            self.__keys.pop(key)

    def __peek_key(self):
        with self.__lock:
            current_cnt = len(self.__keys)
            for _ in range(current_cnt):
                key_inf = self.__keys.popitem(False)
                key, available_time = key_inf[0], key_inf[1]
                self.__keys[key] = available_time
                if available_time < int(time.time() * 1000):
                    return key

    def __block_key(self, key):
        with self.__lock:
            if key in self.__keys:
                self.__keys[key] = int(time.time() * 1000) + OpenAIInvoker.__unblock_time

    def invoke(self, openai_module, tmp_api_key=None, **args):
        if not tmp_api_key:
            key = self.__peek_key()
            if not key:
                raise Exception("No enough available keys")
        else:
            key = tmp_api_key

        try:
            response = openai_module.create(api_key=key, **args)
            return response
        except openai.error.RateLimitError:
            if not tmp_api_key:
                self.__block_key(key)
            raise
            
# Initialize keys
api_keys = ["test", "123"]
invoker = OpenAIInvoker(api_keys)

# Invoke
args = {
    "model": "text-davinci-003",
    "prompt": "Who am I?",
    "temperature": 0.7,
    "max_tokens": 256,
    "top_p": 1,
    "frequency_penalty": 0,
    "presence_penalty": 0
}
    
try:
    response = invoker.invoke(
        OpenAIInvoker.Completion,
        **args
    )
    print(response)
except Exception as invoker_exp:
    print(invoker_exp)