Class PQPool

llms/api_key_pool.py:206–237

class PQPool(Pool):
    def __init__(self, client_args: ClientArgs) -> None:
        super().__init__(client_args)  # unlimited size

        # Initialize the heap with resources
        heapify(self.queue)

    def __enter__(self, **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._cur_resource = heappop(self.queue)

        return self._cur_resource.client_args

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

        # TODO how to update the usage?
        self._cur_usage = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        }
        heappush(self.queue, self._cur_resource.update(self._cur_usage))
        if exc_type is not None:
            # Log the exception or handle it as needed
            logger.error(f"Client Pool Error: {exc_type}, {exc_value}")
            return False
        return True


class RandomPool(Pool):