Feature/allow kb edits

backend/problem/serializers.py

@@ -1,4 +1,5 @@
 from rest_framework import serializers
+from django.utils import timezone
 
 from annotation.serializers import KnowledgeBaseAnnotationSerializer
 from annotation.models import (
@@ -77,11 +78,9 @@ class ProblemInputSerializer(serializers.Serializer):
     base = serializers.IntegerField(required=False, allow_null=True)
 
     def validate_id(self, value):
-        """Validate that the Problem ID, if provided, exists and belongs to a user-created problem."""
+        """Validate that the Problem ID, if provided, exists."""
         if value is not None:
-            if not Problem.objects.filter(
-                id=value, dataset=Problem.Dataset.USER
-            ).exists():
+            if not Problem.objects.filter(id=value).exists():
                 raise serializers.ValidationError(
                     f"Problem with ID {value} does not exist."
                 )
@@ -122,8 +121,7 @@ def create(self, validated_data: dict) -> Problem:
         problem.premises.set(premise_sentences)
 
         kb_items = validated_data.get("kbItems", [])
-        if kb_items:
-            self._create_update_kb_annotations(problem, kb_items)
+        self._handle_kb_annotations(problem, kb_items)
 
         return problem
 
@@ -159,21 +157,40 @@ def _create_update_kb_annotation(
         serializer.is_valid(raise_exception=True)
         serializer.save(problem=problem, session=session, created_by=session.user)
 
-    def _create_update_kb_annotations(
+    def _mark_kb_not_in_input_as_removed(
+        self, problem: Problem, kb_items: list[dict], session: AnnotationSession
+    ) -> None:
+        """
+        Marks KnowledgeBase annotations for a problem that are not included in
+        the provided list of kb_items as removed.
+        """
+        kb_item_ids = {kb_item.get("id") for kb_item in kb_items if kb_item.get("id") is not None}
+
+        annotations_to_delete = KnowledgeBaseAnnotation.objects.filter(
+            problem=problem,
+            removed_at__isnull=True
+        ).exclude(id__in=kb_item_ids)
+
+        for annotation in annotations_to_delete:
+            annotation.removed_at = timezone.now()
+            annotation.removed_by = session.user
+            annotation.save()
+
+    def _handle_kb_annotations(
         self, problem: Problem, kb_items: list[dict]
     ) -> None:
         """
-        Creates or update KnowledgeBase and Label annotations for a problem.
+        Creates, updates and deletes KnowledgeBase annotations for a problem.
         Creates an annotation session if it does not exist.
-
-        TODO: handle deletions!
         """
         request = self.context.get("request", None)
-        if not request or not request.user.is_authenticated:
+        if not request or not request.user.is_authenticated or not request.user.can_edit_kb:
             return
 
         session = AnnotationSession.objects.create(user=request.user)
 
+        self._mark_kb_not_in_input_as_removed(problem, kb_items, session)
+
         for kb_item in kb_items:
             self._create_update_kb_annotation(kb_item, problem, session)
 
@@ -185,8 +202,7 @@ def update(self, instance: Problem, validated_data: dict) -> Problem:
 
         # KB annotations can be made for all problems.
         kb_items = validated_data.get("kbItems", [])
-        if kb_items:
-            self._create_update_kb_annotations(instance, kb_items)
+        self._handle_kb_annotations(instance, kb_items)
 
         # Other fields can only be updated for user-created problems.
         if instance.dataset != Problem.Dataset.USER:

backend/problem/serializers_test.py

@@ -1,12 +1,11 @@
 import pytest
 from rest_framework.exceptions import ValidationError
 
-from annotation.models import AnnotationSession, KnowledgeBaseAnnotation
-
-from .serializers import ProblemInputSerializer, ProblemSerializer
+from annotation.models import KnowledgeBaseAnnotation
+from problem.serializers_test_utils import input_serializer_with_user
+from .serializers import ProblemInputSerializer
 from .models import Problem, Sentence
 
-
 @pytest.fixture
 def hypothesis_sentence(db):
     return Sentence.objects.create(text="Hypothesis")
@@ -38,7 +37,6 @@ def non_user_problem(db, hypothesis_sentence, premise_sentence):
     problem.premises.add(premise_sentence)
     return problem
 
-
 @pytest.mark.django_db
 def test_valid_create_data():
     """Test valid data for creating a problem."""
@@ -89,22 +87,6 @@ def test_invalid_id_non_existent():
         serializer.is_valid(raise_exception=True)
     assert "does not exist" in str(exc_info.value)
 
-
-@pytest.mark.django_db
-def test_invalid_id_not_user_problem(non_user_problem):
-    """Test that a non-user problem ID is invalid."""
-    data = {
-        "id": non_user_problem.pk,
-        "premises": ["premise"],
-        "hypothesis": "hypothesis",
-        "kbItems": [],
-    }
-    serializer = ProblemInputSerializer(data=data)
-    with pytest.raises(ValidationError) as exc_info:
-        serializer.is_valid(raise_exception=True)
-    assert "does not exist" in str(exc_info.value)
-
-
 @pytest.mark.django_db
 def test_empty_premises_invalid():
     """Test that an empty list of premises is invalid."""
@@ -131,3 +113,250 @@ def test_blank_hypothesis_invalid():
     assert not serializer.is_valid()
     assert "hypothesis" in serializer.errors
 
+
+@pytest.mark.django_db
+def test_kb_create_no_permission(user_problem, visitor):
+    """Test that KB annotations cannot be created without permission."""
+    kb_input = [
+        {
+            "entity1": "cat",
+            "entity2": "feline",
+            "relationship": "equal",
+            "notes": "Test note",
+        }
+    ]
+
+    serializer = input_serializer_with_user(visitor)
+    serializer._handle_kb_annotations(user_problem, kb_input) # type: ignore
+
+    assert KnowledgeBaseAnnotation.objects.filter(problem=user_problem).count() == 0, "No KB annotations should be created without permission."
+
+
+@pytest.mark.django_db
+def test_kb_update_no_permission(user_problem, kb_annotation, visitor):
+    """Test that KB annotations cannot be updated without permission."""
+    kb_annotation.problem = user_problem
+    kb_annotation.save()
+    original_entity1 = kb_annotation.entity1
+
+    kb_input = [
+        {
+            "id": kb_annotation.pk,
+            "entity1": "updated_cat",
+            "entity2": kb_annotation.entity2,
+            "relationship": kb_annotation.relationship,
+        }
+    ]
+
+    serializer = input_serializer_with_user(visitor)
+    serializer._handle_kb_annotations(user_problem, kb_input) # type: ignore
+
+    # Verify KB annotation was not updated
+    kb_annotation.refresh_from_db()
+    assert kb_annotation.entity1 == original_entity1, "KB annotation should not have been modified without permission."
+
+
+@pytest.mark.django_db
+def test_kb_mark_removed_no_permission(user_problem, kb_annotation, visitor):
+    """Test that KB annotations cannot be marked as removed without permission."""
+    kb_annotation.problem = user_problem
+    kb_annotation.save()
+
+    kb_input = []  # Empty list should mark any existing KB items as removed.
+
+    serializer = input_serializer_with_user(visitor)
+    serializer._handle_kb_annotations(user_problem, kb_input) # type: ignore
+
+    # Verify KB annotation was not removed
+    kb_annotation.refresh_from_db()
+    assert kb_annotation.removed_at is None, "KB annotation should not have been marked as removed without permission."
+
+@pytest.mark.django_db
+def test_create_single_kb_annotation(user_problem, annotator):
+    """Test creating a single KB annotation for a problem."""
+    kb_input = [
+        {
+            "entity1": "dog",
+            "entity2": "canine",
+            "relationship": "equal",
+            "notes": "Test note",
+        }
+    ]
+
+    serializer = input_serializer_with_user(annotator)
+    serializer._handle_kb_annotations(user_problem, kb_input)
+
+    kb_annotations = KnowledgeBaseAnnotation.objects.filter(
+        problem=user_problem, removed_at__isnull=True
+    )
+    assert kb_annotations.count() == 1, "One KB annotation should have been created."
+
+    kb = kb_annotations.first()
+    assert kb is not None
+    assert kb.entity1 == "dog"
+    assert kb.entity2 == "canine"
+    assert kb.relationship == "equal"
+    assert kb.notes == "Test note"
+    assert kb.created_by == annotator
+
+
+@pytest.mark.django_db
+def test_update_single_kb_annotation(user_problem, kb_annotation, annotator):
+    """Test updating a single KB annotation for a problem."""
+    kb_annotation.problem = user_problem
+    kb_annotation.save()
+
+    kb_input = [
+        {
+            "id": kb_annotation.pk,
+            "entity1": "updated_entity1",
+            "entity2": "updated_entity2",
+            "relationship": "subset",
+            "notes": "Updated note",
+        }
+    ]
+
+    serializer = input_serializer_with_user(annotator)
+    serializer._handle_kb_annotations(user_problem, kb_input)
+
+    # Verify KB annotation was updated
+    kb_annotation.refresh_from_db()
+    assert kb_annotation.entity1 == "updated_entity1"
+    assert kb_annotation.entity2 == "updated_entity2"
+    assert kb_annotation.relationship == "subset"
+    assert kb_annotation.notes == "Updated note"
+    assert kb_annotation.removed_at is None
+
+
+@pytest.mark.django_db
+def test_mark_kb_annotation_as_removed(user_problem, kb_annotation, annotator):
+    """Test marking a KB annotation as removed for a problem."""
+    kb_annotation.problem = user_problem
+    kb_annotation.save()
+
+    kb_input = []  # Empty list should mark existing as removed
+
+    serializer = input_serializer_with_user(annotator)
+    serializer._handle_kb_annotations(user_problem, kb_input)
+
+    # Verify KB annotation was marked as removed
+    kb_annotation.refresh_from_db()
+    assert kb_annotation.removed_at is not None
+    assert kb_annotation.removed_by == annotator
+
+
+@pytest.mark.django_db
+def test_create_and_update_multiple_kb_annotations(
+    user_problem, kb_annotation, annotator
+):
+    """Test creating and updating multiple KB annotations for a problem."""
+    kb_annotation.problem = user_problem
+    kb_annotation.save()
+
+    kb_input = [
+        {
+            "id": kb_annotation.pk,
+            "entity1": "updated_e1",
+            "entity2": "updated_e2",
+            "relationship": "not_equal",
+        },
+        {
+            "entity1": "new_e1",
+            "entity2": "new_e2",
+            "relationship": "subset",
+        },
+        {
+            "entity1": "another_e1",
+            "entity2": "another_e2",
+            "relationship": "superset",
+        },
+    ]
+
+    serializer = input_serializer_with_user(annotator)
+    serializer._handle_kb_annotations(user_problem, kb_input)
+
+    kb_annotations = KnowledgeBaseAnnotation.objects.filter(
+        problem=user_problem, removed_at__isnull=True
+    )
+    assert kb_annotations.count() == 3, "There should be three new active KB annotations after update."
+
+    # Verify the updated annotation
+    kb_annotation.refresh_from_db()
+    assert kb_annotation.entity1 == "updated_e1"
+    assert kb_annotation.entity2 == "updated_e2"
+    assert kb_annotation.relationship == "not_equal"
+
+    # Verify the new annotations
+    new_annotations = kb_annotations.exclude(id=kb_annotation.pk)
+    assert new_annotations.count() == 2, "There should be two new KB annotations after update."
+
+    entities = [(kb.entity1, kb.entity2) for kb in new_annotations]
+    assert ("new_e1", "new_e2") in entities
+    assert ("another_e1", "another_e2") in entities
+
+
+@pytest.mark.django_db
+def test_create_update_and_remove_kb_annotations(
+    user_problem, annotator, annotator_session
+):
+    """Test creating, updating and removing multiple KB annotations for a problem."""
+    # Create initial KB annotations
+    kb1 = KnowledgeBaseAnnotation.objects.create(
+        problem=user_problem,
+        entity1="keep_e1",
+        entity2="keep_e2",
+        relationship="equal",
+        session=annotator_session,
+        created_by=annotator,
+    )
+
+    kb2 = KnowledgeBaseAnnotation.objects.create(
+        problem=user_problem,
+        entity1="remove_e1",
+        entity2="remove_e2",
+        relationship="equal",
+        session=annotator_session,
+        created_by=annotator,
+    )
+
+    # request = Mock(user=annotator)
+    kb_input = [
+        {
+            "id": kb1.pk,
+            "entity1": "updated_keep_e1",
+            "entity2": "updated_keep_e2",
+            "relationship": "subset",
+        },
+        {
+            "entity1": "new_e1",
+            "entity2": "new_e2",
+            "relationship": "superset",
+        },
+    ]
+
+    serializer = input_serializer_with_user(annotator)
+    serializer._handle_kb_annotations(user_problem, kb_input)
+
+    # Verify kb1 was updated.
+    kb1.refresh_from_db()
+    assert kb1.entity1 == "updated_keep_e1"
+    assert kb1.entity2 == "updated_keep_e2"
+    assert kb1.relationship == "subset"
+    assert kb1.removed_at is None
+
+    # Verify kb2 was marked as removed.
+    kb2.refresh_from_db()
+    assert kb2.removed_at is not None
+    assert kb2.removed_by == annotator
+
+    # Verify new annotation was created.
+    active_annotations = KnowledgeBaseAnnotation.objects.filter(
+        problem=user_problem, removed_at__isnull=True
+    )
+    assert active_annotations.count() == 2
+
+    new_annotation = active_annotations.exclude(id=kb1.pk).first()
+    assert new_annotation is not None
+    assert new_annotation.entity1 == "new_e1"
+    assert new_annotation.entity2 == "new_e2"
+    assert new_annotation.relationship == "superset"

backend/problem/serializers_test_utils.py

@@ -0,0 +1,10 @@
+from unittest.mock import Mock
+from problem.serializers import ProblemInputSerializer
+
+
+def input_serializer_with_user(user) -> ProblemInputSerializer:
+    """
+    Helper function to create a ProblemInputSerializer with a mock request containing the specified user.
+    """
+    request = Mock(user=user)
+    return ProblemInputSerializer(context={"request": request})