From a2003cee66bb043f0f4b9b316e929620b0b65d32 Mon Sep 17 00:00:00 2001 From: xainaz Date: Mon, 19 Aug 2024 14:46:19 +0300 Subject: [PATCH 1/5] Add .DS_Store to .gitignore --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index 843c6556..ad7c16c8 100644 --- a/.gitignore +++ b/.gitignore @@ -130,3 +130,5 @@ dmypy.json # Vscode .vscode +.DS_Store + From 14b15eb22e78de8cb0261313bb29318861907a4d Mon Sep 17 00:00:00 2001 From: xainaz Date: Mon, 19 Aug 2024 16:07:41 +0300 Subject: [PATCH 2/5] Initial commit, added wallet factory and wallet module. No test is added yet --- aixplain/factories/wallet_factoy.py | 28 ++++++++++++++++++++++++++++ aixplain/modules/wallet.py | 7 +++++++ 2 files changed, 35 insertions(+) create mode 100644 aixplain/factories/wallet_factoy.py create mode 100644 aixplain/modules/wallet.py diff --git a/aixplain/factories/wallet_factoy.py b/aixplain/factories/wallet_factoy.py new file mode 100644 index 00000000..23503d77 --- /dev/null +++ b/aixplain/factories/wallet_factoy.py @@ -0,0 +1,28 @@ +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 + resp = None + + def get(cls) -> Wallet: + try: + # Check for code 200, other code will be caught when trying to return a Wallet object + url = f"{cls.backend_url}/sdk/billing/wallet" + if cls.aixplain_key != "": + headers = {"x-aixplain-key": f"{cls.aixplain_key}", "Content-Type": "application/json"} + else: + headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"} + logging.info(f"Start fetching billing information from - {url} ") + headers = {"Content-Type": "application/json", "x-api-key": config.TEAM_API_KEY} + r = _request_with_retry("get", url, headers=headers) + resp = r.json() + resp["api_key"] = config.TEAM_API_KEY + return Wallet(totalBalance=resp["totalBalance"], reservedBalance=resp["reservedBalance"]) + + except Exception: + raise Exception(f"Failed to get wallet. Please check API key.") diff --git a/aixplain/modules/wallet.py b/aixplain/modules/wallet.py new file mode 100644 index 00000000..4707e663 --- /dev/null +++ b/aixplain/modules/wallet.py @@ -0,0 +1,7 @@ +from typing import Union + + +class Wallet: + def __init__(self, totalBalance: Union[int, float], reservedBalance: Union[int, float]): + self.totalBalance = totalBalance + self.reservedBalance = reservedBalance From dbb879726803d2bca4de2fb6c7e413d6b76a099e Mon Sep 17 00:00:00 2001 From: xainaz Date: Mon, 19 Aug 2024 17:41:36 +0300 Subject: [PATCH 3/5] Added test --- tests/unit/wallet_test.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 tests/unit/wallet_test.py diff --git a/tests/unit/wallet_test.py b/tests/unit/wallet_test.py new file mode 100644 index 00000000..3b24f27a --- /dev/null +++ b/tests/unit/wallet_test.py @@ -0,0 +1,13 @@ +from aixplain.factories import WalletFactory +import pytest + +def test_wallet(): + factory = WalletFactory() + with pytest.raises(Exception): + wallet = factory.get() + total_balance_str= f"Total Balance: {wallet.totalBalance}." + reserved_balance_str= f"Reserved Balance: {wallet.reservedBalance}." + assert str(total_balance_str) == "Total Balance: 5." + assert str(reserved_balance_str=="Reserced Balance: 0.") + + From fb49910acbe530560aebd8165b1c1daae86685d9 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Tue, 20 Aug 2024 10:46:32 -0300 Subject: [PATCH 4/5] Formatting Wallet tests --- aixplain/factories/__init__.py | 1 + aixplain/factories/wallet_factoy.py | 20 ++++++------- aixplain/modules/agent/tool/model_tool.py | 3 +- aixplain/modules/wallet.py | 35 ++++++++++++++++++++--- tests/unit/wallet_test.py | 23 ++++++++------- 5 files changed, 55 insertions(+), 27 deletions(-) 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 index 23503d77..59ec7c14 100644 --- a/aixplain/factories/wallet_factoy.py +++ b/aixplain/factories/wallet_factoy.py @@ -7,22 +7,20 @@ class WalletFactory: aixplain_key = config.AIXPLAIN_API_KEY backend_url = config.BACKEND_URL - resp = None + @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" - if cls.aixplain_key != "": - headers = {"x-aixplain-key": f"{cls.aixplain_key}", "Content-Type": "application/json"} - else: - headers = {"Authorization": f"Token {config.TEAM_API_KEY}", "Content-Type": "application/json"} - logging.info(f"Start fetching billing information from - {url} ") + + 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() - resp["api_key"] = config.TEAM_API_KEY - return Wallet(totalBalance=resp["totalBalance"], reservedBalance=resp["reservedBalance"]) - - except Exception: - raise Exception(f"Failed to get wallet. Please check API key.") + 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/agent/tool/model_tool.py b/aixplain/modules/agent/tool/model_tool.py index 79e4601d..e15a8bea 100644 --- a/aixplain/modules/agent/tool/model_tool.py +++ b/aixplain/modules/agent/tool/model_tool.py @@ -24,7 +24,6 @@ from aixplain.enums.function import Function from aixplain.enums.supplier import Supplier -from aixplain.factories.model_factory import ModelFactory from aixplain.modules.agent.tool import Tool from aixplain.modules.model import Model @@ -76,4 +75,4 @@ def __init__( model = model.id self.supplier = supplier self.model = model - self.function = function \ No newline at end of file + self.function = function diff --git a/aixplain/modules/wallet.py b/aixplain/modules/wallet.py index 4707e663..d7c63524 100644 --- a/aixplain/modules/wallet.py +++ b/aixplain/modules/wallet.py @@ -1,7 +1,34 @@ -from typing import Union +__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, totalBalance: Union[int, float], reservedBalance: Union[int, float]): - self.totalBalance = totalBalance - self.reservedBalance = reservedBalance + 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 index 3b24f27a..48ee19ab 100644 --- a/tests/unit/wallet_test.py +++ b/tests/unit/wallet_test.py @@ -1,13 +1,16 @@ -from aixplain.factories import WalletFactory -import pytest +__author__ = "aixplain" -def test_wallet(): - factory = WalletFactory() - with pytest.raises(Exception): - wallet = factory.get() - total_balance_str= f"Total Balance: {wallet.totalBalance}." - reserved_balance_str= f"Reserved Balance: {wallet.reservedBalance}." - assert str(total_balance_str) == "Total Balance: 5." - assert str(reserved_balance_str=="Reserced Balance: 0.") +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"] From 13415d47d5b0eb2112a2edfcc14916b5e45cb6a5 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Tue, 20 Aug 2024 11:09:23 -0300 Subject: [PATCH 5/5] wallet_factoy -> wallet_factory --- aixplain/factories/__init__.py | 2 +- aixplain/factories/{wallet_factoy.py => wallet_factory.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename aixplain/factories/{wallet_factoy.py => wallet_factory.py} (100%) 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_factoy.py b/aixplain/factories/wallet_factory.py similarity index 100% rename from aixplain/factories/wallet_factoy.py rename to aixplain/factories/wallet_factory.py