rabbitmq logic update

worker.py

@@ -14,10 +14,33 @@ from mineru_single import Processor
 class RabbitMQWorker:
     def __init__(self, num_workers: int = 1):
         self.num_workers = num_workers
-        self.public_connection = pika.BlockingConnection(pika.URLParameters(os.getenv("RABBITMQ_URL", "amqp://guest:guest@localhost:5672/")))
-        self.public_channel = self.public_connection.channel()
+        self.rabbit_url = os.getenv("RABBITMQ_URL", "amqp://guest:guest@localhost:5672/")
         self.processor = Processor()
 
+    def publish_message(self, body_dict: dict, headers: dict):
+        """Create a new connection for each publish operation"""
+        try:
+            connection = pika.BlockingConnection(pika.URLParameters(self.rabbit_url))
+            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
+                )
+            )
+            
+            connection.close()
+            return True
+        except Exception as e:
+            print(f"Error publishing message: {e}")
+            return False
+
     def callback(self, ch, method, properties, body):
         """Handle incoming RabbitMQ messages"""
         thread_id = threading.current_thread().name
@@ -25,37 +48,38 @@ class RabbitMQWorker:
 
         print(f"[Worker {thread_id}] Received message: {body}, headers: {headers}")
 
-        if headers.get("request_type") == "process_files":
-            contexts = []
-            body_dict = json.loads(body)
-            for file in body_dict.get("input_files", []):
-                contexts.append({"key": file["key"], "body": self.processor.process(file["url"])})
-            body_dict["md_context"] = contexts
-            json_body = json.dumps(body_dict)
-            self.public_channel.queue_declare(queue="ml_server", durable=True)
-            self.public_channel.basic_publish(
-                exchange="",
-                routing_key="ml_server",
-                body=json_body,
-                properties=pika.BasicProperties(headers=headers)
-            )
-            print(f"[Worker {thread_id}] Contexts: {contexts}")
-            
-        else:
-            print(f"[Worker {thread_id}] Unknown process")
-            return
-
-    def worker(self, channel):
-        """Worker process to consume messages"""
         try:
-            channel.start_consuming()
+            if headers.get("request_type") == "process_files":
+                contexts = []
+                body_dict = json.loads(body)
+                
+                # Process files
+                for file in body_dict.get("input_files", []):
+                    try:
+                        context = {"key": file["key"], "body": self.processor.process(file["url"])}
+                        contexts.append(context)
+                    except Exception as e:
+                        print(f"Error processing file {file['key']}: {e}")
+                        contexts.append({"key": file["key"], "body": f"Error: {str(e)}"})
+
+                body_dict["md_context"] = contexts
+                
+                # Publish results
+                if self.publish_message(body_dict, headers):
+                    print(f"[Worker {thread_id}] Successfully published results")
+                else:
+                    print(f"[Worker {thread_id}] Failed to publish results")
+                
+                print(f"[Worker {thread_id}] Contexts: {contexts}")
+            else:
+                print(f"[Worker {thread_id}] Unknown process")
+                
         except Exception as e:
-            print(f"[Worker] Error: {e}")
+            print(f"Error in callback: {e}")
 
     def connect_to_rabbitmq(self):
         """Establish connection to RabbitMQ"""
-        rabbit_url = os.getenv("RABBITMQ_URL", "amqp://guest:guest@localhost:5672/")
-        connection = pika.BlockingConnection(pika.URLParameters(rabbit_url))
+        connection = pika.BlockingConnection(pika.URLParameters(self.rabbit_url))
         channel = connection.channel()
 
         channel.queue_declare(queue="gpu_server", durable=True)
@@ -70,10 +94,15 @@ class RabbitMQWorker:
     def start(self):
         """Start the worker threads"""
         print(f"Starting {self.num_workers} workers")
-        with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
-            for _ in range(self.num_workers):
-                connection, channel = self.connect_to_rabbitmq()
-                executor.submit(self.worker, channel)
+        while True:
+            try:
+                with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
+                    for _ in range(self.num_workers):
+                        connection, channel = self.connect_to_rabbitmq()
+                        executor.submit(self.worker, channel)
+            except Exception as e:
+                print(f"Connection lost, reconnecting... Error: {e}")
+                time.sleep(5)  # Wait before reconnecting
 
 def main():
     worker = RabbitMQWorker()