Artifex/control_plane/agents/models.py

from __future__ import annotations

import uuid
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")