Artifex/agents/progeny.py

from __future__ import annotations

from dataclasses import dataclass

from control_plane.agents.models import Agent, AgentPlan, AgentVersion, BenchmarkRun, ProgenySignal, PromotionStatus
from control_plane.events.bus import EventBus
from control_plane.events.models import EventType
from control_plane.projects.models import Project, Task, Milestone


@dataclass(frozen=True)
class BenchmarkDecision:
    decision: str
    metrics: dict[str, float]


class ProgenyService:
    def __init__(self, bus: EventBus | None = None) -> None:
        self.bus = bus or EventBus()

    def create_candidate_from_plan(self, plan: AgentPlan) -> AgentVersion:
        agent = plan.agent or Agent.objects.create(name=plan.name, role=plan.role)
        next_version = (agent.versions.order_by("-version").first().version + 1) if agent.versions.exists() else 1
        version = AgentVersion.objects.create(
            agent=agent,
            version=next_version,
            model=plan.model,
            system_contract=plan.system_contract,
            capabilities=plan.capabilities,
            tools=plan.tools,
            permissions=plan.permissions,
            context_policy=plan.context_policy,
            workflow=plan.workflow,
            retry_policy={"max_retries": 2},
            evaluator=plan.success_criteria,
            promotion_status=PromotionStatus.CHALLENGER,
        )
        if agent.champion_version_id is None:
            agent.champion_version = version
            version.promotion_status = PromotionStatus.CHAMPION
            version.save(update_fields=["promotion_status", "updated_at"])
            agent.save(update_fields=["champion_version", "updated_at"])
        self.bus.publish(EventType.AGENT_CREATED, actor="progeny", payload={"agent": agent.name, "version": version.version})
        return version

    def replay_benchmark(self, champion: AgentVersion, challenger: AgentVersion, benchmark_set: list[dict[str, object]]) -> BenchmarkRun:
        metrics = self._score(challenger, benchmark_set)
        champion_metrics = self._score(champion, benchmark_set)
        decision = "PROMOTE" if metrics["completion_rate"] >= champion_metrics["completion_rate"] and metrics["test_pass_rate"] >= champion_metrics["test_pass_rate"] else "REJECT"
        return BenchmarkRun.objects.create(
            champion=champion,
            challenger=challenger,
            benchmark_set=benchmark_set,
            metrics={"champion": champion_metrics, "challenger": metrics},
            decision=decision,
        )

    def promote_or_reject(self, run: BenchmarkRun) -> BenchmarkDecision:
        challenger = run.challenger
        agent = challenger.agent
        challenger_metrics = run.metrics["challenger"]
        if run.decision == "PROMOTE":
            if agent.champion_version_id:
                old = agent.champion_version
                old.promotion_status = PromotionStatus.CANDIDATE
                old.save(update_fields=["promotion_status", "updated_at"])
            challenger.promotion_status = PromotionStatus.CHAMPION
            challenger.save(update_fields=["promotion_status", "updated_at"])
            agent.champion_version = challenger
            agent.save(update_fields=["champion_version", "updated_at"])
            self.bus.publish(EventType.AGENT_PROMOTED, actor="progeny", payload={"agent": agent.name, "version": challenger.version})
            return BenchmarkDecision("PROMOTED", challenger_metrics)
        challenger.promotion_status = PromotionStatus.REJECTED
        challenger.save(update_fields=["promotion_status", "updated_at"])
        return BenchmarkDecision("REJECTED", challenger_metrics)

    def create_reviewer_signal(self, project: Project, task: Task, milestone: Milestone, agent_version: AgentVersion, status: str, findings: list[dict[str, object]], summary: str) -> ProgenySignal:
        severity = "high" if status in ["REWORK_REQUIRED", "REJECTED"] else "info"
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            source="reviewer",
            severity=severity,
            failure_category=status,
            summary=summary,
            evidence={"findings": findings},
            status="OPEN",
            grouping_key=f"reviewer:{task.id}:{status}",
            model=agent_version.model,
        )
        self.bus.publish("PROGENY_SIGNAL_CREATED", project=project, task=task, actor="progeny", payload={"signal_id": str(signal.id), "source": "reviewer", "status": status})
        return signal

    def create_judge_signal(self, project: Project, task: Task, milestone: Milestone, agent_version: AgentVersion, result: str, evidence: list[dict[str, object]], summary: str) -> ProgenySignal:
        severity = "high" if result == "FAIL" else "info"
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            source="judge",
            severity=severity,
            failure_category=result,
            summary=summary,
            evidence={"evidence": evidence},
            status="OPEN",
            grouping_key=f"judge:{task.id}:{result}",
            model=agent_version.model,
        )
        self.bus.publish("PROGENY_SIGNAL_CREATED", project=project, task=task, actor="progeny", payload={"signal_id": str(signal.id), "source": "judge", "result": result})
        return signal

    def create_model_output_signal(self, project: Project, task: Task, milestone: Milestone, agent_version: AgentVersion, error: str, raw_output: str) -> ProgenySignal:
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            source="model_output",
            severity="high",
            failure_category="MALFORMED_OUTPUT",
            summary=f"Model output malformed: {error}",
            evidence={"raw_output": raw_output[:1000]},
            status="OPEN",
            grouping_key=f"model_output:{task.id}",
            model=agent_version.model,
        )
        self.bus.publish("PROGENY_SIGNAL_CREATED", project=project, task=task, actor="progeny", payload={"signal_id": str(signal.id), "source": "model_output"})
        return signal

    def create_retry_exhausted_signal(self, project: Project, task: Task, milestone: Milestone, agent_version: AgentVersion, attempts: int) -> ProgenySignal:
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            source="retry",
            severity="critical",
            failure_category="RETRY_EXHAUSTED",
            summary=f"Task {task.id} exhausted {attempts} retries",
            evidence={"attempts": attempts},
            status="OPEN",
            grouping_key=f"retry:{task.id}",
            model=agent_version.model,
        )
        self.bus.publish("PROGENY_SIGNAL_CREATED", project=project, task=task, actor="progeny", payload={"signal_id": str(signal.id), "source": "retry"})
        return signal

    def _score(self, version: AgentVersion, benchmark_set: list[dict[str, object]]) -> dict[str, float]:
        if not benchmark_set:
            return {"completion_rate": 0.0, "test_pass_rate": 0.0, "review_acceptance": 0.0, "tokens": 0.0, "runtime": 0.0}
        base = 1.0 if "do not self-certify" in version.system_contract.lower() else 0.8
        return {
            "completion_rate": base,
            "test_pass_rate": base,
            "review_acceptance": base,
            "tokens": float(len(version.system_contract.split())),
            "runtime": float(len(benchmark_set)),
        }