102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
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 apps.orgs.models import Organization
|
|
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')
|
|
organization = ForeignKey(Organization, on_delete = CASCADE, related_name = 'agents', null = True, blank = True)
|
|
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, blank = True)
|
|
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} ({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}"
|
|
|
|
class Meta:
|
|
ordering = ['timestamp']
|
|
verbose_name = "Agent Event"
|
|
verbose_name_plural = "Agent Events"
|