Class RandomPool

llms/api_key_pool.py:240–266

class RandomPool(Pool):
    def __init__(self, client_args: ClientArgs) -> None:
        super().__init__(client_args)
        logger.info(f"initializing RandomPool with {len(client_args)} api keys")
        logger.info(f"client args: {client_args}")

    def __enter__(self, *args, **kwargs) -> ClientArgs:
        """Get a client from the pool and return it as a context manager."""

        if len(self) == 0:
            raise APIKeyPoolEmptyError("api key pool is empty")

        self._index = random.randint(0, len(self) - 1)
        return self.queue[self._index].client_args

    def __exit__(self, exc_type, exc_value, exc_traceback) -> bool:
        """Put the resource back in the pool."""

        if exc_value is not None:
            if (
                exc_type is RateLimitError
                and exc_value.status_code == 429
                and judge_quota_exceeded(exc_value.message)
            ):
                self._remove_cur_arg()
            return False
        return True


# class SpinPool(metaclass=Singleton):