ai-ticket-platform/backend/apps/agents/models.py

import uuid

from django.conf import settings
from django.db import models
from django.db.models import Q

from apps.modelhub.models import ModelDeployment, ModelInvocation
from apps.skills.models import SkillVersion


class AgentDefinition(models.Model):
    code = models.SlugField(max_length=80, unique=True)
    name = models.CharField(max_length=120)
    description = models.TextField(blank=True)
    active_version = models.ForeignKey(
        "AgentVersion",
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="active_for_definitions",
    )
    created_by = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="created_agents",
    )
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ["name"]


class AgentVersion(models.Model):
    definition = models.ForeignKey(
        AgentDefinition,
        on_delete=models.CASCADE,
        related_name="versions",
    )
    version = models.PositiveIntegerField()
    status = models.CharField(max_length=24, default="draft")
    model_deployment = models.ForeignKey(
        ModelDeployment,
        on_delete=models.PROTECT,
        related_name="agent_versions",
    )
    system_prompt = models.TextField(blank=True)
    input_schema = models.JSONField(default=dict)
    example_input = models.JSONField(default=dict, blank=True)
    temperature_override = models.FloatField(null=True, blank=True)
    max_output_tokens_override = models.PositiveIntegerField(null=True, blank=True)
    created_by = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="created_agent_versions",
    )
    published_by = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="published_agent_versions",
    )
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)
    published_at = models.DateTimeField(null=True, blank=True)

    class Meta:
        ordering = ["-version"]
        constraints = [
            models.UniqueConstraint(
                fields=["definition", "version"],
                name="unique_agent_version",
            ),
            models.UniqueConstraint(
                fields=["definition"],
                condition=Q(status="draft"),
                name="unique_agent_draft",
            ),
        ]


class AgentStepBinding(models.Model):
    agent_version = models.ForeignKey(
        AgentVersion,
        on_delete=models.CASCADE,
        related_name="steps",
    )
    code = models.SlugField(max_length=80)
    order = models.PositiveIntegerField()
    skill_version = models.ForeignKey(
        SkillVersion,
        on_delete=models.PROTECT,
        related_name="agent_bindings",
    )
    input_mapping = models.JSONField(default=dict)

    class Meta:
        ordering = ["order"]
        constraints = [
            models.UniqueConstraint(
                fields=["agent_version", "code"],
                name="unique_agent_step_code",
            ),
            models.UniqueConstraint(
                fields=["agent_version", "order"],
                name="unique_agent_step_order",
            ),
        ]


class AgentRun(models.Model):
    public_id = models.UUIDField(default=uuid.uuid4, unique=True, editable=False)
    definition = models.ForeignKey(
        AgentDefinition,
        on_delete=models.PROTECT,
        related_name="runs",
    )
    agent_version = models.ForeignKey(
        AgentVersion,
        on_delete=models.PROTECT,
        related_name="runs",
    )
    sandbox = models.BooleanField(default=False)
    status = models.CharField(max_length=24, default="queued")
    input_payload = models.JSONField(default=dict)
    final_output = models.JSONField(null=True, blank=True)
    config_snapshot = models.JSONField(default=dict)
    current_step = models.CharField(max_length=80, blank=True)
    total_steps = models.PositiveIntegerField(default=0)
    input_tokens = models.PositiveIntegerField(default=0)
    output_tokens = models.PositiveIntegerField(default=0)
    duration_ms = models.PositiveIntegerField(default=0)
    task_id = models.CharField(max_length=120, blank=True)
    cancellation_requested_at = models.DateTimeField(null=True, blank=True)
    error_code = models.CharField(max_length=80, blank=True)
    error_message = models.TextField(blank=True)
    created_by = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="agent_runs",
    )
    created_at = models.DateTimeField(auto_now_add=True)
    started_at = models.DateTimeField(null=True, blank=True)
    finished_at = models.DateTimeField(null=True, blank=True)

    class Meta:
        ordering = ["-created_at"]
        indexes = [
            models.Index(fields=["status", "created_at"]),
            models.Index(fields=["definition", "created_at"]),
        ]


class AgentRunStep(models.Model):
    run = models.ForeignKey(AgentRun, on_delete=models.CASCADE, related_name="steps")
    code = models.CharField(max_length=80)
    name = models.CharField(max_length=120)
    order = models.PositiveIntegerField()
    skill_version = models.ForeignKey(
        SkillVersion,
        on_delete=models.PROTECT,
        related_name="run_steps",
    )
    status = models.CharField(max_length=24, default="queued")
    input_payload = models.JSONField(default=dict)
    output_payload = models.JSONField(null=True, blank=True)
    input_tokens = models.PositiveIntegerField(default=0)
    output_tokens = models.PositiveIntegerField(default=0)
    duration_ms = models.PositiveIntegerField(default=0)
    error_code = models.CharField(max_length=80, blank=True)
    error_message = models.TextField(blank=True)
    started_at = models.DateTimeField(null=True, blank=True)
    finished_at = models.DateTimeField(null=True, blank=True)

    class Meta:
        ordering = ["order"]
        constraints = [
            models.UniqueConstraint(fields=["run", "order"], name="unique_run_step_order")
        ]


class AgentRunAttempt(models.Model):
    step = models.ForeignKey(
        AgentRunStep,
        on_delete=models.CASCADE,
        related_name="attempts",
    )
    model_invocation = models.ForeignKey(
        ModelInvocation,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="agent_attempts",
    )
    attempt_type = models.CharField(max_length=32, default="initial")
    status = models.CharField(max_length=24)
    raw_output = models.TextField(blank=True)
    validation_error = models.TextField(blank=True)
    input_tokens = models.PositiveIntegerField(default=0)
    output_tokens = models.PositiveIntegerField(default=0)
    duration_ms = models.PositiveIntegerField(default=0)
    created_at = models.DateTimeField(auto_now_add=True)

    class Meta:
        ordering = ["created_at"]