Added optional account_address to get_plan_balance, fixed create task event handling

payments_py/ai_query_api.py

@@ -94,7 +94,8 @@ async def create_task(self, did: str, task: Task, _callback: Optional[Any]=None)
         token = self.get_service_token(did)
         result = self.post(endpoint, task, headers={'Authorization': f'Bearer {token.accessToken}'})
         if(result.status_code == 201 and _callback):
-            await self.subscribe_task_logs(_callback, [result["task"]["task_id"]])
+            tasks = result.json()
+            await self.subscribe_task_logs(_callback, [tasks["task"]["task_id"]])
         return result
 
     def create_steps(self, did: str, task_id: str, steps: List[Step]):
@@ -221,8 +222,7 @@ async def subscribe_task_logs(self, callback: Any, tasks: List[str]):
             if not tasks:
                 raise Exception('No task rooms to join in configuration')
             await self.connect_socket()
-
-            await self.socket_client.on('_connected', self._on_connected(callback, tasks))
+            self.socket_client.on('_connected', self._on_connected(callback, tasks))
         except Exception as error:
             raise Exception(f"Unable to initialize websocket client: {self.web_socket_host} - {str(error)}")
 
@@ -235,6 +235,6 @@ async def handle_connected_event(*args):
             async def handle_task_log_event(data: Any):
                 callback(data)
 
-            await self.socket_client.on('task-log', handle_task_log_event)
+            self.socket_client.on('task-log', handle_task_log_event)
 
         return handle_connected_event    

payments_py/payments.py

@@ -1,6 +1,7 @@
 import os
 from typing import List, Optional
 import requests
+import jwt
 
 from payments_py.data_models import BalanceResultDto, BurnResultDto, CreateAssetResultDto, DownloadFileResultDto, MintResultDto, OrderPlanResultDto, ServiceTokenResultDto
 from payments_py.environments import Environment
@@ -49,6 +50,8 @@ def __init__(self, nvm_api_key: str, environment: Environment,
         self.environment = environment
         self.app_id = app_id
         self.version = version
+        decoded_jwt = jwt.decode(self.nvm_api_key, options={"verify_signature": False})
+        self.account_address = decoded_jwt.get('sub')
         if ai_protocol:
             self.ai_protocol = AIQueryApi(self.backend_options)
 
@@ -551,13 +554,13 @@ def get_asset_ddo(self, did: str):
         response = self.get(f"{self.environment.value['backend']}/api/v1/payments/asset/ddo/{did}")
         return response
 
-    def get_plan_balance(self, plan_did: str, account_address: str) -> BalanceResultDto:
+    def get_plan_balance(self, plan_did: str, account_address: Optional[str] = None) -> BalanceResultDto:
         """
         Get the balance of an account for a Payment Plan.
 
         Args:
             plan_did (str): The DID of the plan.
-            account_address (str): The account address.
+            account_address (Optional[str]): The account address. Defaults to `self.account_address` if not provided.
 
         Returns:
             BalanceResultDto: The response from the API call formatted as a BalanceResultDto.
@@ -579,13 +582,18 @@ def get_plan_balance(self, plan_did: str, account_address: str) -> BalanceResult
                 "balance": 10000000
             }
         """
+        # Use self.account_address if account_address is not provided
+        account_address = account_address or self.account_address
+
         body = {
             "subscriptionDid": plan_did,
-            **{snake_to_camel(k): v for k, v in locals().items() if v is not None and k != 'self'}
+            "accountAddress": account_address,
         }
+
         url = (f"{self.environment.value['backend']}/api/v1/payments/subscription/balance")
         response = self.post(url, body)
         response.raise_for_status()
+
         balance = {
             "planType": response.json()['subscriptionType'],
             "isOwner": response.json()['isOwner'],
@@ -594,6 +602,9 @@ def get_plan_balance(self, plan_did: str, account_address: str) -> BalanceResult
         }
         return BalanceResultDto.model_validate(balance)
 
+    def get_service_access_config(self, service_did: str) -> ServiceTokenResultDto:
+        return self.get_service_token(service_did)
+
     def get_service_token(self, service_did: str) -> ServiceTokenResultDto:
         """
         Get the required configuration for accessing a remote service agent.

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "payments-py"
-version = "0.5.2"
+version = "0.5.3"
 description = ""
 authors = ["enrique <enrique@nevermined.io>"]
 readme = "README.md"