Dynavera/apps/mlstore/models.py

from django.db.models import BigAutoField, BooleanField, CASCADE, CharField, DateTimeField, ForeignKey, JSONField, Model, TextField, UUIDField, IntegerField
from pgvector.django import VectorField
from apps.users.mixins import TimeStampMixin
from apps.users.models import User
from apps.orgs.models import Organization, Role, TrainingFile
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"


class RoleRagDocument(TimeStampMixin, Model):

    id = BigAutoField(primary_key = True)
    uuid = UUIDField(default = uuid4, editable = False, unique = True)
    role = ForeignKey(Role, on_delete = CASCADE, related_name = 'rag_documents')
    training_file = ForeignKey(TrainingFile, on_delete = CASCADE, related_name = 'rag_documents', null = True, blank = True)

    content = TextField()
    content_hash = CharField(max_length = 64, db_index = True)
    embedding = VectorField(dimensions = 1536, null = True, blank = True)
    metadata = JSONField(default = dict, blank = True)
    chunk_index = IntegerField(default = 0)
    is_active = BooleanField(default = True)

    class Meta:
        verbose_name = "Role RAG Document"
        verbose_name_plural = "Role RAG Documents"

    def __str__(self) -> str:
        return f"{self.role.name} - chunk {self.chunk_index}"