Artifex/graph/bootstrap.py

from __future__ import annotations

from django.utils import timezone

from graph.agent_control import agent_investigation_graph_v1
from graph.crypto_venture_cohort import crypto_venture_cohort_graph_v1
from graph.lifecycle import (
    project_evolution_graph_v1,
    project_exploration_graph_v1,
    project_extension_graph_v1,
)
from graph.models import (
    ExecutionGraphDefinition,
    ExecutionGraphVersion,
    ExecutionGraphVersionStatus,
)
from graph.roadmap import project_roadmap_review_graph_v1
from graph.scenario_lab import scenario_lab_graph_v1
from graph.steward import steward_run_graph_v1
from graph.story_authoring import story_authoring_graph_v1, story_authoring_graph_v2
from graph.task_execution import task_execution_graph_v1
from graph.venture_cohort import venture_discovery_cohort_graph_v1
from graph.venture_discovery import venture_discovery_graph_v1


def champion_task_execution_graph_v1() -> ExecutionGraphVersion:
    spec = task_execution_graph_v1()
    definition, _ = ExecutionGraphDefinition.objects.get_or_create(
        name=spec.name,
        defaults={"graph_type": spec.graph_type, "description": str(spec.metadata.get("description", ""))},
    )
    version, created = ExecutionGraphVersion.objects.get_or_create(
        graph=definition,
        version=spec.version,
        defaults={
            "status": ExecutionGraphVersionStatus.CHAMPION,
            "graph_spec": spec.to_dict(),
            "metadata": {"immutable_after_use": True},
            "promoted_at": timezone.now(),
        },
    )
    if not created and version.status != ExecutionGraphVersionStatus.CHAMPION:
        version.status = ExecutionGraphVersionStatus.CHAMPION
        version.promoted_at = timezone.now()
        version.save(update_fields=["status", "promoted_at"])
    return version


def champion_steward_run_graph_v1() -> ExecutionGraphVersion:
    spec = steward_run_graph_v1()
    definition, _ = ExecutionGraphDefinition.objects.get_or_create(
        name=spec.name,
        defaults={"graph_type": spec.graph_type, "description": str(spec.metadata.get("description", ""))},
    )
    version, created = ExecutionGraphVersion.objects.get_or_create(
        graph=definition,
        version=spec.version,
        defaults={
            "status": ExecutionGraphVersionStatus.CHAMPION,
            "graph_spec": spec.to_dict(),
            "metadata": {"immutable_after_use": True},
            "promoted_at": timezone.now(),
        },
    )
    if not created and version.status != ExecutionGraphVersionStatus.CHAMPION:
        version.status = ExecutionGraphVersionStatus.CHAMPION
        version.promoted_at = timezone.now()
        version.save(update_fields=["status", "promoted_at"])
    return version


def _champion_graph(spec) -> ExecutionGraphVersion:
    definition, _ = ExecutionGraphDefinition.objects.get_or_create(
        name=spec.name,
        defaults={"graph_type": spec.graph_type, "description": str(spec.metadata.get("description", ""))},
    )
    ExecutionGraphVersion.objects.filter(graph=definition, status=ExecutionGraphVersionStatus.CHAMPION).exclude(version=spec.version).update(status=ExecutionGraphVersionStatus.RETIRED)
    version, created = ExecutionGraphVersion.objects.get_or_create(
        graph=definition,
        version=spec.version,
        defaults={"status": ExecutionGraphVersionStatus.CHAMPION, "graph_spec": spec.to_dict(), "metadata": {"immutable_after_use": True}, "promoted_at": timezone.now()},
    )
    if not created and version.status != ExecutionGraphVersionStatus.CHAMPION:
        ExecutionGraphVersion.objects.filter(graph=definition, status=ExecutionGraphVersionStatus.CHAMPION).exclude(pk=version.pk).update(status=ExecutionGraphVersionStatus.RETIRED)
        version.status = ExecutionGraphVersionStatus.CHAMPION
        version.promoted_at = timezone.now()
        version.save(update_fields=["status", "promoted_at"])
    return version


def champion_project_extension_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(project_extension_graph_v1())


def champion_project_evolution_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(project_evolution_graph_v1())


def champion_project_exploration_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(project_exploration_graph_v1())


def champion_project_roadmap_review_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(project_roadmap_review_graph_v1())


def champion_scenario_lab_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(scenario_lab_graph_v1())


def champion_agent_investigation_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(agent_investigation_graph_v1())


def champion_venture_discovery_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(venture_discovery_graph_v1())


def champion_venture_discovery_cohort_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(venture_discovery_cohort_graph_v1())


def champion_crypto_venture_cohort_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(crypto_venture_cohort_graph_v1())


def champion_story_authoring_graph_v1() -> ExecutionGraphVersion:
    return _champion_graph(story_authoring_graph_v1())


def champion_story_authoring_graph_v2() -> ExecutionGraphVersion:
    return _champion_graph(story_authoring_graph_v2())