Artifex/graph/models.py

from __future__ import annotations

from django.db import models


class ExecutionGraphVersionStatus(models.TextChoices):
    DRAFT = "DRAFT"
    CHALLENGER = "CHALLENGER"
    CHAMPION = "CHAMPION"
    RETIRED = "RETIRED"


class GraphRunStatus(models.TextChoices):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    PAUSED = "PAUSED"
    COMPLETE = "COMPLETE"
    FAILED = "FAILED"
    CANCELLED = "CANCELLED"


class GraphNodeRunStatus(models.TextChoices):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    PAUSED = "PAUSED"
    COMPLETE = "COMPLETE"
    FAILED = "FAILED"
    SKIPPED = "SKIPPED"


class ExecutionGraphDefinition(models.Model):
    name = models.CharField(max_length=160, unique=True)
    graph_type = models.CharField(max_length=80)
    description = models.TextField(blank=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    def __str__(self) -> str:
        return self.name


class ExecutionGraphVersion(models.Model):
    graph = models.ForeignKey(ExecutionGraphDefinition, on_delete=models.CASCADE, related_name="versions")
    version = models.PositiveIntegerField()
    status = models.CharField(max_length=32, choices=ExecutionGraphVersionStatus.choices, default=ExecutionGraphVersionStatus.DRAFT)
    graph_spec = models.JSONField(default=dict)
    metadata = models.JSONField(default=dict, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)
    promoted_at = models.DateTimeField(null=True, blank=True)

    class Meta:
        constraints = [models.UniqueConstraint(fields=["graph", "version"], name="unique_execution_graph_version")]

    def __str__(self) -> str:
        return f"{self.graph.name} v{self.version}"


class GraphRun(models.Model):
    execution_graph_version = models.ForeignKey(ExecutionGraphVersion, on_delete=models.PROTECT, related_name="runs")
    project = models.ForeignKey("projects.Project", on_delete=models.CASCADE, null=True, blank=True, related_name="graph_runs")
    milestone = models.ForeignKey("projects.Milestone", on_delete=models.SET_NULL, null=True, blank=True, related_name="graph_runs")
    feature = models.ForeignKey("projects.Feature", on_delete=models.SET_NULL, null=True, blank=True, related_name="graph_runs")
    task = models.ForeignKey("projects.Task", on_delete=models.SET_NULL, null=True, blank=True, related_name="graph_runs")
    task_attempt = models.ForeignKey("projects.TaskAttempt", on_delete=models.SET_NULL, null=True, blank=True, related_name="graph_runs")
    status = models.CharField(max_length=32, choices=GraphRunStatus.choices, default=GraphRunStatus.PENDING)
    started_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)
    current_node = models.CharField(max_length=120, blank=True)
    failure_reason = models.TextField(blank=True)
    metadata = models.JSONField(default=dict, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

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


class GraphNodeRun(models.Model):
    graph_run = models.ForeignKey(GraphRun, on_delete=models.CASCADE, related_name="node_runs")
    node_id = models.CharField(max_length=120)
    node_type = models.CharField(max_length=120)
    visit_index = models.PositiveIntegerField(default=1)
    status = models.CharField(max_length=32, choices=GraphNodeRunStatus.choices, default=GraphNodeRunStatus.PENDING)
    started_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)
    agent_version = models.ForeignKey("agents.AgentVersion", on_delete=models.SET_NULL, null=True, blank=True, related_name="graph_node_runs")
    model_request = models.ForeignKey("resources.ModelRequest", on_delete=models.SET_NULL, null=True, blank=True, related_name="graph_node_runs")
    input_metadata = models.JSONField(default=dict, blank=True)
    output_metadata = models.JSONField(default=dict, blank=True)
    failure_evidence = models.JSONField(default=dict, blank=True)
    telemetry = models.JSONField(default=dict, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        constraints = [models.UniqueConstraint(fields=["graph_run", "node_id", "visit_index"], name="unique_graph_node_run_visit")]
        indexes = [models.Index(fields=["node_id", "status"])]


class GraphEdgeTraversal(models.Model):
    graph_run = models.ForeignKey(GraphRun, on_delete=models.CASCADE, related_name="edge_traversals")
    source_node = models.CharField(max_length=120)
    target_node = models.CharField(max_length=120)
    condition = models.CharField(max_length=120, blank=True)
    result = models.CharField(max_length=120, blank=True)
    metadata = models.JSONField(default=dict, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)

    class Meta:
        indexes = [models.Index(fields=["graph_run", "created_at"])]


class GraphApprovalStatus(models.TextChoices):
    PENDING = "PENDING"
    APPROVED = "APPROVED"
    REJECTED = "REJECTED"


class GraphApproval(models.Model):
    graph_run = models.ForeignKey(GraphRun, on_delete=models.CASCADE, related_name="approvals")
    node_run = models.ForeignKey(GraphNodeRun, on_delete=models.CASCADE, null=True, blank=True, related_name="approvals")
    status = models.CharField(max_length=32, choices=GraphApprovalStatus.choices, default=GraphApprovalStatus.PENDING)
    reason = models.CharField(max_length=160)
    payload = models.JSONField(default=dict, blank=True)
    requested_by = models.CharField(max_length=120, default="graph_runtime")
    decided_by = models.CharField(max_length=120, blank=True)
    decided_at = models.DateTimeField(null=True, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)