Skip to content

Models

Task

Bases: Model

Source code in API/orchestrator/models.py
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
class Task(models.Model):
    user = models.ForeignKey(
        User,
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name="tasks",
        help_text="Select the user",
    )
    name = models.CharField(
        max_length=100, help_text="A unique name to track the cluster of tasks"
    )

    task_name = models.CharField(
        max_length=100,
        help_text="The name of the task",
    )
    parameters = models.JSONField(
        default=dict,
        blank=True,
        null=True,
        help_text="Enter the parameters for the task",
    )
    result_status = models.CharField(
        max_length=100,
        choices=[
            ("pending", "Pending"),
            ("completed", "Completed"),
            ("failed", "Failed"),
            ("started", "Started"),
            ("cancelled", "Cancelled"),
        ],
        default="pending",
    )
    result_json = models.JSONField(
        default=dict,
        blank=True,
        null=True,
        help_text="The result of the task",
    )
    description = models.TextField(blank=True, null=True)
    track_id = models.CharField(
        max_length=100,
        blank=True,
        null=True,
        help_text="The tracking ID of the task, will start with T-{cluster_name}-{id}",
    )
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    def __str__(self):
        return self.name

    @classmethod
    def create_task(
        cls,
        user: Optional[User],
        name: str,
        task_name: str,
        parameters: dict,
        description: str = "",
        track_id: Optional[str] = None,
    ):
        """
        Create a task
        Args:
            user (User): The user who created the task
            name (str): The name of the task
            task_name (str): The name of the task
            parameters (dict): The parameters for the task
            description (str): The description of the task
            track_id (str): The tracking ID of the task, will start with T-{cluster_name}-{id}

        Returns:

        """
        task = cls(
            user=user,
            name=name,
            task_name=task_name,
            parameters=parameters,
            description=description,
            track_id=track_id,
        )
        task.save()
        return task

    @staticmethod
    def init_track_id(name: str) -> str:
        """
        Initialize the track ID
        Args:
            name (str): The name of the task

        Returns:
            str: The track ID
        """
        uid = str(uuid4())
        # replace the - with ""
        uid = uid.replace("-", "")
        return f"T-{name}-{uid}"

    # override the save method, to call the chain
    def save(self, *args, **kwargs):
        # if it is updated, then we need to call the chain
        if self.result_status == "completed":
            completed_task.send(sender=self, data=self.__dict__)
        super().save(*args, **kwargs)

    @staticmethod
    def get_task_name_choices():
        """
        Get dynamic task name choices
        Returns:
            list: List of tuples containing task name choices
        """
        # Here you can fetch the choices from an external source or database
        return [
            ("quantization_llm", "Quantization LLM"),
            ("hf_llm", "HF LLM"),
            ("emotion_detection", "Emotion Detection"),
            ("speech2text", "Speech2Text"),
            ("text2speech", "Text2Speech"),
            ("general_ml", "General ML"),
            ("openai_speech2text", "OpenAI Speech2Text"),
            ("openai_gpt_4o", "OpenAI GPT4o"),
            ("openai_gpt_35", "OpenAI GPT3.5"),
            ("openai_text2speech", "OpenAI Text2Speech"),
            ("openai_gpt_4o_text_and_image", "OpenAI GPT4o Text and Image"),
            ("openai_gpt_4o_text_only", "OpenAI GPT4o Text"),
            ("rag", "RAG"),
        ]

    @staticmethod
    def task_ml_task_mapping() -> dict:
        return {
            "quantization_llm": "text_generation",
            "hf_llm": "text_generation",
            "emotion_detection": "emotion_detection",
            "speech2text": "speech2text",
            "text2speech": "text2speech",
            "general_ml": "general_ml",
            "openai_speech2text": "speech2text",
            "openai_gpt_4o": "text_generation",
            "openai_gpt_35": "text_generation",
            "openai_text2speech": "text2speech",
            "openai_gpt_4o_text_and_image": "text_generation",
            "openai_gpt_4o_text_only": "text_generation",
            "rag": "rag",
        }

create_task(user, name, task_name, parameters, description='', track_id=None) classmethod

Create a task Args: user (User): The user who created the task name (str): The name of the task task_name (str): The name of the task parameters (dict): The parameters for the task description (str): The description of the task track_id (str): The tracking ID of the task, will start with T-{cluster_name}-{id}

Returns:

Source code in API/orchestrator/models.py
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
@classmethod
def create_task(
    cls,
    user: Optional[User],
    name: str,
    task_name: str,
    parameters: dict,
    description: str = "",
    track_id: Optional[str] = None,
):
    """
    Create a task
    Args:
        user (User): The user who created the task
        name (str): The name of the task
        task_name (str): The name of the task
        parameters (dict): The parameters for the task
        description (str): The description of the task
        track_id (str): The tracking ID of the task, will start with T-{cluster_name}-{id}

    Returns:

    """
    task = cls(
        user=user,
        name=name,
        task_name=task_name,
        parameters=parameters,
        description=description,
        track_id=track_id,
    )
    task.save()
    return task

get_task_name_choices() staticmethod

Get dynamic task name choices Returns: list: List of tuples containing task name choices

Source code in API/orchestrator/models.py
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
@staticmethod
def get_task_name_choices():
    """
    Get dynamic task name choices
    Returns:
        list: List of tuples containing task name choices
    """
    # Here you can fetch the choices from an external source or database
    return [
        ("quantization_llm", "Quantization LLM"),
        ("hf_llm", "HF LLM"),
        ("emotion_detection", "Emotion Detection"),
        ("speech2text", "Speech2Text"),
        ("text2speech", "Text2Speech"),
        ("general_ml", "General ML"),
        ("openai_speech2text", "OpenAI Speech2Text"),
        ("openai_gpt_4o", "OpenAI GPT4o"),
        ("openai_gpt_35", "OpenAI GPT3.5"),
        ("openai_text2speech", "OpenAI Text2Speech"),
        ("openai_gpt_4o_text_and_image", "OpenAI GPT4o Text and Image"),
        ("openai_gpt_4o_text_only", "OpenAI GPT4o Text"),
        ("rag", "RAG"),
    ]

init_track_id(name) staticmethod

Initialize the track ID Args: name (str): The name of the task

Returns:

Name Type Description
str str

The track ID

Source code in API/orchestrator/models.py
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
@staticmethod
def init_track_id(name: str) -> str:
    """
    Initialize the track ID
    Args:
        name (str): The name of the task

    Returns:
        str: The track ID
    """
    uid = str(uuid4())
    # replace the - with ""
    uid = uid.replace("-", "")
    return f"T-{name}-{uid}"