refine

Author: ccssu

src/bizyair/commands/processors/prompt_processor.py

@@ -137,6 +137,7 @@ def process(
                     "last_node_id": last_node_id,
                 }
             )
+            # dict_keys(['139', '141', '142', '138', '143', '150'])
         return pre_prompt, str(last_node_id)
 
     def validate_input(

src/bizyair/commands/servers/prompt_server.py

@@ -79,6 +79,7 @@ def execute(
         data = json.dumps({"prompt": prompt, "last_node_id": last_node_id}).encode(
             "utf-8"
         )
+
         subscriber.prompt = prompt
         subscriber.last_node_id = last_node_id
         self.mediator.start_sse_client(

src/bizyair/commands/servers/pub_sub_sse.py

@@ -99,16 +99,18 @@ def connect():
                 sse_client = SSEClient(response)
                 for event in sse_client.events():
                     data = json.loads(event.data)
-                    import pprint
-
-                    pprint.pprint(truncate_long_strings(data))
+                    # import pprint
+                    # pprint.pprint(truncate_long_strings(data))
                     self.publish(data, subscriber=subscriber)
             except Exception as e:
                 print(f"Error connecting to {url}: {e}")
             finally:
                 print(f"SSE client for {subscriber.name} disconnected")
                 self.publish(self.STOP_SIGNAL, subscriber=subscriber)
 
+        import pdb
+
+        pdb.set_trace()
         self.executor.submit(connect)
 
 
@@ -119,14 +121,21 @@ def __init__(self, name, mediator):
         self.messages = queue.Queue()
         self.is_subscribed = True  # 订阅状态监测
         self.tmp_result = {}
+        self.prompt = {}
 
     def receive(self, message):
         self.messages.put(message)
 
-    def is_empty(self):
-        return self.messages.empty()
+    def is_active(self, node_id):
+        if node_id not in self.prompt:
+            return False
+        return (
+            self.is_subscribed
+            or not self.messages.empty()
+            or node_id in self.tmp_result
+        )
 
-    def pop(self, timeout=5):
+    def pop(self, timeout=86400):
         try:
             if self.is_subscribed or not self.messages.empty():
                 return self.messages.get(timeout=timeout)
@@ -140,16 +149,19 @@ def unsubscribe(self):
         self.mediator.unsubscribe(self)
         self.is_subscribed = False  # 更新订阅状态为未订阅
 
-    def get_result(self, node_id, timeout=12):
+    def get_result(self, node_id, timeout=86400):
         while True:
             if node_id in self.tmp_result:
                 out = self.tmp_result[node_id]
                 del self.tmp_result[node_id]
                 return out
 
             result = self.pop(timeout=timeout)
-            if result is None:
-                return None
+            import pprint
+
+            print(f"{node_id=}")
+            pprint.pprint(truncate_long_strings(result))
+
             try:
                 if (
                     "message" in result

src/bizyair/nodes_base.py

@@ -145,15 +145,11 @@ def _merge_results(self, result: List[List[Any]], node_ios: List[BizyAirNodeIO])
 
     def _should_pre_run(self):
         subscriber = BizyAirBaseNode.subscriber
-        return (
-            subscriber is not None
-            and not subscriber.is_empty()
-            and self.assigned_id in subscriber.prompt
-        )
+        return subscriber is not None and self.assigned_id in subscriber.prompt
 
-    def _pre_run(self):
+    def _pre_run(self, timeout=86400):
         subscriber = BizyAirBaseNode.subscriber
-        result = subscriber.get_result(self.assigned_id)
+        result = subscriber.get_result(self.assigned_id, timeout=timeout)
         return decode_data(result)
 
     def default_function(self, **kwargs):