MinerU / test_listener.py
SkyNait's picture
listener
25c35b4
import pika
import os
import json
import dotenv
import threading
import time
dotenv.load_dotenv()
RABBITMQ_URL = os.getenv("RABBITMQ_URL")
def send_topic_extraction_request(payload: dict):
"""Simulate ml_server sending a topic extraction request to gpu_server."""
params = pika.URLParameters(RABBITMQ_URL)
params.heartbeat = 5
params.blocked_connection_timeout = 2
connection = pika.BlockingConnection(params)
channel = connection.channel()
channel.queue_declare(queue="gpu_server", durable=True)
message = json.dumps(payload).encode("utf-8")
channel.basic_publish(
exchange="",
routing_key="gpu_server",
body=message,
properties=pika.BasicProperties(delivery_mode=2)
)
print("Topic extraction request sent from ml_server to gpu_server.")
connection.close()
def listen_for_results():
"""Simulate ml_server listening for topic extraction results on ml_server queue."""
params = pika.URLParameters(RABBITMQ_URL)
params.heartbeat = 5
params.blocked_connection_timeout = 2
connection = pika.BlockingConnection(params)
channel = connection.channel()
channel.queue_declare(queue="ml_server", durable=True)
def callback(ch, method, properties, body):
try:
result = json.loads(body)
print("Received topic extraction result:")
print(json.dumps(result, indent=2))
ch.basic_ack(delivery_tag=method.delivery_tag)
except Exception as e:
print("Error processing message:", e)
ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
channel.basic_consume(queue="ml_server", on_message_callback=callback)
print("Listening for topic extraction results on ml_server queue...")
channel.start_consuming()
if __name__ == "__main__":
payload = {
"pattern": "topic_extraction",
"data": {
"input_files": [
{
"key": "",
"url": "https://qualifications.pearson.com/content/dam/pdf/A%20Level/Mathematics/2017/specification-and-sample-assesment/a-level-l3-mathematics-specification-issue4.pdf",
"type": "specification",
"page": [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 34, 35, 36, 37, 38, 39, 40, 41, 42]
}
],
"topics": [
{
"title": "",
"id": 123
}
]
}
}
producer_thread = threading.Thread(target=send_topic_extraction_request, args=(payload,))
producer_thread.start()
time.sleep(1)
listen_for_results()