fix: Delete from storage when files deleted from Admin site (#970)

Co-authored-by: Martyna Marcinkowska <61530975+tectonia@users.noreply.github.com>

Author: komalg1

code/backend/batch/batch_push_results.py

@@ -60,4 +60,4 @@ def _process_document_deleted_event(message_body) -> None:
     search_handler = Search.get_search_handler(env_helper)
 
     blob_url = message_body.get("data", {}).get("url", "")
-    search_handler.delete_by_source(f"{blob_url}_SAS_TOKEN_PLACEHOLDER_")
+    search_handler.delete_from_index(blob_url)

code/backend/batch/utilities/helpers/azure_blob_storage_client.py

@@ -156,7 +156,23 @@ def delete_file(self, file_name):
         blob_client = self.blob_service_client.get_blob_client(
             container=self.container_name, blob=file_name
         )
-        blob_client.delete_blob()
+        if blob_client.exists():
+            blob_client.delete_blob()
+
+    def delete_files(self, files, integrated_vectorization: bool):
+        """
+        Deletes files from the Azure Blob Storage container.
+
+        Args:
+            files (list[str]): The names of the files to delete.
+
+        Returns:
+            None
+        """
+        for filename, ids in files.items():
+            if not integrated_vectorization:
+                filename = filename.split("/")[-1]
+            self.delete_file(filename)
 
     def get_all_files(self):
         # Get all files in the container from Azure Blob Storage

code/backend/batch/utilities/search/azure_search_handler.py

@@ -1,4 +1,5 @@
 from typing import List
+
 from .search_handler_base import SearchHandlerBase
 from ..helpers.llm_helper import LLMHelper
 from ..helpers.azure_computer_vision_client import AzureComputerVisionClient
@@ -63,6 +64,14 @@ def delete_files(self, files):
 
         return ", ".join(files_to_delete)
 
+    def search_by_blob_url(self, blob_url):
+        return self.search_client.search(
+            "*",
+            select="id, title",
+            include_total_count=True,
+            filter=f"source eq '{blob_url}_SAS_TOKEN_PLACEHOLDER_'",
+        )
+
     def query_search(self, question) -> List[SourceDocument]:
         encoding = tiktoken.get_encoding(self._ENCODER_NAME)
         tokenised_question = encoding.encode(question)

code/backend/batch/utilities/search/integrated_vectorization_search_handler.py

@@ -57,6 +57,16 @@ def output_results(self, results):
 
         return files
 
+    def search_by_blob_url(self, blob_url: str):
+        if self._check_index_exists():
+            title = blob_url.split(f"{self.env_helper.AZURE_BLOB_CONTAINER_NAME}/")[1]
+            return self.search_client.search(
+                "*",
+                select="id, chunk_id, title",
+                include_total_count=True,
+                filter=f"title eq '{title}'",
+            )
+
     def delete_files(self, files):
         ids_to_delete = []
         files_to_delete = []

code/backend/batch/utilities/search/search_handler_base.py

@@ -22,6 +22,11 @@ def get_unique_files(self, results, facet_key: str):
             return [facet["value"] for facet in results.get_facets()[facet_key]]
         return []
 
+    def delete_from_index(self, blob_url) -> None:
+        documents = self.search_by_blob_url(blob_url)
+        files_to_delete = self.output_results(documents)
+        self.delete_files(files_to_delete)
+
     @abstractmethod
     def create_search_client(self) -> SearchClient:
         pass
@@ -50,24 +55,6 @@ def delete_files(self, files):
     def query_search(self, question) -> list[SourceDocument]:
         pass
 
-    def delete_by_source(self, source) -> None:
-        if source is None:
-            return
-
-        documents = self._get_documents_by_source(source)
-        if documents is None:
-            return
-
-        files_to_delete = self.output_results(documents)
-        self.delete_files(files_to_delete)
-
-    def _get_documents_by_source(self, source):
-        if source is None:
-            return None
-
-        return self.search_client.search(
-            "*",
-            select="id, title",
-            include_total_count=True,
-            filter=f"source eq '{source}'",
-        )
+    @abstractmethod
+    def search_by_blob_url(self, blob_url):
+        pass

code/backend/pages/03_Delete_Data.py

@@ -5,6 +5,7 @@
 import logging
 from batch.utilities.helpers.env_helper import EnvHelper
 from batch.utilities.search.search import Search
+from batch.utilities.helpers.azure_blob_storage_client import AzureBlobStorageClient
 
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
 env_helper: EnvHelper = EnvHelper()
@@ -61,11 +62,15 @@
                 st.info("No files selected")
                 st.stop()
             else:
-                files_to_delete = search_handler.delete_files(
-                    selected_files,
+                blob_client = AzureBlobStorageClient()
+                blob_client.delete_files(
+                    selected_files, env_helper.AZURE_SEARCH_USE_INTEGRATED_VECTORIZATION
                 )
-                if len(files_to_delete) > 0:
-                    st.success("Deleted files: " + str(files_to_delete))
+                if len(selected_files) > 0:
+                    st.success(
+                        "Deleted files from storage. Deleting from the index is an asynchronous process and may take a few minutes to complete."
+                        + ", ".join([name for name, ids in selected_files.items()])
+                    )
 
 except Exception:
     logger.error(traceback.format_exc())

code/tests/search_utilities/test_azure_search_handler.py

@@ -363,3 +363,26 @@ def test_semantic_search_with_advanced_image_processing(
         query_answer="extractive",
         top=handler.env_helper.AZURE_SEARCH_TOP_K,
     )
+
+
+def test_delete_from_index(handler, mock_search_client):
+    # given
+    blob_url = "https://example.com/blob"
+    filter_value = f"source eq '{blob_url}_SAS_TOKEN_PLACEHOLDER_'"
+    documents = [
+        {"id": 1, "title": "file1"},
+        {"id": 2, "title": "file2"},
+        {"id": 3, "title": "file1"},
+        {"id": 4, "title": "file3"},
+    ]
+    handler.search_client.search.return_value = documents
+    ids_to_delete = [{"id": 1}, {"id": 3}, {"id": 2}, {"id": 4}]
+
+    # when
+    handler.delete_from_index(blob_url)
+
+    # then
+    handler.search_client.search.assert_called_once_with(
+        "*", select="id, title", include_total_count=True, filter=filter_value
+    )
+    handler.search_client.delete_documents.assert_called_once_with(ids_to_delete)

code/tests/search_utilities/test_integrated_vectorization_search_handler.py

@@ -254,3 +254,28 @@ def test_query_search_converts_results_to_source_documents(handler):
 
     # then
     assert actual_results == expected_results
+
+
+def test_delete_from_index(env_helper_mock, handler, search_client_mock):
+    # given
+    env_helper_mock.AZURE_BLOB_CONTAINER_NAME = "documents"
+    blob_url = "https://example.com/documents/file1.txt"
+    title = "file1.txt"
+    documents = [
+        {"chunk_id": "123_chunk", "title": title},
+        {"chunk_id": "789_chunk", "title": title},
+    ]
+    search_client_mock.search.return_value = documents
+    ids_to_delete = [{"chunk_id": "123_chunk"}, {"chunk_id": "789_chunk"}]
+
+    # when
+    handler.delete_from_index(blob_url)
+
+    # then
+    search_client_mock.search.assert_called_once_with(
+        "*",
+        select="id, chunk_id, title",
+        include_total_count=True,
+        filter=f"title eq '{title}'",
+    )
+    search_client_mock.delete_documents.assert_called_once_with(ids_to_delete)

code/tests/test_batch_push_results.py

@@ -123,6 +123,6 @@ def test_batch_push_results_with_blob_deleted_event_uses_search_to_delete_with_s
     )
 
     batch_push_results.build().get_user_function()(mock_queue_message)
-    mock_get_search_handler.delete_by_source.assert_called_once_with(
-        "https://test.test/test/test_filename.pdf_SAS_TOKEN_PLACEHOLDER_"
+    mock_get_search_handler.delete_from_index.assert_called_once_with(
+        "https://test.test/test/test_filename.pdf"
     )