app.py

import os
import time
import logging
from aimessage import start_kafka_consumer_thread
from aitask import model, tokenizer

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


if __name__ == "__main__":
    consumers_config = [
        {"group_id": "scan-group-1", "topic": "scan-cyber-1"},
        {"group_id": "scan-group-2", "topic": "scan-cyber-2"},
        {"group_id": "scan-group-3", "topic": "scan-cyber-3"},
        {"group_id": "scan-group-4", "topic": "scan-cyber-4"},
        {"group_id": "scan-group-5", "topic": "scan-cyber-5"},
        {"group_id": "scan-group-6", "topic": "scan-cyber-6"},
        {"group_id": "scan-group-7", "topic": "scan-cyber-7"},
        {"group_id": "scan-group-8", "topic": "scan-cyber-8"},
    ]

    consumer_threads = []
    for config in consumers_config:
        consumer_thread = start_kafka_consumer_thread(config["group_id"], config["topic"])
        consumer_threads.append(consumer_thread)

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        logger.info("Main thread interrupted by user. Exiting.")
        for thread in consumer_threads:
            thread.join()
        logger.info("All threads have been cleanly shut down.")