Add ID field for complaint types

StanfordVL · Mar 5, 2025 · 7aa020e · 7aa020e
1 parent 9befec9
commit 7aa020e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/bddl/knowledge_base/models.py b/bddl/knowledge_base/models.py
@@ -968,17 +968,21 @@ def __str__(self):
 
 @dataclass(eq=False, order=False)
 class ComplaintType(Model):
+    id: str = UUIDField()
     message: str
 
     complaints_fk: OneToMany = OneToManyField("Complaint", "complaint_type")
 
     class Meta:
-        pk = "message"
+        pk = "id"
         ordering = ["message"]
 
     @cached_property
     def objects(self):
         return [complaint.object for complaint in self.complaints]
+
+    def __str__(self):
+        return self.message
 
 @dataclass(eq=False, order=False)
 class Complaint(Model):

diff --git a/bddl/knowledge_base/processing.py b/bddl/knowledge_base/processing.py
@@ -401,7 +401,11 @@ def create_complaints(self):
                 continue
 
             # Create the relevant objects
-            complaint_type, created = ComplaintType.get_or_create(message=complaint_message)
+            complaint_types_with_message = [ct for ct in ComplaintType.all_objects() if ct.message == complaint_message]
+            if len(complaint_types_with_message) == 0:
+                complaint_type = ComplaintType.create(message=complaint_message)
+            else:
+                complaint_type, = complaint_types_with_message
             Complaint.create(object=obj, complaint_type=complaint_type, content=complaint_content)
 
     # TODO: Move to cached property on Synset