annotations/Azaion.AI/main.pyx

import queue
import threading
cimport constants
from api_client cimport ApiClient
from annotation cimport Annotation
from inference import Inference
from processor_command cimport FileCommand, CommandType, ProcessorType
from remote_handlers cimport SocketHandler, RabbitHandler
import argparse

cdef class ParsedArguments:
    cdef str email, password, folder;
    cdef bint persist_token

    def __init__(self, str email, str password, str folder, bint persist_token):
        self.email = email
        self.password = password
        self.folder = folder
        self.persist_token = persist_token

cdef class CommandProcessor:
    cdef ApiClient api_client
    cdef SocketHandler socket_handler
    cdef RabbitHandler rabbit_handler
    cdef object command_queue
    cdef bint running

    def __init__(self, args: ParsedArguments):
        self.api_client = ApiClient(args.email, args.password, args.folder)
        self.socket_handler = SocketHandler(self.on_message)
        self.socket_handler.start()
        self.rabbit_handler = RabbitHandler(self.api_client, self.on_message)
        self.rabbit_handler.start()
        self.command_queue = queue.Queue(maxsize=constants.QUEUE_MAXSIZE)
        self.running = True

    def start(self):
        threading.Thread(target=self.process_queue, daemon=True).start()

    cdef on_message(self, FileCommand cmd):
        try:
            if cmd.command_type == CommandType.INFERENCE:
                self.command_queue.put(cmd)
            elif cmd.command_type == CommandType.LOAD:
                threading.Thread(target=self.process_load, args=[cmd], daemon=True).start()
        except Exception as e:
            print(f"Error handling client: {e}")

    cdef on_annotations(self, cmd: FileCommand, annotations: [Annotation]):
        handler = self.socket_handler if cmd.processor_type == ProcessorType.SOCKET else self.rabbit_handler
        handler.send(annotations)


    cdef process_queue(self):
        while self.running:
            try:
                command = self.command_queue.get()
                model = self.api_client.load_ai_model()
                Inference(model, self.on_annotations).run_inference(command)
            except Exception as e:
                print(f"Error processing queue: {e}")

    cdef process_load(self, FileCommand command):
        response = self.api_client.load_file(command)
        handler = self.socket_handler if command.processor_type == ProcessorType.SOCKET else self.rabbit_handler
        handler.send(response)


    def stop(self):
        self.running = False

def parse_arguments():
    parser = argparse.ArgumentParser(description="Command Processor")
    parser.add_argument("-e", "--email",  type=str, default="", help="Email")
    parser.add_argument("-p", "--pw", type=str, default="", help="Password")
    parser.add_argument("-f", "--folder", type=str, default="", help="Folder to API inner folder to download file from")
    parser.add_argument("-t", "--persist_token", type=bool, default=True, help="True for persisting token from API")
    cdef args = parser.parse_args()

    cdef str email = args.email
    cdef str password = args.pw
    cdef str folder = args.folder
    cdef bint persist_token = args.persist_token

    return ParsedArguments(email, password, folder, persist_token)

if __name__ == '__main__':
    args = parse_arguments()
    processor = CommandProcessor(args)
    try:
        processor.start()
    except KeyboardInterrupt:
        processor.stop()