diff --git a/aixplain/factories/__init__.py b/aixplain/factories/__init__.py index 08cb8d4a..70361e77 100644 --- a/aixplain/factories/__init__.py +++ b/aixplain/factories/__init__.py @@ -30,4 +30,4 @@ from .model_factory import ModelFactory from .pipeline_factory import PipelineFactory from .finetune_factory import FinetuneFactory -from .wallet_factoy import WalletFactory +from .wallet_factory import WalletFactory diff --git a/aixplain/factories/wallet_factory.py b/aixplain/factories/wallet_factory.py new file mode 100644 index 00000000..59ec7c14 --- /dev/null +++ b/aixplain/factories/wallet_factory.py @@ -0,0 +1,26 @@ +import aixplain.utils.config as config +from aixplain.modules.wallet import Wallet +from aixplain.utils.file_utils import _request_with_retry +import logging + + +class WalletFactory: + aixplain_key = config.AIXPLAIN_API_KEY + backend_url = config.BACKEND_URL + + @classmethod + def get(cls) -> Wallet: + """Get wallet information""" + try: + resp = None + # Check for code 200, other code will be caught when trying to return a Wallet object + url = f"{cls.backend_url}/sdk/billing/wallet" + + headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"} + logging.info(f"Start fetching billing information from - {url} - {headers}") + headers = {"Content-Type": "application/json", "x-api-key": config.TEAM_API_KEY} + r = _request_with_retry("get", url, headers=headers) + resp = r.json() + return Wallet(total_balance=resp["totalBalance"], reserved_balance=resp["reservedBalance"]) + except Exception as e: + raise Exception(f"Failed to get the wallet credit information. Error: {str(e)}")