From d8b63f13da302a6e2d2cb2cf6183aa572d26b558 Mon Sep 17 00:00:00 2001 From: Chris Smit Date: Tue, 17 Dec 2024 10:11:31 +0200 Subject: [PATCH 1/5] Pipeline archive updates Pipelines should only be archived when not being referenced by other records. Versions should also be archived when archiving the working version of the pipeline or node --- apps/experiments/models.py | 13 ++++++-- apps/experiments/tests/test_models.py | 14 +++++++++ apps/pipelines/models.py | 45 +++++++++++++++++++++++++-- 3 files changed, 68 insertions(+), 4 deletions(-) diff --git a/apps/experiments/models.py b/apps/experiments/models.py index f053b784d..7e7a283eb 100644 --- a/apps/experiments/models.py +++ b/apps/experiments/models.py @@ -769,13 +769,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 diff --git a/apps/experiments/tests/test_models.py b/apps/experiments/tests/test_models.py index 34e2901f2..040635bce 100644 --- a/apps/experiments/tests/test_models.py +++ b/apps/experiments/tests/test_models.py @@ -1012,6 +1012,20 @@ def test_archive_working_experiment(self, experiment): assert second_version.is_archived is True assert ScheduledMessage.objects.filter(experiment=experiment).exists() is False + def test_archive_with_assistant(self, experiment): + """ + Archiving an assistant experiment should only archive the assistant when it is not still being referenced by + another experiment or pipeline + """ + pytest.fail("Not implemented") + + def test_archive_with_pipeline(self, experiment): + """ + Archiving a pipeline experiment should only archive the pipeline when it is not still being referenced by + another experiment or static trigger action + """ + pytest.fail("Not implemented") + 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 diff --git a/apps/pipelines/models.py b/apps/pipelines/models.py index dae7f3f3c..bf84d7b04 100644 --- a/apps/pipelines/models.py +++ b/apps/pipelines/models.py @@ -262,13 +262,40 @@ 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 self.get_related_actions_queryset().exists(): + 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_related_actions_queryset(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 + ) + class Node(BaseModel, VersionsMixin): flow_id = models.CharField(max_length=128, db_index=True) # The ID assigned by react-flow @@ -311,6 +338,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" From 717651bdedacdd14ea50daa5490b62c30ceeaf12 Mon Sep 17 00:00:00 2001 From: Chris Smit Date: Tue, 17 Dec 2024 10:17:54 +0200 Subject: [PATCH 2/5] Tiny test refactor --- apps/experiments/tests/test_models.py | 40 ++++++++------------------- 1 file changed, 11 insertions(+), 29 deletions(-) diff --git a/apps/experiments/tests/test_models.py b/apps/experiments/tests/test_models.py index 040635bce..f310b9caf 100644 --- a/apps/experiments/tests/test_models.py +++ b/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,34 +1022,6 @@ def test_archive_working_experiment(self, experiment): assert second_version.is_archived is True assert ScheduledMessage.objects.filter(experiment=experiment).exists() is False - def test_archive_with_assistant(self, experiment): - """ - Archiving an assistant experiment should only archive the assistant when it is not still being referenced by - another experiment or pipeline - """ - pytest.fail("Not implemented") - - def test_archive_with_pipeline(self, experiment): - """ - Archiving a pipeline experiment should only archive the pipeline when it is not still being referenced by - another experiment or static trigger action - """ - pytest.fail("Not implemented") - - 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 - - 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, - } - @pytest.mark.django_db() class TestExperimentObjectManager: From ed54697ba035c40f69a797737b1f9d6706cb7235 Mon Sep 17 00:00:00 2001 From: Chris Smit Date: Tue, 17 Dec 2024 12:02:23 +0200 Subject: [PATCH 3/5] Add tests --- apps/experiments/tests/test_models.py | 56 +++++++++++++++++++++++++++ apps/pipelines/tests/test_models.py | 32 +++++++++++++++ 2 files changed, 88 insertions(+) diff --git a/apps/experiments/tests/test_models.py b/apps/experiments/tests/test_models.py index f310b9caf..5910a2e94 100644 --- a/apps/experiments/tests/test_models.py +++ b/apps/experiments/tests/test_models.py @@ -1022,6 +1022,62 @@ def test_archive_working_experiment(self, experiment): assert second_version.is_archived is True assert ScheduledMessage.objects.filter(experiment=experiment).exists() is False + @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 + """ + + 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() class TestExperimentObjectManager: diff --git a/apps/pipelines/tests/test_models.py b/apps/pipelines/tests/test_models.py index 1901ea8e0..b67aee5ba 100644 --- a/apps/pipelines/tests/test_models.py +++ b/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 From 62df7784c49c2114d394c81b643ae9ea873e02f8 Mon Sep 17 00:00:00 2001 From: Chris Smit Date: Tue, 17 Dec 2024 14:53:55 +0200 Subject: [PATCH 4/5] Show referenced objects when trying to archive a pipeline --- apps/assistants/tables.py | 2 +- apps/assistants/views.py | 1 + apps/pipelines/models.py | 14 +++++--- apps/pipelines/tables.py | 8 +++-- apps/pipelines/views.py | 35 ++++++++++++++++--- .../partials/referenced_objects.html | 11 +++++- 6 files changed, 58 insertions(+), 13 deletions(-) diff --git a/apps/assistants/tables.py b/apps/assistants/tables.py index 32ce0a4ea..49423c83d 100644 --- a/apps/assistants/tables.py +++ b/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( diff --git a/apps/assistants/views.py b/apps/assistants/views.py index 82d5a0d03..7cb1d0674 100644 --- a/apps/assistants/views.py +++ b/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, }, diff --git a/apps/pipelines/models.py b/apps/pipelines/models.py index bf84d7b04..48c1932a1 100644 --- a/apps/pipelines/models.py +++ b/apps/pipelines/models.py @@ -270,7 +270,7 @@ def archive(self) -> bool: if self.get_related_experiments_queryset().exists(): return False - if self.get_related_actions_queryset().exists(): + if self.get_static_trigger_experiments(): return False super().archive() @@ -289,11 +289,17 @@ def get_node_param_values(self, node_cls, param_name: str) -> list: def get_related_experiments_queryset(self) -> models.QuerySet: return self.experiment_set.filter(is_archived=False) - def get_related_actions_queryset(self) -> models.QuerySet: + def get_static_trigger_experiments(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 + return ( + EventAction.objects.filter( + action_type=EventActionType.PIPELINE_START, + params__pipeline_id=self.id, + static_trigger__is_archived=False, + ) + .annotate(experiment=models.F("static_trigger__experiment")) + .values("experiment") ) diff --git a/apps/pipelines/tables.py b/apps/pipelines/tables.py index 7dea0ec32..443ccfee7 100644 --- a/apps/pipelines/tables.py +++ b/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", ), ] ) diff --git a/apps/pipelines/views.py b/apps/pipelines/views.py index a3cd17257..6a7ddb762 100644 --- a/apps/pipelines/views.py +++ b/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_experiments() + 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): diff --git a/templates/assistants/partials/referenced_objects.html b/templates/assistants/partials/referenced_objects.html index 390a23f40..cbfffd117 100644 --- a/templates/assistants/partials/referenced_objects.html +++ b/templates/assistants/partials/referenced_objects.html @@ -1,5 +1,5 @@
-

This assistant is still referenced by other objects. You must archive those first.

+

This {{ object_name }} is still referenced by other objects. You must archive those first.

{% if experiments %}

Experiments

    @@ -17,4 +17,13 @@

    Pipelines

    {% endfor %}
{% endif %} + + {% if static_trigger_experiments %} +

Static Trigger Experiments

+
    + {% for experiment in static_trigger_experiments %} +
  • {% include "generic/chip.html" with chip=experiment %}
    • + {% endfor %} +
+ {% endif %}
From 50dc064d187468e59aaffbf6228189d3d8461879 Mon Sep 17 00:00:00 2001 From: Chris Smit Date: Tue, 17 Dec 2024 15:03:11 +0200 Subject: [PATCH 5/5] Naming update --- apps/pipelines/models.py | 8 ++++---- apps/pipelines/views.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/apps/pipelines/models.py b/apps/pipelines/models.py index 48c1932a1..a5d018e26 100644 --- a/apps/pipelines/models.py +++ b/apps/pipelines/models.py @@ -270,7 +270,7 @@ def archive(self) -> bool: if self.get_related_experiments_queryset().exists(): return False - if self.get_static_trigger_experiments(): + if len(self.get_static_trigger_experiment_ids()) > 0: return False super().archive() @@ -289,7 +289,7 @@ def get_node_param_values(self, node_cls, param_name: str) -> list: def get_related_experiments_queryset(self) -> models.QuerySet: return self.experiment_set.filter(is_archived=False) - def get_static_trigger_experiments(self) -> models.QuerySet: + def get_static_trigger_experiment_ids(self) -> models.QuerySet: from apps.events.models import EventAction, EventActionType return ( @@ -298,8 +298,8 @@ def get_static_trigger_experiments(self) -> models.QuerySet: params__pipeline_id=self.id, static_trigger__is_archived=False, ) - .annotate(experiment=models.F("static_trigger__experiment")) - .values("experiment") + .annotate(trigger_experiment_id=models.F("static_trigger__experiment")) + .values("trigger_experiment_id") ) diff --git a/apps/pipelines/views.py b/apps/pipelines/views.py index 6a7ddb762..d3386bb7c 100644 --- a/apps/pipelines/views.py +++ b/apps/pipelines/views.py @@ -98,7 +98,7 @@ def delete(self, request, team_slug: str, pk: int): for experiment in pipeline.get_related_experiments_queryset() ] - query = pipeline.get_static_trigger_experiments() + 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()