Dynavera/apps/mlstore/models.py

from django.db.models import BigAutoField, CASCADE, CharField, DateTimeField, ForeignKey, JSONField, Model, TextField, UUIDField
from apps.users.mixins import TimeStampMixin
from apps.users.models import User
from uuid import uuid4

class AgentModel(Model):

    id = BigAutoField(primary_key = True)
    uuid = UUIDField(default = uuid4, unique = True, editable = False)
    name = CharField(max_length = 255)
    version = CharField(max_length = 50)
    path = CharField(max_length=1024, blank=True, default='')

    class Meta:
        verbose_name = 'Model'
        verbose_name_plural = 'Models'

    def __str__(self):
        return self.name

class Agent(TimeStampMixin, Model):

    STATUS_CHOICES = [
        ('idle', 'Idle'),
        ('running', 'Running'),
        ('paused', 'Paused'),
        ('completed', 'Completed'),
        ('failed', 'Failed'),
    ]

    id   = BigAutoField(primary_key = True)
    uuid = UUIDField(default = uuid4, unique = True, editable = False)
    
    model = ForeignKey(AgentModel, on_delete = CASCADE, related_name = 'agents')
    status = CharField(max_length = 20, choices = STATUS_CHOICES, default = 'idle')

    description = TextField(blank = True, default = '')

    started_at   = DateTimeField(null = True, blank = True)
    completed_at = DateTimeField(null = True, blank = True)

    class Meta:
        verbose_name = 'Agent Instance'
        verbose_name_plural = 'Agent Instances'

    def __str__(self):
        return f'{self.model.name} - {self.uuid}'

class AgentRun(TimeStampMixin, Model):

    RUN_CHOICES = [
        ('queued', 'Queued'),
        ('running', 'Running'),
        ('completed', 'Completed'),
        ('failed', 'Failed'),
    ]
    
    id = BigAutoField(primary_key = True)
    uuid = UUIDField(default = uuid4, editable = False, unique = True)
    agent = ForeignKey(Agent, on_delete = CASCADE, related_name = 'runs')
    user = ForeignKey(User, on_delete = CASCADE, related_name = 'agent_runs')
    status = CharField(max_length = 20, choices = RUN_CHOICES, default = 'queued')

    input_data  = JSONField(default = dict)
    output_data = JSONField(default = dict, blank = True)
    error_message = TextField(blank = True, default = "")
    started_at = DateTimeField(null = True, blank = True)
    completed_at = DateTimeField(null = True, blank = True)

    def __str__(self) -> str:
        return f"Execution {self.uuid} - {self.agent.name} ({self.status})"

    class Meta:
        verbose_name = "Agent Run"
        verbose_name_plural = "Agent Runs"

class AgentEvent(Model):
	EVENT_TYPES = [
		('started', 'Started'),
		('message', 'Message'),
		('progress', 'Progress'),
		('completed', 'Completed'),
		('error', 'Error'),
		('step', 'Step'),
	]

	uuid = UUIDField(default = uuid4, editable = False, unique = True)
	execution = ForeignKey(AgentRun, on_delete = CASCADE, related_name = 'events')
	event_type = CharField(max_length = 20, choices = EVENT_TYPES)

	content = JSONField()
	timestamp = DateTimeField(auto_now_add = True)

	def __str__(self) -> str:
		return f"{self.id} - {self.event_type} - {self.execution.agent.name}"

	class Meta:
		ordering = ['timestamp']
		verbose_name = "Agent Event"
		verbose_name_plural = "Agent Events"