Eng 398 sdk get users credits - Initial

.gitignore

@@ -130,3 +130,5 @@ dmypy.json
 
 # Vscode
 .vscode
+.DS_Store
+

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

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)}")

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

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"]