Implement Credits payment on Notify endpoint

src/aleph/vm/orchestrator/payment.py

@@ -19,7 +19,7 @@
 logger = logging.getLogger(__name__)
 
 
-async def fetch_balance_of_address(address: str) -> Decimal:
+async def get_address_balance(address: str) -> dict:
     """
     Get the balance of the user from the PyAleph API.
 
@@ -36,39 +36,36 @@ async def fetch_balance_of_address(address: str) -> Decimal:
 
         # Consider the balance as null if the address is not found
         if resp.status == 404:
-            return Decimal(0)
+            return {}
 
         # Raise an error if the request failed
         resp.raise_for_status()
 
-        resp_data = await resp.json()
-        return resp_data["balance"]
+        return await resp.json()
 
 
-async def fetch_credit_balance_of_address(address: str) -> Decimal:
+async def fetch_balance_of_address(address: str) -> Decimal:
     """
     Get the balance of the user from the PyAleph API.
+    """
 
-    API Endpoint:
-        GET /api/v0/addresses/{address}/balance
+    resp_data = await get_address_balance(address)
+    if hasattr(resp_data, "balance"):
+        return Decimal(resp_data["balance"])
 
-    For more details, see the PyAleph API documentation:
-    https://github.com/aleph-im/pyaleph/blob/master/src/aleph/web/controllers/routes.py#L62
-    """
+    return Decimal(0)
 
-    async with aiohttp.ClientSession() as session:
-        url = f"{settings.API_SERVER}/api/v0/addresses/{address}/credit_balance"
-        resp = await session.get(url)
 
-        # Consider the balance as null if the address is not found
-        if resp.status == 404:
-            return Decimal(0)
+async def fetch_credit_balance_of_address(address: str) -> Decimal:
+    """
+    Get the credits of the user from the PyAleph API.
+    """
 
-        # Raise an error if the request failed
-        resp.raise_for_status()
+    resp_data = await get_address_balance(address)
+    if hasattr(resp_data, "credit_balance"):
+        return Decimal(resp_data["credit_balance"])
 
-        resp_data = await resp.json()
-        return resp_data["credits"]
+    return Decimal(0)
 
 
 async def fetch_execution_price(
@@ -82,7 +79,7 @@ async def fetch_execution_price(
         resp.raise_for_status()
 
         resp_data = await resp.json()
-        required_credits: float = resp_data["required_credits"]  # Field not defined yet on API side.
+        required_credits: float = resp_data["cost"]
         payment_type: str | None = resp_data["payment_type"]
 
         if payment_type_required and payment_type is None:

src/aleph/vm/orchestrator/views/__init__.py

@@ -35,6 +35,7 @@
 from aleph.vm.orchestrator.payment import (
     InvalidAddressError,
     InvalidChainError,
+    fetch_credit_balance_of_address,
     fetch_execution_price,
     get_stream,
 )
@@ -645,6 +646,27 @@ async def notify_allocation(request: web.Request):
                     f"Present: {active_flow_per_month} / month (flow = {community_flow})\n"
                     f"Address: {community_wallet}",
                 )
+    elif payment_type == PaymentType.credit:
+        if is_confidential:
+            logger.debug(f"Confidential instance {item_hash} using Credits")
+        if have_gpu:
+            logger.debug(f"GPU Instance {item_hash} using Credits")
+
+        credit_balance = await fetch_credit_balance_of_address(message.content.address)
+
+        required_credits = await fetch_execution_price(item_hash, [PaymentType.credit])
+        logger.debug(
+            f"Required credit balance for Address {message.content.address} executions: {required_credits}, {item_hash}"
+        )
+
+        if required_credits > credit_balance:
+            return web.HTTPPaymentRequired(
+                reason="Insufficient balance",
+                text="Insufficient credits for this instance\n\n"
+                f"Required: {required_credits} credits \n"
+                f"Current user credits: {credit_balance}",
+            )
+
     else:
         return web.HTTPBadRequest(reason="Invalid payment method")