Class Pool

llms/api_key_pool.py:164–203

class Pool:
    def __init__(self, client_args: ClientArgs) -> None:
        assert "api_keys" in client_args or "api_key" in client_args
        assert not ("api_keys" in client_args and "api_key" in client_args)

        if "api_key" in client_args:
            self.queue: List[Resource] = [Resource(**client_args)]
        elif "api_keys" in client_args:
            self.queue: List[Resource] = [
                Resource(
                    api_key=api_key,
                    base_url=client_args["base_url"],
                    organization=client_args.get("organization", None),
                    project=client_args.get("project", None),
                    max_retries=client_args.get("max_retries", 2),
                    timeout=client_args.get("timeout", None),
                )
                for api_key in client_args["api_keys"]
            ]

        self._index: int = 0

    def __len__(self) -> int:
        return len(self.queue)

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

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

    def _remove_cur_arg(self) -> None:
        logger.info(f"removing api key {self.queue[self._index]}")
        _ = self.client_args.pop(self._index)
        if len(self) == 0:
            raise APIKeyPoolEmptyError("api key pool is empty")


class PQPool(Pool):