100 lines
3 KiB
Python
100 lines
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 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)
|
||
|
|
|
||
|
|
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"
|