Artifex/control_plane/agents/models.py

from __future__ import annotations

import uuid
from django.core.exceptions import ValidationError
from django.db import models

from control_plane.common import TimestampedModel


class AgentRole(models.TextChoices):
    PROJECT_ARCHAEOLOGIST = "PROJECT_ARCHAEOLOGIST"
    PLANNER = "PLANNER"
    CODER = "CODER"
    REVIEWER = "REVIEWER"
    PROJECT_JUDGE = "PROJECT_JUDGE"


class PromotionStatus(models.TextChoices):
    CANDIDATE = "CANDIDATE"
    CHALLENGER = "CHALLENGER"
    CHAMPION = "CHAMPION"
    REJECTED = "REJECTED"


class Agent(TimestampedModel):
    name = models.CharField(max_length=200, unique=True)
    role = models.CharField(max_length=80, choices=AgentRole.choices)
    description = models.TextField(blank=True)
    champion_version = models.ForeignKey(
        "AgentVersion", on_delete=models.SET_NULL, null=True, blank=True, related_name="championed_by"
    )


class AgentVersion(TimestampedModel):
    agent = models.ForeignKey(Agent, on_delete=models.CASCADE, related_name="versions")
    version = models.PositiveIntegerField()
    model = models.CharField(max_length=120)
    system_contract = models.TextField()
    capabilities = models.JSONField(default=list, blank=True)
    tools = models.JSONField(default=list, blank=True)
    permissions = models.JSONField(default=dict, blank=True)
    context_policy = models.JSONField(default=dict, blank=True)
    workflow = models.JSONField(default=dict, blank=True)
    retry_policy = models.JSONField(default=dict, blank=True)
    evaluator = models.JSONField(default=dict, blank=True)
    benchmark_status = models.CharField(max_length=32, default="UNBENCHMARKED")
    promotion_status = models.CharField(max_length=32, choices=PromotionStatus.choices, default=PromotionStatus.CANDIDATE)

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


class AgentPlan(TimestampedModel):
    agent = models.ForeignKey(Agent, on_delete=models.CASCADE, related_name="plans", null=True, blank=True)
    name = models.CharField(max_length=200)
    role = models.CharField(max_length=80, choices=AgentRole.choices)
    model = models.CharField(max_length=120)
    system_contract = models.TextField()
    capabilities = models.JSONField(default=list, blank=True)
    tools = models.JSONField(default=list, blank=True)
    permissions = models.JSONField(default=dict, blank=True)
    context_policy = models.JSONField(default=dict, blank=True)
    workflow = models.JSONField(default=dict, blank=True)
    benchmarks = models.JSONField(default=list, blank=True)
    success_criteria = models.JSONField(default=dict, blank=True)
    created_by = models.CharField(max_length=120, default="sol")


class AgentRun(TimestampedModel):
    agent_version = models.ForeignKey(AgentVersion, on_delete=models.PROTECT, related_name="runs")
    project = models.ForeignKey("projects.Project", on_delete=models.CASCADE, null=True, blank=True)
    task = models.ForeignKey("projects.Task", on_delete=models.SET_NULL, null=True, blank=True)
    status = models.CharField(max_length=32, default="QUEUED")
    input_contract = models.JSONField(default=dict, blank=True)
    output_contract = models.JSONField(default=dict, blank=True)
    metrics = models.JSONField(default=dict, blank=True)


class BenchmarkRun(TimestampedModel):
    champion = models.ForeignKey(AgentVersion, on_delete=models.PROTECT, related_name="champion_benchmarks")
    challenger = models.ForeignKey(AgentVersion, on_delete=models.PROTECT, related_name="challenger_benchmarks")
    benchmark_set = models.JSONField(default=list, blank=True)
    metrics = models.JSONField(default=dict, blank=True)
    decision = models.CharField(max_length=32, default="PENDING")


class ProgenySignal(TimestampedModel):
    project = models.ForeignKey("projects.Project", on_delete=models.SET_NULL, null=True, blank=True)
    task = models.ForeignKey("projects.Task", on_delete=models.SET_NULL, null=True, blank=True)
    milestone = models.ForeignKey("projects.Milestone", on_delete=models.SET_NULL, null=True, blank=True)
    agent_version = models.ForeignKey(AgentVersion, on_delete=models.SET_NULL, null=True, blank=True)
    graph_run = models.ForeignKey("graph.GraphRun", on_delete=models.SET_NULL, null=True, blank=True)
    graph_node_run = models.ForeignKey("graph.GraphNodeRun", on_delete=models.SET_NULL, null=True, blank=True)
    execution_graph_version = models.ForeignKey("graph.ExecutionGraphVersion", on_delete=models.SET_NULL, null=True, blank=True)
    source = models.CharField(max_length=80)
    severity = models.CharField(max_length=32, default="INFO")
    failure_category = models.CharField(max_length=80, blank=True)
    summary = models.TextField()
    evidence = models.JSONField(default=dict, blank=True)
    status = models.CharField(max_length=32, default="OPEN")
    grouping_key = models.CharField(max_length=120, blank=True)
    model = models.CharField(max_length=120, blank=True)


class ProgenyInvestigation(TimestampedModel):
    status = models.CharField(max_length=32, default="OPEN")
    signals = models.ManyToManyField(ProgenySignal, related_name="investigations", blank=True)
    signal_clusters = models.JSONField(default=list, blank=True)
    affected_projects = models.JSONField(default=list, blank=True)
    affected_agents = models.JSONField(default=list, blank=True)
    affected_graph_versions = models.JSONField(default=list, blank=True)
    affected_nodes = models.JSONField(default=list, blank=True)
    hypotheses = models.JSONField(default=list, blank=True)
    recommended_target = models.CharField(max_length=80)
    confidence = models.FloatField(default=0.0)
    recommended_route = models.CharField(max_length=120)
    proposed_experiments = models.JSONField(default=list, blank=True)
    expected_impact = models.TextField(blank=True)
    estimated_cost = models.CharField(max_length=80, blank=True)


class ImprovementCandidate(TimestampedModel):
    status = models.CharField(max_length=32, default="PROPOSED")
    target_type = models.CharField(max_length=80)
    target_id = models.CharField(max_length=120, blank=True)
    target_label = models.CharField(max_length=240, blank=True)
    hypothesis = models.TextField()
    recommended_route = models.CharField(max_length=120, blank=True)
    evidence = models.JSONField(default=dict, blank=True)
    investigation = models.ForeignKey(ProgenyInvestigation, on_delete=models.SET_NULL, null=True, blank=True, related_name="improvement_candidates")
    execution_graph_version = models.ForeignKey("graph.ExecutionGraphVersion", on_delete=models.SET_NULL, null=True, blank=True, related_name="improvement_candidates")
    agent_version = models.ForeignKey(AgentVersion, on_delete=models.SET_NULL, null=True, blank=True, related_name="improvement_candidates")


class ReplayDataset(TimestampedModel):
    name = models.CharField(max_length=200)
    description = models.TextField(blank=True)
    version = models.PositiveIntegerField(default=1)
    status = models.CharField(max_length=32, default="DRAFT")
    selection_criteria = models.JSONField(default=dict, blank=True)
    frozen_at = models.DateTimeField(null=True, blank=True)
    metadata = models.JSONField(default=dict, blank=True)

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


class ReplayCase(TimestampedModel):
    replay_dataset = models.ForeignKey(ReplayDataset, on_delete=models.CASCADE, related_name="cases")
    source_project = models.ForeignKey("projects.Project", on_delete=models.SET_NULL, null=True, blank=True, related_name="replay_cases")
    source_task = models.ForeignKey("projects.Task", on_delete=models.SET_NULL, null=True, blank=True, related_name="replay_cases")
    source_task_attempt = models.ForeignKey("projects.TaskAttempt", on_delete=models.SET_NULL, null=True, blank=True, related_name="replay_cases")
    source_graph_run = models.ForeignKey("graph.GraphRun", on_delete=models.SET_NULL, null=True, blank=True, related_name="replay_cases")
    task_type = models.CharField(max_length=80)
    project_type = models.CharField(max_length=100, blank=True)
    original_goal = models.TextField()
    acceptance_criteria = models.JSONField(default=list, blank=True)
    repository_path = models.TextField(blank=True)
    repository_baseline_ref = models.CharField(max_length=80)
    expected_evaluator_inputs = models.JSONField(default=dict, blank=True)
    failure_classification = models.CharField(max_length=120, blank=True)
    selection_metadata = models.JSONField(default=dict, blank=True)

    def save(self, *args: object, **kwargs: object) -> None:
        if self.replay_dataset.status == "FROZEN":
            raise ValidationError("Frozen replay dataset membership is immutable; create a new dataset version.")
        super().save(*args, **kwargs)


class ProgenyExperiment(TimestampedModel):
    investigation = models.ForeignKey(ProgenyInvestigation, on_delete=models.SET_NULL, null=True, blank=True, related_name="experiments")
    improvement_candidate = models.ForeignKey(ImprovementCandidate, on_delete=models.SET_NULL, null=True, blank=True, related_name="experiments")
    target_type = models.CharField(max_length=80)
    target_identifier = models.CharField(max_length=240, blank=True)
    replay_dataset = models.ForeignKey(ReplayDataset, on_delete=models.PROTECT, related_name="experiments")
    hypothesis = models.TextField()
    success_criteria = models.JSONField(default=dict, blank=True)
    status = models.CharField(max_length=32, default="DRAFT")
    started_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)
    metadata = models.JSONField(default=dict, blank=True)


class ExperimentVariant(TimestampedModel):
    experiment = models.ForeignKey(ProgenyExperiment, on_delete=models.CASCADE, related_name="variants")
    role = models.CharField(max_length=32)
    target_type = models.CharField(max_length=80)
    target_reference = models.CharField(max_length=240)
    execution_graph_version = models.ForeignKey("graph.ExecutionGraphVersion", on_delete=models.SET_NULL, null=True, blank=True, related_name="experiment_variants")
    agent_version = models.ForeignKey(AgentVersion, on_delete=models.SET_NULL, null=True, blank=True, related_name="experiment_variants")
    configuration_snapshot = models.JSONField(default=dict)
    metadata = models.JSONField(default=dict, blank=True)

    def save(self, *args: object, **kwargs: object) -> None:
        if not self._state.adding:
            previous = ExperimentVariant.objects.get(pk=self.pk)
            if previous.configuration_snapshot != self.configuration_snapshot or previous.target_reference != self.target_reference:
                raise ValidationError("Experiment variant configuration is immutable; create a new variant.")
        super().save(*args, **kwargs)


class ReplayRun(TimestampedModel):
    experiment = models.ForeignKey(ProgenyExperiment, on_delete=models.CASCADE, related_name="replay_runs")
    variant = models.ForeignKey(ExperimentVariant, on_delete=models.CASCADE, related_name="replay_runs")
    replay_case = models.ForeignKey(ReplayCase, on_delete=models.CASCADE, related_name="replay_runs")
    graph_run = models.ForeignKey("graph.GraphRun", on_delete=models.SET_NULL, null=True, blank=True, related_name="replay_runs")
    replay_task = models.ForeignKey("projects.Task", on_delete=models.SET_NULL, null=True, blank=True, related_name="replay_runs")
    started_at = models.DateTimeField(null=True, blank=True)
    completed_at = models.DateTimeField(null=True, blank=True)
    status = models.CharField(max_length=32, default="PENDING")
    failure_classification = models.CharField(max_length=80, blank=True)
    commit_candidate_sha = models.CharField(max_length=64, blank=True)
    telemetry = models.JSONField(default=dict, blank=True)
    failure_evidence = models.JSONField(default=dict, blank=True)
    metadata = models.JSONField(default=dict, blank=True)


class ReplayResult(TimestampedModel):
    replay_run = models.OneToOneField(ReplayRun, on_delete=models.CASCADE, related_name="result")
    completion_status = models.CharField(max_length=32)
    tests_status = models.CharField(max_length=32, blank=True)
    reviewer_status = models.CharField(max_length=32, blank=True)
    judge_status = models.CharField(max_length=32, blank=True)
    accepted_candidate = models.BooleanField(default=False)
    metrics = models.JSONField(default=dict, blank=True)
    safety = models.JSONField(default=dict, blank=True)
    evidence = models.JSONField(default=dict, blank=True)


class ExperimentComparison(TimestampedModel):
    experiment = models.OneToOneField(ProgenyExperiment, on_delete=models.CASCADE, related_name="comparison")
    champion_variant = models.ForeignKey(ExperimentVariant, on_delete=models.SET_NULL, null=True, blank=True, related_name="champion_comparisons")
    challenger_variant = models.ForeignKey(ExperimentVariant, on_delete=models.SET_NULL, null=True, blank=True, related_name="challenger_comparisons")
    aggregate_metrics = models.JSONField(default=dict, blank=True)
    paired_outcomes = models.JSONField(default=dict, blank=True)
    regression_cases = models.JSONField(default=list, blank=True)
    verdict = models.CharField(max_length=80, default="INCONCLUSIVE")
    reasons = models.JSONField(default=list, blank=True)
    approved_at = models.DateTimeField(null=True, blank=True)
    rejected_at = models.DateTimeField(null=True, blank=True)
    decided_by = models.CharField(max_length=120, blank=True)