Feature/user problems

backend/problem/models.py

@@ -73,13 +73,16 @@ def serialize(self) -> dict:
             case _:
                 serialized_extra_data = {}
 
+        kb_items = KnowledgeBase.objects.filter(problem=self)
+
         return {
             "id": self.pk,
             "dataset": self.dataset,
             "premises": [premise.text for premise in self.premises.all()],
             "hypothesis": self.hypothesis.text,
             "entailmentLabel": self.entailment_label,
             "extraData": serialized_extra_data,
+            "kbItems": [item.serialize() for item in kb_items],
         }
 
 
@@ -105,3 +108,11 @@ class Relationship(models.TextChoices):
         on_delete=models.CASCADE,
         related_name="knowledge_bases",
     )
+
+    def serialize(self) -> dict:
+        return {
+            "id": self.pk,
+            "entity1": self.entity1,
+            "entity2": self.entity2,
+            "relationship": self.relationship,
+        }

backend/problem/problem_details.py

@@ -10,16 +10,16 @@
 
 @dataclass
 class RelatedProblemIds:
-    first: Optional[str] = None
-    previous: Optional[str] = None
-    next: Optional[str] = None
-    last: Optional[str] = None
-    total: Optional[int] = None
+    first: int | None = None
+    previous: int | None = None
+    next: int | None = None
+    last: int | None = None
+    total: int | None = None
 
 
 def get_related_problem_ids(
     problem_qs: QuerySet[Problem],
-    problem_id: Optional[int],
+    problem_id: int | None = None,
 ) -> RelatedProblemIds:
     """
     Retrieves the IDs of surrounding problem objects
@@ -44,10 +44,10 @@ def get_related_problem_ids(
             problem = None
 
     return RelatedProblemIds(
-        first=str(first_problem.pk) if first_problem else None,
-        previous=str(previous_problem.pk) if previous_problem else None,
-        next=str(next_problem.pk) if next_problem else None,
-        last=str(last_problem.pk) if last_problem else None,
+        first=first_problem.pk if first_problem else None,
+        previous=previous_problem.pk if previous_problem else None,
+        next=next_problem.pk if next_problem else None,
+        last=last_problem.pk if last_problem else None,
         total=total,
     )
 
@@ -70,7 +70,7 @@ def get_filters(query_params: QueryDict) -> Q | None:
         filters &= Q(dataset=dataset)
     if entailment_label:
         filters &= Q(entailment_label=entailment_label)
-    if gold is not None:
+    if gold is not None and gold != "":
         raise NotImplementedError()
     if text:
         filters &= Q(

backend/problem/views/problem.py

@@ -1,30 +1,29 @@
-from dataclasses import dataclass
-from typing import Optional
+from dataclasses import asdict, dataclass
+from typing import TypedDict
 
 from django.http import JsonResponse
-from rest_framework.views import APIView, Request
+from rest_framework.views import APIView
 
+from langpro_annotator.logger import logger
 from problem.problem_details import get_filters, get_related_problem_ids
-from problem.models import Problem
+from problem.models import KnowledgeBase, Problem, Sentence
 
 
 @dataclass
 class ProblemResponse:
-    id: Optional[int] = None
-    index: Optional[int] = None
-    problem: Optional[Problem] = None
-    error: Optional[str] = None
+    problem: Problem | None = None
+    index: int | None = None
+    error: str | None = None
 
-    first: Optional[str] = None
-    previous: Optional[str] = None
-    next: Optional[str] = None
-    last: Optional[str] = None
-    total: Optional[int] = None
+    first: int | None = None
+    previous: int | None = None
+    next: int | None = None
+    last: int | None = None
+    total: int | None = None
 
     def json_response(self, status=200) -> JsonResponse:
         return JsonResponse(
             {
-                "id": self.id,
                 "index": self.index,
                 "problem": self.problem.serialize() if self.problem else None,
                 "error": self.error,
@@ -38,8 +37,21 @@ def json_response(self, status=200) -> JsonResponse:
         )
 
 
+@dataclass
+class SaveProblemResponse:
+    id: int | None = None
+    error: str | None = None
+
+    def json_response(self, status=200) -> JsonResponse:
+        return JsonResponse(asdict(self), status=status)
+
+
 class ProblemView(APIView):
-    def get(self, request: Request, problem_id: int | None = None):
+    def get(self, request, problem_id: int | None = None) -> JsonResponse:
+        """
+        If a Problem ID is provided, retrieves the requested Problem.
+        Otherwise, simply returns the first problem of the QS.
+        """
         filters = get_filters(request.query_params)
 
         qs = Problem.objects.all()
@@ -63,12 +75,202 @@ def get(self, request: Request, problem_id: int | None = None):
         related_problem_ids = get_related_problem_ids(qs, problem_id)
 
         return ProblemResponse(
-            id=problem.pk if problem else None,
-            index=problem_index,
             problem=problem,
+            index=problem_index,
             first=related_problem_ids.first,
             previous=related_problem_ids.previous,
             next=related_problem_ids.next,
             last=related_problem_ids.last,
             total=related_problem_ids.total,
         ).json_response(status=200)
+
+    def post(self, request, problem_id: int | None = None) -> JsonResponse:
+        """
+        If the Problem ID is None, attempts to create a new Problem;
+        else the associated Problem is updated.
+        """
+        parse_input = request.data
+
+        try:
+            validated_input = validate_input(parse_input)
+        except ValueError as e:
+            logger.error(f"Input validation error: {e}")
+            return SaveProblemResponse(
+                id=None,
+                error=str(e),
+            ).json_response(status=400)
+
+        problem: Problem | None = None
+        error: str | None = None
+
+        if problem_id is None:
+            try:
+                problem = create_problem_from_input(validated_input)
+            except Exception as e:
+                error = f"Error creating problem: {str(e)}"
+        else:
+            try:
+                problem = update_problem_from_input(validated_input)
+            except Exception as e:
+                error = f"Error updating problem: {str(e)}"
+
+        if problem is None or error is not None:
+            return SaveProblemResponse(
+                id=None,
+                error=error,
+            ).json_response(status=500)
+
+        return SaveProblemResponse(id=problem.pk, error=None).json_response()
+
+
+class KBItem(TypedDict):
+    id: int
+    entity1: str
+    relationship: str
+    entity2: str
+
+
+class ParseInput(TypedDict):
+    id: int
+    premises: list[str]
+    hypothesis: str
+    kbItems: list[KBItem]
+
+
+def validate_input(parse_input: dict) -> ParseInput:
+    """
+    Validate the parse input data.
+    """
+    if not isinstance(parse_input, dict):
+        raise ValueError("Input must be a dictionary")
+
+    if "id" not in parse_input or (
+        parse_input["id"] is not None and not isinstance(parse_input["id"], int)
+    ):
+        raise ValueError("Invalid problem 'id' field")
+
+    if "premises" not in parse_input or not isinstance(parse_input["premises"], list):
+        raise ValueError("Missing or invalid 'premises' field")
+
+    if "hypothesis" not in parse_input or not isinstance(
+        parse_input["hypothesis"], str
+    ):
+        raise ValueError("Missing or invalid 'hypothesis' field")
+
+    if "kbItems" not in parse_input or not isinstance(parse_input["kbItems"], list):
+        raise ValueError("Missing or invalid 'kbItems' field")
+
+    for item in parse_input["kbItems"]:
+        if not isinstance(item, dict):
+            raise ValueError("Each kbItem must be a dictionary")
+        if "id" not in item or (
+            item["id"] is not None and not isinstance(item["id"], int)
+        ):
+            raise ValueError("Invalid 'id' in kbItem")
+        if "entity1" not in item or not isinstance(item["entity1"], str):
+            raise ValueError("Missing or invalid 'entity1' in kbItem")
+        if "relationship" not in item or not isinstance(item["relationship"], str):
+            raise ValueError("Missing or invalid 'relationship' in kbItem")
+        if item["relationship"] not in KnowledgeBase.Relationship.values:
+            raise ValueError(f"Invalid 'relationship' in kbItem.")
+        if "entity2" not in item or not isinstance(item["entity2"], str):
+            raise ValueError("Missing or invalid 'entity2' in kbItem")
+
+    return ParseInput(
+        id=parse_input["id"],
+        premises=parse_input["premises"],
+        hypothesis=parse_input["hypothesis"],
+        kbItems=parse_input["kbItems"],
+    )
+
+
+def create_problem_from_input(parse_input: ParseInput) -> Problem:
+    """
+    Save a new Problem instance from the given parse input data.
+    """
+    try:
+        premise_sentences = [
+            Sentence.objects.get_or_create(text=premise)[0]
+            for premise in parse_input["premises"]
+        ]
+    except Exception as e:
+        raise ValueError(f"Error creating premise sentences: {e}")
+
+    try:
+        hypothesis_sentence = Sentence.objects.get_or_create(
+            text=parse_input["hypothesis"]
+        )[0]
+    except Exception as e:
+        raise ValueError(f"Error creating hypothesis sentence: {e}")
+
+    problem = Problem.objects.create(
+        hypothesis=hypothesis_sentence,
+        dataset=Problem.Dataset.USER,
+        # TODO: Determine entailment label based on LangPro parser output.
+        entailment_label=Problem.EntailmentLabel.UNKNOWN,
+        extra_data={},
+    )
+
+    problem.premises.set(premise_sentences)
+
+    update_or_create_kb_items(problem=problem, kb_items=parse_input["kbItems"])
+
+    return problem
+
+
+def update_or_create_kb_items(problem: Problem, kb_items: list[KBItem]) -> None:
+    kb_ids: list[str] = []
+    for item in kb_items:
+        id = item.get("id")
+        entity1 = item.get("entity1")
+        relationship = item.get("relationship")
+        entity2 = item.get("entity2")
+
+        if id is None:
+            try:
+                kb = KnowledgeBase.objects.create(
+                    entity1=entity1,
+                    relationship=relationship,
+                    entity2=entity2,
+                    problem=problem,
+                )
+                kb_ids.append(kb.pk)
+            except Exception as e:
+                raise ValueError(f"Error creating knowledge base items: {e}.")
+        else:
+            try:
+                kb = KnowledgeBase.objects.get(id=id, problem_id=problem.pk)
+            except KnowledgeBase.DoesNotExist:
+                raise ValueError(f"Unable to find knowledge base item with id {id}.")
+            kb.entity1 = entity1
+            kb.relationship = relationship
+            kb.entity2 = entity2
+            kb.save()
+            kb_ids.append(kb.pk)
+
+    # Delete existing knowledge bases associated to this problem that are
+    # not included in the input.
+    KnowledgeBase.objects.filter(problem_id=problem.pk).exclude(id__in=kb_ids).delete()
+
+
+def update_problem_from_input(parse_input: ParseInput) -> Problem:
+    try:
+        problem = Problem.objects.get(id=parse_input["id"])
+    except Problem.DoesNotExist:
+        raise ValueError(f"Cannot find Problem with ID: {parse_input["id"]}")
+
+    problem.hypothesis = Sentence.objects.get_or_create(text=parse_input["hypothesis"])[
+        0
+    ]
+    problem.save()
+
+    premises: list[Sentence] = []
+    for input_premise in parse_input["premises"]:
+        premise = Sentence.objects.get_or_create(text=input_premise)[0]
+        premises.append(premise)
+
+    problem.premises.set(premises)
+
+    update_or_create_kb_items(problem, parse_input["kbItems"])
+
+    return problem