diff --git a/.gitignore b/.gitignore index 843c6556..ad7c16c8 100644 --- a/.gitignore +++ b/.gitignore @@ -130,3 +130,5 @@ dmypy.json # Vscode .vscode +.DS_Store + diff --git a/aixplain/factories/__init__.py b/aixplain/factories/__init__.py index 7b876899..08cb8d4a 100644 --- a/aixplain/factories/__init__.py +++ b/aixplain/factories/__init__.py @@ -30,3 +30,4 @@ from .model_factory import ModelFactory from .pipeline_factory import PipelineFactory from .finetune_factory import FinetuneFactory +from .wallet_factoy import WalletFactory diff --git a/aixplain/factories/wallet_factoy.py b/aixplain/factories/wallet_factoy.py new file mode 100644 index 00000000..59ec7c14 --- /dev/null +++ b/aixplain/factories/wallet_factoy.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)}") diff --git a/aixplain/modules/wallet.py b/aixplain/modules/wallet.py new file mode 100644 index 00000000..d7c63524 --- /dev/null +++ b/aixplain/modules/wallet.py @@ -0,0 +1,34 @@ +__author__ = "aixplain" + +""" +Copyright 2024 The aiXplain SDK authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Author: aiXplain Team +Date: August 20th 2024 +Description: + Wallet Class +""" + + +class Wallet: + def __init__(self, total_balance: float, reserved_balance: float): + """Create a Wallet with the necessary information + + Args: + total_balance (float): total credit balance + reserved_balance (float): reserved credit balance + """ + self.total_balance = total_balance + self.reserved_balance = reserved_balance diff --git a/tests/unit/wallet_test.py b/tests/unit/wallet_test.py new file mode 100644 index 00000000..48ee19ab --- /dev/null +++ b/tests/unit/wallet_test.py @@ -0,0 +1,16 @@ +__author__ = "aixplain" + +from aixplain.factories import WalletFactory +import aixplain.utils.config as config +import requests_mock + + +def test_wallet_service(): + with requests_mock.Mocker() as mock: + url = f"{config.BACKEND_URL}/sdk/billing/wallet" + headers = {"x-api-key": config.TEAM_API_KEY, "Content-Type": "application/json"} + ref_response = {"totalBalance": 5, "reservedBalance": "0"} + mock.get(url, headers=headers, json=ref_response) + wallet = WalletFactory.get() + assert wallet.total_balance == ref_response["totalBalance"] + assert wallet.reserved_balance == ref_response["reservedBalance"]