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