From 800528f45bb1e63af0ee70e1f86e137e18af9bac Mon Sep 17 00:00:00 2001
From: Maxwell-Lindsey <mal13@illinois.edu>
Date: Sun, 25 Feb 2024 21:50:45 -0600
Subject: [PATCH] added functionality for adding or removing a single doc_group
 for imoproved readability

---
 ai_ta_backend/vector_database.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/ai_ta_backend/vector_database.py b/ai_ta_backend/vector_database.py
index 7c223add..99cba2f8 100644
--- a/ai_ta_backend/vector_database.py
+++ b/ai_ta_backend/vector_database.py
@@ -1737,7 +1737,22 @@ def add_documents_to_doc_group(self, course_name: str, docs: MaterialDocument |
     if ret == '':
       return "Success"
     return ret
+  
+  def append_doc_group(self, course_name: str, doc: MaterialDocument, doc_group: str):
+    """
+    Append a document group to a document's doc_groups array.
+    """
+    if doc_group not in doc.doc_groups:
+      doc.doc_groups.append(doc_group)
+    self.add_documents_to_doc_group(course_name, doc)
 
+  def remove_doc_group(self, course_name: str, doc: MaterialDocument, doc_group: str):
+    """
+    Remove a document group from a document's doc_groups array.
+    """
+    if doc_group in doc.doc_groups:
+      doc.doc_groups.remove(doc_group)
+    self.add_documents_to_doc_group(course_name, doc)
 
 if __name__ == '__main__':
   pass