Update worker.py

worker.py

@@ -17,33 +17,53 @@ class RabbitMQWorker:
         self.rabbit_url = os.getenv("RABBITMQ_URL")
         self.processor = Processor()
 
-    def publish_message(self, body_dict: dict, headers: dict):
-        """Create a new connection for each publish operation"""
-        try:
+        self.publisher_connection = None
+        self.publisher_channel = None
+
+
+    def setup_publisher(self):
+        if not self.publiserh_connection or self.publisher_connection.is_closed:
             connection_params = pika.URLParameters(self.rabbit_url)
-            connection_params.heartbeat = 10
-            connection_params.blocked_connection_timeout = 5
+            connection_params.heartbeat = 30  # Match the consumer heartbeat
+            connection_params.blocked_connection_timeout = 300  # Increase timeout
             
-            connection = pika.BlockingConnection(connection_params)
-            channel = connection.channel()
-            
-            channel.queue_declare(queue="ml_server", durable=True)
-            
-            channel.basic_publish(
-                exchange="",
-                routing_key="ml_server",
-                body=json.dumps(body_dict),
-                properties=pika.BasicProperties(
-                    delivery_mode=2,  # make message persistent
-                    headers=headers
+            self.publisher_connection = pika.BlockingConnection(connection_params)
+            self.publisher_channel = self.publisher_connection.channel()
+            self.publisher_channel.queue_declare(queue="ml_server", durable=True)
+    
+    def publish_message(self, body_dict: dict, headers: dict):
+        """Use persistent connection for publishing"""
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                # Ensure publisher connection is setup
+                self.setup_publisher()
+                
+                self.publisher_channel.basic_publish(
+                    exchange="",
+                    routing_key="ml_server",
+                    body=json.dumps(body_dict),
+                    properties=pika.BasicProperties(
+                        delivery_mode=2,
+                        headers=headers
+                    )
                 )
-            )
-            
-            connection.close()
-            return True
-        except Exception as e:
-            print(f"Error publishing message: {e}")
-            return False
+                return True
+            except Exception as e:
+                print(f"Publish attempt {attempt + 1} failed: {e}")
+                # Close failed connection
+                if self.publisher_connection and not self.publisher_connection.is_closed:
+                    try:
+                        self.publisher_connection.close()
+                    except:
+                        pass
+                self.publisher_connection = None
+                self.publisher_channel = None
+                
+                if attempt == max_retries - 1:
+                    print(f"Failed to publish after {max_retries} attempts")
+                    return False
+                time.sleep(2) 
 
     def callback(self, ch, method, properties, body):
         """Handle incoming RabbitMQ messages"""