feat: simplify some methods

nevermined-io · Oct 11, 2024 · ae95303 · ae95303
1 parent 68ad4fb
commit ae95303
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 39 deletions.
diff --git a/payments_py/ai_query_api.py b/payments_py/ai_query_api.py
@@ -38,15 +38,11 @@ def __init__(self, opts: BackendApiOptions):
     async def subscribe(self, callback: Any, did: Optional[str]=None, events: Optional[str]=None):
         await self._subscribe(callback, did, events)
         print('query-api:: Connected to the server')
-        pending_steps = self.get_steps(AgentExecutionStatus.Pending)
-        await self._emit_events(pending_steps.json()['steps'])
-
-        # if callback is not None:
-        #     await callback(pending_steps.json()['steps'])
-        # print('query-api:: Pending steps:', pending_steps.json())
-
-
-
+        try: 
+            pending_steps = self.get_steps(AgentExecutionStatus.Pending)
+            await self._emit_events(pending_steps.json()['steps'])
+        except Exception as e:
+            print('query-api:: Unable to get pending events', e)
 
     def create_task(self, did: str, task: Any, jwt: Optional[str] = None):
         """

diff --git a/payments_py/nvm_backend.py b/payments_py/nvm_backend.py
@@ -8,7 +8,7 @@
 
 from payments_py.environments import Environment
 
-sio = socketio.AsyncClient(logger=False, engineio_logger=False)
+sio = socketio.AsyncClient(logger=True, engineio_logger=True)
 
 
 class BackendApiOptions:
@@ -113,9 +113,9 @@ async def event_handler(data):
 
     async def _emit_events(self, data: Any):
         await self.connect_socket()
-        print(f"nvm-backend:: Emitting steps that were pending: {data}")
         for x in data:
-            await self.socket_client.emit('step-updated', data=x)
+            print(f"nvm-backend:: Emitting step: {x}")
+            await self.socket_client.emit(event='_emit-steps', data=x)
 
 
     async def join_room(self, room_ids: Optional[Union[str, List[str]]] = None):

diff --git a/tests/protocol_test.py b/tests/protocol_test.py
@@ -41,32 +41,32 @@ async def eventsReceived(data):
     payments_builder = Payments(nvm_api_key=nvm_api_key, environment=Environment.appStaging, app_id="your_app_id", version="1.0.0", ai_protocol=True, web_socket_options={'bearer_token': nvm_api_key})
     global response_data
     print('eventsReceived::', len(data))
-    # if isinstance(data, list):
-    #     print('eventsReceived::', 'pending data:', len(data))
-    #     for step in data:
-    #         print('eventsReceived::', 'step:', step)
-    #         result = payments_builder.ai_protocol.update_step(did=step['did'], 
-    #                                                           task_id=step['task_id'], 
-    #                                                           step_id=step['step_id'], 
-    #                                                           step={'step_id': step['step_id'],
-    #                                                                 'task_id': step['task_id'], 
-    #                                                                 'step_status': AgentExecutionStatus.Completed.value,
-    #                                                                 'output': 'success',
-    #                                                                 'is_last': True
-    #                                                                 })
-    #         print(result.json())
-
-    # else:
-    print('eventsReceived::', 'parsing event with did:', data)
-    result = payments_builder.ai_protocol.update_step(did=data["did"], 
-                                                            task_id=data["task_id"], 
-                                                            step_id=data['step_id'], 
-                                                            step={'step_id': data['step_id'],
-                                                                'task_id': data["task_id"], 
-                                                                'step_status': AgentExecutionStatus.Completed.value,
-                                                                'output': 'success',
-                                                                'is_last': True
-                                                                })
+    if isinstance(data, list):
+        print('eventsReceived::', 'pending data:', len(data))
+        for step in data:
+            print('eventsReceived::', 'step:', step)
+            result = payments_builder.ai_protocol.update_step(did=step['did'], 
+                                                              task_id=step['task_id'], 
+                                                              step_id=step['step_id'], 
+                                                              step={'step_id': step['step_id'],
+                                                                    'task_id': step['task_id'], 
+                                                                    'step_status': AgentExecutionStatus.Completed.value,
+                                                                    'output': 'success',
+                                                                    'is_last': True
+                                                                    })
+            print(result.json())
+
+    else:
+        print('eventsReceived::', 'parsing event with did:', data)
+        result = payments_builder.ai_protocol.update_step(did=data["did"], 
+                                                                task_id=data["task_id"], 
+                                                                step_id=data['step_id'], 
+                                                                step={'step_id': data['step_id'],
+                                                                    'task_id': data["task_id"], 
+                                                                    'step_status': AgentExecutionStatus.Completed.value,
+                                                                    'output': 'success',
+                                                                    'is_last': True
+                                                                    })
     response_data = data
     response_event.set()
     print(result.json())
@@ -117,7 +117,7 @@ async def test_AIQueryApi_create_task(ai_query_api_build_fixture, ai_query_api_s
 
     subscriber.ai_protocol.create_task(agent.did, {'query': 'sample_query', 'name': 'sample_task', 'additional_params': {'param1': 'value1', 'param2': 'value2'}})
 
-    await asyncio.wait_for(response_event.wait(), timeout=20)
+    await asyncio.wait_for(response_event.wait(), timeout=30)
 
     assert response_data is not None, "Builder did not receive the event from subscriber"
     print('Task received by builder:', response_data)