Artifex/agents/progeny.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from typing import Any

from control_plane.agents.models import Agent, AgentPlan, AgentVersion, BenchmarkRun, ImprovementCandidate, ProgenyInvestigation, 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
from graph.models import ExecutionGraphVersion, GraphNodeRun, GraphRun


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


@dataclass(frozen=True)
class ProgenySignalGroup:
    grouping_key: str
    occurrence_count: int
    affected_projects: list[int]
    affected_agents: list[int]
    affected_graph_versions: list[int]
    affected_graph_nodes: list[str]
    first_seen: datetime
    last_seen: datetime
    severity: str
    failure_category: str


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,
        *,
        graph_run: GraphRun | None = None,
        graph_node_run: GraphNodeRun | None = None,
        execution_graph_version: ExecutionGraphVersion | None = None,
        metadata: dict[str, object] | None = None,
    ) -> ProgenySignal:
        severity = "high" if status in ["REWORK_REQUIRED", "REJECTED"] else "info"
        lineage = self._lineage(graph_run, graph_node_run, execution_graph_version)
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            graph_run=graph_run,
            graph_node_run=graph_node_run,
            execution_graph_version=lineage["execution_graph_version"],
            source="reviewer",
            severity=severity,
            failure_category=status,
            summary=summary,
            evidence={"findings": findings, **lineage["evidence"], **(metadata or {})},
            status="OPEN",
            grouping_key=self._grouping_key("reviewer", status, agent_version, lineage),
            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,
        *,
        graph_run: GraphRun | None = None,
        graph_node_run: GraphNodeRun | None = None,
        execution_graph_version: ExecutionGraphVersion | None = None,
        metadata: dict[str, object] | None = None,
    ) -> ProgenySignal:
        severity = "high" if result == "FAIL" else "info"
        lineage = self._lineage(graph_run, graph_node_run, execution_graph_version)
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            graph_run=graph_run,
            graph_node_run=graph_node_run,
            execution_graph_version=lineage["execution_graph_version"],
            source="judge",
            severity=severity,
            failure_category=result,
            summary=summary,
            evidence={"evidence": evidence, **lineage["evidence"], **(metadata or {})},
            status="OPEN",
            grouping_key=self._grouping_key("judge", result, agent_version, lineage),
            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,
        *,
        graph_run: GraphRun | None = None,
        graph_node_run: GraphNodeRun | None = None,
        execution_graph_version: ExecutionGraphVersion | None = None,
    ) -> ProgenySignal:
        lineage = self._lineage(graph_run, graph_node_run, execution_graph_version)
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            graph_run=graph_run,
            graph_node_run=graph_node_run,
            execution_graph_version=lineage["execution_graph_version"],
            source="model_output",
            severity="high",
            failure_category="MODEL_OUTPUT_INVALID",
            summary=f"Model output malformed: {error}",
            evidence={"raw_output": raw_output[:1000], "error": error, **lineage["evidence"]},
            status="OPEN",
            grouping_key=self._grouping_key("model_output", "MODEL_OUTPUT_INVALID", agent_version, lineage),
            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_provider_signal(
        self,
        project: Project | None,
        task: Task | None,
        milestone: Milestone | None,
        agent_version: AgentVersion | None,
        category: str,
        summary: str,
        evidence: dict[str, object],
        *,
        severity: str = "high",
        graph_run: GraphRun | None = None,
        graph_node_run: GraphNodeRun | None = None,
        execution_graph_version: ExecutionGraphVersion | None = None,
        model: str = "",
    ) -> ProgenySignal:
        lineage = self._lineage(graph_run, graph_node_run, execution_graph_version)
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            graph_run=graph_run,
            graph_node_run=graph_node_run,
            execution_graph_version=lineage["execution_graph_version"],
            source="provider",
            severity=severity,
            failure_category=category,
            summary=summary,
            evidence={**evidence, **lineage["evidence"]},
            status="OPEN",
            grouping_key=self._grouping_key("provider", category, agent_version, lineage),
            model=model or (agent_version.model if agent_version else ""),
        )
        self.bus.publish("PROGENY_SIGNAL_CREATED", project=project, task=task, actor="progeny", payload={"signal_id": str(signal.id), "source": "provider", "category": category})
        return signal

    def create_retry_exhausted_signal(
        self,
        project: Project,
        task: Task,
        milestone: Milestone,
        agent_version: AgentVersion,
        attempts: int,
        *,
        graph_run: GraphRun | None = None,
        graph_node_run: GraphNodeRun | None = None,
        execution_graph_version: ExecutionGraphVersion | None = None,
        evidence: dict[str, object] | None = None,
    ) -> ProgenySignal:
        lineage = self._lineage(graph_run, graph_node_run, execution_graph_version)
        signal = ProgenySignal.objects.create(
            project=project,
            task=task,
            milestone=milestone,
            agent_version=agent_version,
            graph_run=graph_run,
            graph_node_run=graph_node_run,
            execution_graph_version=lineage["execution_graph_version"],
            source="retry",
            severity="critical",
            failure_category="TASK_RETRY_EXHAUSTED",
            summary=f"Task {task.id} exhausted {attempts} retries",
            evidence={"attempts": attempts, **(evidence or {}), **lineage["evidence"]},
            status="OPEN",
            grouping_key=self._grouping_key("retry", "TASK_RETRY_EXHAUSTED", agent_version, lineage),
            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 create_graph_runtime_signal(
        self,
        graph_run: GraphRun,
        category: str,
        summary: str,
        evidence: dict[str, object],
        *,
        graph_node_run: GraphNodeRun | None = None,
        severity: str = "high",
    ) -> ProgenySignal:
        lineage = self._lineage(graph_run, graph_node_run, graph_run.execution_graph_version)
        signal = ProgenySignal.objects.create(
            project=graph_run.project,
            task=graph_run.task,
            milestone=graph_run.milestone,
            graph_run=graph_run,
            graph_node_run=graph_node_run,
            execution_graph_version=graph_run.execution_graph_version,
            source="graph_runtime",
            severity=severity,
            failure_category=category,
            summary=summary,
            evidence={**evidence, **lineage["evidence"]},
            status="OPEN",
            grouping_key=self._grouping_key("graph_runtime", category, None, lineage),
        )
        self.bus.publish("PROGENY_SIGNAL_CREATED", project=graph_run.project, task=graph_run.task, actor="progeny", payload={"signal_id": str(signal.id), "source": "graph_runtime", "category": category})
        return signal

    def query_inbox(self, **filters: object):
        signals = ProgenySignal.objects.select_related("project", "task", "agent_version", "execution_graph_version", "graph_node_run").all()
        status = filters.get("status", "OPEN")
        if status:
            signals = signals.filter(status=status)
        if filters.get("source"):
            signals = signals.filter(source=filters["source"])
        if filters.get("project"):
            signals = signals.filter(project=filters["project"])
        if filters.get("agent"):
            signals = signals.filter(agent_version__agent=filters["agent"])
        if filters.get("agent_version"):
            signals = signals.filter(agent_version=filters["agent_version"])
        if filters.get("model"):
            signals = signals.filter(model=filters["model"])
        if filters.get("execution_graph"):
            signals = signals.filter(execution_graph_version__graph=filters["execution_graph"])
        if filters.get("execution_graph_version"):
            signals = signals.filter(execution_graph_version=filters["execution_graph_version"])
        if filters.get("graph_node"):
            signals = signals.filter(graph_node_run__node_id=filters["graph_node"])
        if filters.get("severity"):
            signals = signals.filter(severity=filters["severity"])
        if filters.get("failure_category"):
            signals = signals.filter(failure_category=filters["failure_category"])
        if filters.get("start"):
            signals = signals.filter(created_at__gte=filters["start"])
        if filters.get("end"):
            signals = signals.filter(created_at__lte=filters["end"])
        return signals.order_by("-created_at")

    def group_unresolved_signals(self, **filters: object) -> list[ProgenySignalGroup]:
        signals = list(self.query_inbox(**filters))
        buckets: dict[str, list[ProgenySignal]] = {}
        for signal in signals:
            buckets.setdefault(signal.grouping_key or self._fallback_grouping_key(signal), []).append(signal)
        groups: list[ProgenySignalGroup] = []
        severity_rank = {"info": 0, "INFO": 0, "low": 1, "medium": 2, "high": 3, "critical": 4}
        for key, bucket in buckets.items():
            ordered = sorted(bucket, key=lambda signal: signal.created_at)
            groups.append(
                ProgenySignalGroup(
                    grouping_key=key,
                    occurrence_count=len(bucket),
                    affected_projects=sorted({str(signal.project_id) for signal in bucket if signal.project_id}),
                    affected_agents=sorted({str(signal.agent_version_id) for signal in bucket if signal.agent_version_id}),
                    affected_graph_versions=sorted({str(signal.execution_graph_version_id) for signal in bucket if signal.execution_graph_version_id}),
                    affected_graph_nodes=sorted({signal.graph_node_run.node_id for signal in bucket if signal.graph_node_run_id}),
                    first_seen=ordered[0].created_at,
                    last_seen=ordered[-1].created_at,
                    severity=max((signal.severity for signal in bucket), key=lambda value: severity_rank.get(value, 0)),
                    failure_category=ordered[-1].failure_category,
                )
            )
        return sorted(groups, key=lambda group: (group.severity == "critical", group.occurrence_count, group.last_seen), reverse=True)

    def create_smart_investigation(self, grouping_key: str) -> ProgenyInvestigation:
        signals = list(self.query_inbox().filter(grouping_key=grouping_key).order_by("created_at"))
        if not signals:
            raise ValueError(f"No open signals for grouping key {grouping_key}")
        analysis = self._analyze_signals(signals)
        investigation = ProgenyInvestigation.objects.create(
            signal_clusters=[{"grouping_key": grouping_key, "signal_ids": [str(signal.id) for signal in signals], "occurrence_count": len(signals)}],
            affected_projects=sorted({str(signal.project_id) for signal in signals if signal.project_id}),
            affected_agents=sorted({str(signal.agent_version_id) for signal in signals if signal.agent_version_id}),
            affected_graph_versions=sorted({str(signal.execution_graph_version_id) for signal in signals if signal.execution_graph_version_id}),
            affected_nodes=sorted({signal.graph_node_run.node_id for signal in signals if signal.graph_node_run_id}),
            hypotheses=analysis["hypotheses"],
            recommended_target=str(analysis["target"]),
            confidence=float(analysis["confidence"]),
            recommended_route=str(analysis["route"]),
            proposed_experiments=analysis["experiments"],
            expected_impact=str(analysis["impact"]),
            estimated_cost=str(analysis["cost"]),
        )
        investigation.signals.set(signals)
        return investigation

    def create_improvement_candidate(self, investigation: ProgenyInvestigation, hypothesis: str | None = None) -> ImprovementCandidate:
        signal = investigation.signals.select_related("execution_graph_version", "agent_version").first()
        target_type = investigation.recommended_target
        execution_graph_version = signal.execution_graph_version if signal and target_type in {"WORKFLOW_GRAPH", "GRAPH_NODE"} else None
        agent_version = signal.agent_version if signal and target_type in {"AGENT", "REVIEWER", "JUDGE"} else None
        target_id = ""
        target_label = target_type
        if execution_graph_version is not None:
            target_id = str(execution_graph_version.id)
            target_label = f"{execution_graph_version.graph.name} v{execution_graph_version.version}"
        elif agent_version is not None:
            target_id = str(agent_version.id)
            target_label = f"{agent_version.agent.name} v{agent_version.version}"
        return ImprovementCandidate.objects.create(
            investigation=investigation,
            target_type=target_type,
            target_id=target_id,
            target_label=target_label,
            hypothesis=hypothesis or str((investigation.hypotheses or [{}])[0].get("hypothesis", "Improve target based on Progeny investigation evidence.")),
            recommended_route=investigation.recommended_route,
            evidence={"investigation_id": str(investigation.id), "signal_clusters": investigation.signal_clusters},
            execution_graph_version=execution_graph_version,
            agent_version=agent_version,
        )

    def list_investigations(self, status: str | None = None):
        investigations = ProgenyInvestigation.objects.all()
        if status:
            investigations = investigations.filter(status=status)
        return investigations.order_by("-created_at")

    def list_improvements(self, status: str | None = None):
        candidates = ImprovementCandidate.objects.select_related("investigation", "execution_graph_version", "agent_version").all()
        if status:
            candidates = candidates.filter(status=status)
        return candidates.order_by("-created_at")

    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)),
        }

    def _lineage(
        self,
        graph_run: GraphRun | None,
        graph_node_run: GraphNodeRun | None,
        execution_graph_version: ExecutionGraphVersion | None,
    ) -> dict[str, Any]:
        version = execution_graph_version or (graph_run.execution_graph_version if graph_run else None)
        evidence: dict[str, object] = {}
        if graph_run is not None:
            evidence["graph_run_id"] = graph_run.id
        if version is not None:
            evidence["execution_graph_version_id"] = version.id
            evidence["execution_graph"] = version.graph.name
            evidence["execution_graph_version"] = version.version
        if graph_node_run is not None:
            evidence["graph_node_run_id"] = graph_node_run.id
            evidence["node_id"] = graph_node_run.node_id
            evidence["node_type"] = graph_node_run.node_type
            evidence["visit_index"] = graph_node_run.visit_index
        return {"execution_graph_version": version, "evidence": evidence}

    def _grouping_key(self, source: str, category: str, agent_version: AgentVersion | None, lineage: dict[str, Any]) -> str:
        evidence = lineage["evidence"]
        graph = evidence.get("execution_graph", "no_graph")
        graph_version = evidence.get("execution_graph_version", "no_version")
        node_type = evidence.get("node_type", "no_node")
        agent = f"agent:{agent_version.id}" if agent_version else "agent:none"
        return f"{source}:{category}:{graph}:v{graph_version}:{node_type}:{agent}"

    def _fallback_grouping_key(self, signal: ProgenySignal) -> str:
        node_type = signal.graph_node_run.node_type if signal.graph_node_run_id else "no_node"
        graph_version = signal.execution_graph_version.version if signal.execution_graph_version_id else "no_version"
        fingerprint = str(signal.evidence.get("fingerprint") or signal.evidence.get("type") or signal.summary[:80]).lower()
        return f"{signal.source}:{signal.failure_category}:v{graph_version}:{node_type}:{fingerprint}"

    def _analyze_signals(self, signals: list[ProgenySignal]) -> dict[str, object]:
        corpus_parts: list[str] = []
        for signal in signals:
            corpus_parts.extend(
                [
                    signal.source,
                    signal.failure_category,
                    signal.summary,
                    str(signal.evidence),
                    signal.graph_node_run.node_type if signal.graph_node_run_id else "",
                ]
            )
        corpus = " ".join(corpus_parts).lower()
        sources = {signal.source for signal in signals}
        node_types = {signal.graph_node_run.node_type for signal in signals if signal.graph_node_run_id}
        graph_versions = {str(signal.execution_graph_version_id) for signal in signals if signal.execution_graph_version_id}
        target = "NO_SYSTEMIC_CHANGE"
        route = "record evidence, no intervention"
        confidence = 0.45
        experiments = ["Review representative signal evidence manually before changing production behavior."]
        impact = "Avoid unnecessary system changes when evidence is project-specific."
        cost = "low"

        if "unsupported operation" in corpus or "missing capability" in corpus:
            target = "TOOL_POLICY"
            route = "Progeny"
            confidence = 0.82
            experiments = ["Replay affected task with candidate tool policy that grants the missing operation."]
            impact = "Reduce repeated task failures caused by unavailable safe mutations."
        elif "malformed json" in corpus or "model_output_invalid" in corpus:
            target = "MODEL"
            route = "Model Studio / Model Router"
            confidence = 0.78
            experiments = ["Replay prompts with stricter response-format contract and compare valid-output rate."]
            impact = "Reduce invalid model responses before they reach mutation tools."
        elif "path" in corpus or "timeout" in corpus or "provider_unavailable" in corpus or "provider_timeout" in corpus or "environment" in corpus:
            target = "INFRASTRUCTURE"
            route = "Steward"
            confidence = 0.76
            experiments = ["Replay with captured environment and provider health checks before changing agents."]
            impact = "Separate environmental breakage from agent/model quality issues."
        elif "graph_runtime" in sources or (len(graph_versions) == 1 and len(node_types) == 1 and len(signals) > 1):
            target = "GRAPH_NODE" if node_types else "WORKFLOW_GRAPH"
            route = "Progeny Graph Evolution"
            confidence = 0.74
            experiments = ["Replay the cluster against a challenger graph version with adjusted node policy or transition handling."]
            impact = "Reduce recurring workflow-node failures without changing task implementation agents."
            cost = "medium"
        elif "false" in corpus and "review" in corpus:
            target = "REVIEWER"
            route = "Progeny"
            confidence = 0.7
            experiments = ["Replay accepted diffs against a reviewer challenger with calibrated route-detection criteria."]
            impact = "Reduce false rejections while preserving quality gates."
        elif signals[0].task_id and len({signal.task_id for signal in signals}) == 1 and len(signals) == 1:
            target = "PROJECT_INTENT"
            route = "Repair"
            confidence = 0.55
            experiments = ["Inspect project-specific assertion and task acceptance criteria before system changes."]
            impact = "Resolve the isolated task without overfitting global behavior."
        hypothesis = {
            "target": target,
            "hypothesis": f"Evidence from {len(signals)} signal(s) points to {target} as the likely root-cause target.",
            "supporting_evidence": [str(signal.id) for signal in signals],
            "graph_nodes": sorted(node_types),
            "graph_versions": sorted(graph_versions),
        }
        return {"target": target, "route": route, "confidence": confidence, "experiments": experiments, "impact": impact, "cost": cost, "hypotheses": [hypothesis]}