| 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() |