Merge pull request #999 from dimagi/cs/issue_988

Archive pipeline versions when archiving experiment versions

apps/assistants/tables.py

@@ -28,7 +28,7 @@ class OpenAiAssistantTable(tables.Table):
                 title="Archive",
                 icon_class="fa-solid fa-box-archive",
                 required_permissions=["assistants.delete_openaiassistant"],
-                confirm_message="This will only not delete the assistant from OpenAI.",
+                confirm_message="This will not delete the assistant from OpenAI.",
                 hx_method="delete",
             ),
             actions.AjaxAction(

apps/assistants/views.py

@@ -223,6 +223,7 @@ def delete(self, request, team_slug: str, pk: int):
             response = render_to_string(
                 "assistants/partials/referenced_objects.html",
                 context={
+                    "object_name": "assistant",
                     "experiments": experiments,
                     "pipeline_nodes": pipeline_nodes,
                 },

apps/experiments/models.py

@@ -770,13 +770,22 @@ def compare_with_latest(self):
 
     @transaction.atomic()
     def archive(self):
+        """
+        Archive the experiment and all versions in the case where this is the working version. The linked assistant and
+        pipeline for the working version should not be archived.
+        """
         super().archive()
+        self.static_triggers.update(is_archived=True)
+
         if self.is_working_version:
             self.delete_experiment_channels()
             self.versions.update(is_archived=True, audit_action=AuditAction.AUDIT)
             self.scheduled_messages.all().delete()
-        elif self.assistant:
-            self.assistant.archive()
+        else:
+            if self.assistant:
+                self.assistant.archive()
+            elif self.pipeline:
+                self.pipeline.archive()
 
     def delete_experiment_channels(self):
         from apps.channels.models import ExperimentChannel

apps/experiments/tests/test_models.py

@@ -994,7 +994,17 @@ def test_get_version(self, experiment):
     def test_archive_working_experiment(self, experiment):
         """Test that archiving an experiment archives all versions and deletes its scheduled messages"""
         session = ExperimentSessionFactory(experiment=experiment)
-        event_action, _ = self._construct_event_action(time_period=TimePeriod.DAYS, experiment_id=experiment.id)
+        event_action = EventActionFactory(
+            action_type=EventActionType.SCHEDULETRIGGER,
+            params={
+                "name": "Test",
+                "time_period": TimePeriod.DAYS,
+                "frequency": 1,
+                "repetitions": 1,
+                "prompt_text": "hi",
+                "experiment_id": experiment.id,
+            },
+        )
         ScheduledMessageFactory(
             experiment=experiment, team=experiment.team, participant=session.participant, action=event_action
         )
@@ -1012,19 +1022,61 @@ def test_archive_working_experiment(self, experiment):
         assert second_version.is_archived is True
         assert ScheduledMessage.objects.filter(experiment=experiment).exists() is False
 
-    def _construct_event_action(self, time_period: TimePeriod, experiment_id: int, frequency=1, repetitions=1) -> tuple:
-        params = self._get_params(experiment_id, time_period, frequency, repetitions)
-        return EventActionFactory(params=params, action_type=EventActionType.SCHEDULETRIGGER), params
+    @patch("apps.assistants.tasks.delete_openai_assistant_task.delay")
+    @patch("apps.assistants.sync.push_assistant_to_openai", Mock())
+    def test_archive_with_assistant(self, delete_openai_assistant_task):
+        """
+        Archiving an assistant experiment should only archive the assistant when it is not still being referenced by
+        another experiment or pipeline or when referenced by the working version
+        """
 
-    def _get_params(self, experiment_id: int, time_period: TimePeriod = TimePeriod.DAYS, frequency=1, repetitions=1):
-        return {
-            "name": "Test",
-            "time_period": time_period,
-            "frequency": frequency,
-            "repetitions": repetitions,
-            "prompt_text": "hi",
-            "experiment_id": experiment_id,
-        }
+        assistant = OpenAiAssistantFactory()
+        experiment1 = ExperimentFactory(assistant=assistant)
+        experiment2 = ExperimentFactory(assistant=assistant)
+
+        # Version assistant should be archived and removed from OpenAI. Only the version points to it
+        new_version = experiment1.create_new_version()
+        assistant_version = new_version.assistant
+        new_version.archive()
+        delete_openai_assistant_task.assert_called_with(assistant_version.id)
+        delete_openai_assistant_task.reset_mock()
+        self._assert_archived(assistant_version, True)
+
+        experiment1.archive()
+        self._assert_archived(experiment1, True)
+        self._assert_archived(assistant, False)
+
+        experiment2.archive()
+        self._assert_archived(experiment2, True)
+        self._assert_archived(assistant, False)
+
+        delete_openai_assistant_task.assert_not_called()
+
+    @patch("apps.assistants.tasks.delete_openai_assistant_task.delay")
+    @patch("apps.assistants.sync.push_assistant_to_openai", Mock())
+    def test_archive_with_pipeline(self, delete_openai_assistant_task):
+        assistant = OpenAiAssistantFactory()
+        pipeline = PipelineFactory()
+        NodeFactory(pipeline=pipeline, type="AssistantNode", params={"assistant_id": assistant.id})
+        experiment = ExperimentFactory(pipeline=pipeline)
+
+        # For a version, the pipeline should be archived as well as the assistant that it references
+        new_version = experiment.create_new_version()
+        assert assistant.versions.count() == 1
+        assistant_version = assistant.versions.first()
+        new_version.archive()
+        self._assert_archived(new_version.pipeline, True)
+        self._assert_archived(assistant_version, True)
+        delete_openai_assistant_task.assert_called_with(assistant_version.id)
+        delete_openai_assistant_task.reset_mock()
+
+        experiment.archive()
+        self._assert_archived(experiment.pipeline, False)
+        self._assert_archived(assistant, False)
+
+    def _assert_archived(self, model_obj, archived: bool):
+        model_obj.refresh_from_db(fields=["is_archived"])
+        assert model_obj.is_archived == archived
 
 
 @pytest.mark.django_db()

apps/pipelines/models.py

@@ -262,13 +262,46 @@ def create_new_version(self, *args, **kwargs):
         return pipeline_version
 
     @transaction.atomic()
-    def archive(self):
+    def archive(self) -> bool:
+        """
+        Archive this record only when it is not still being referenced by other records. If this record is the working
+        version, all of its versions will be archived as well. The same goes for its nodes.
+        """
+        if self.get_related_experiments_queryset().exists():
+            return False
+
+        if len(self.get_static_trigger_experiment_ids()) > 0:
+            return False
+
         super().archive()
-        self.node_set.update(is_archived=True)
+        for node in self.node_set.all():
+            node.archive()
+
+        if self.is_working_version:
+            for version in self.versions.filter(is_archived=False):
+                version.archive()
+
+        return True
 
     def get_node_param_values(self, node_cls, param_name: str) -> list:
         return list(self.node_set.filter(type=node_cls.__name__).values_list(f"params__{param_name}", flat=True))
 
+    def get_related_experiments_queryset(self) -> models.QuerySet:
+        return self.experiment_set.filter(is_archived=False)
+
+    def get_static_trigger_experiment_ids(self) -> models.QuerySet:
+        from apps.events.models import EventAction, EventActionType
+
+        return (
+            EventAction.objects.filter(
+                action_type=EventActionType.PIPELINE_START,
+                params__pipeline_id=self.id,
+                static_trigger__is_archived=False,
+            )
+            .annotate(trigger_experiment_id=models.F("static_trigger__experiment"))
+            .values("trigger_experiment_id")
+        )
+
 
 class Node(BaseModel, VersionsMixin):
     flow_id = models.CharField(max_length=128, db_index=True)  # The ID assigned by react-flow
@@ -311,6 +344,20 @@ def create_new_version(self, save: bool = True):
             new_version.save()
         return new_version
 
+    def archive(self):
+        """
+        Archiving a node will also archive the assistant if it is an assistant node. The node's versions will be
+        archived when the pipeline they belong to is archived.
+        """
+        from apps.assistants.models import OpenAiAssistant
+
+        super().archive()
+        if self.is_a_version and self.type == "AssistantNode":
+            assistant_id = self.params.get("assistant_id")
+            if assistant_id:
+                assistant = OpenAiAssistant.objects.get(id=assistant_id)
+                assistant.archive()
+
 
 class PipelineRunStatus(models.TextChoices):
     RUNNING = "running", "Running"

apps/pipelines/tables.py

@@ -16,9 +16,13 @@ class PipelineTable(tables.Table):
     actions = actions.ActionsColumn(
         actions=[
             actions.edit_action(url_name="pipelines:edit"),
-            actions.delete_action(
-                url_name="pipelines:delete",
+            actions.AjaxAction(
+                "pipelines:delete",
+                title="Archive",
+                icon_class="fa-solid fa-box-archive",
+                required_permissions=["pipelines.delete_pipeline"],
                 confirm_message="This will delete the pipeline and any associated logs. Are you sure?",
+                hx_method="delete",
             ),
         ]
     )

apps/pipelines/tests/test_models.py

@@ -4,9 +4,11 @@
 import pytest
 
 from apps.channels.models import ExperimentChannel
+from apps.events.models import EventActionType
 from apps.experiments.models import Experiment, ExperimentSession, Participant
 from apps.pipelines.tests.utils import create_runnable, end_node, llm_response_with_prompt_node, start_node
 from apps.utils.factories.assistants import OpenAiAssistantFactory
+from apps.utils.factories.events import EventActionFactory, ExperimentFactory, StaticTriggerFactory
 from apps.utils.factories.pipelines import NodeFactory, PipelineFactory
 from apps.utils.factories.service_provider_factories import (
     LlmProviderFactory,
@@ -125,3 +127,33 @@ def test_simple_invoke_with_pipeline(self, get_llm_service):
         ] == expected_call_messages
 
         assert ExperimentSession.objects.count() == 0
+
+    @pytest.mark.django_db()
+    def test_archive_pipeline(self):
+        assistant = OpenAiAssistantFactory()
+        pipeline = PipelineFactory()
+        NodeFactory(pipeline=pipeline, type="AssistantNode", params={"assistant_id": assistant.id})
+        start_pipeline__action = EventActionFactory(
+            action_type=EventActionType.PIPELINE_START,
+            params={
+                "pipeline_id": pipeline.id,
+            },
+        )
+        experiment1 = ExperimentFactory()
+        experiment2 = ExperimentFactory()
+        static_trigger = StaticTriggerFactory(experiment=experiment2, action=start_pipeline__action)
+
+        # Experiment and Static trigger still uses it
+        assert pipeline.archive() is False
+
+        experiment1.archive()
+        # Static trigger from experiment2 still uses it
+        assert pipeline.archive() is False
+
+        static_trigger.archive()
+        # Nothing uses it, so archive it
+        assert pipeline.archive() is True
+
+        # Double check that the node didn't archive the assistant
+        assistant.refresh_from_db()
+        assert assistant.is_archived is False

apps/pipelines/views.py

@@ -6,9 +6,10 @@
 from django.contrib import messages
 from django.contrib.auth.decorators import permission_required
 from django.contrib.auth.mixins import PermissionRequiredMixin
-from django.db.models import Count, QuerySet
+from django.db.models import Count, QuerySet, Subquery
 from django.http import Http404, HttpResponse, JsonResponse
 from django.shortcuts import get_object_or_404, redirect, render
+from django.template.loader import render_to_string
 from django.template.response import TemplateResponse
 from django.urls import reverse
 from django.views import View
@@ -18,7 +19,7 @@
 from django_tables2 import SingleTableView
 
 from apps.assistants.models import OpenAiAssistant
-from apps.experiments.models import AgentTools, SourceMaterial
+from apps.experiments.models import AgentTools, Experiment, SourceMaterial
 from apps.pipelines.flow import FlowPipelineData
 from apps.pipelines.models import Pipeline, PipelineRun
 from apps.pipelines.nodes.base import OptionsSource
@@ -28,6 +29,8 @@
 from apps.teams.decorators import login_and_team_required
 from apps.teams.mixins import LoginAndTeamRequiredMixin
 
+from ..generics.chips import Chip
+
 
 class PipelineHome(LoginAndTeamRequiredMixin, TemplateView, PermissionRequiredMixin):
     permission_required = "pipelines.view_pipeline"
@@ -86,9 +89,31 @@ class DeletePipeline(LoginAndTeamRequiredMixin, View, PermissionRequiredMixin):
 
     def delete(self, request, team_slug: str, pk: int):
         pipeline = get_object_or_404(Pipeline.objects.prefetch_related("node_set"), id=pk, team=request.team)
-        pipeline.archive()
-        messages.success(request, f"{pipeline.name} deleted")
-        return HttpResponse()
+        if pipeline.archive():
+            messages.success(request, "Pipeline Archived")
+            return HttpResponse()
+        else:
+            experiments = [
+                Chip(label=experiment.name, url=experiment.get_absolute_url())
+                for experiment in pipeline.get_related_experiments_queryset()
+            ]
+
+            query = pipeline.get_static_trigger_experiment_ids()
+            static_trigger_experiments = [
+                Chip(label=experiment.name, url=experiment.get_absolute_url())
+                for experiment in Experiment.objects.filter(id__in=Subquery(query)).all()
+            ]
+
+            response = render_to_string(
+                "assistants/partials/referenced_objects.html",
+                context={
+                    "object_name": "pipeline",
+                    "experiments": experiments,
+                    "static_trigger_experiments": static_trigger_experiments,
+                },
+            )
+
+        return HttpResponse(response, headers={"HX-Reswap": "none"}, status=400)
 
 
 def _pipeline_node_parameter_values(team, llm_providers, llm_provider_models):

templates/assistants/partials/referenced_objects.html

@@ -1,5 +1,5 @@
 <div>
-  <p>This assistant is still referenced by other objects. You must archive those first.</p>
+  <p>This {{ object_name }} is still referenced by other objects. You must archive those first.</p>
   {% if experiments %}
     <h2 class="font-semibold my-2">Experiments</h2>
     <ul class="list-disc list-inside">
@@ -17,4 +17,13 @@ <h2 class="font-semibold my-2">Pipelines</h2>
       {% endfor %}
     </ul>
   {% endif %}
+
+  {% if static_trigger_experiments %}
+    <h2 class="font-semibold my-2">Static Trigger Experiments</h2>
+    <ul class="list-disc list-inside">
+      {% for experiment in static_trigger_experiments %}
+        <li>{% include "generic/chip.html" with chip=experiment %}</li>
+      {% endfor %}
+    </ul>
+  {% endif %}
 </div>