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