Restructered indexing API with non atomic requests

backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py

@@ -302,7 +302,7 @@ def index_document(
         """
         tool: CustomTool = CustomTool.objects.get(pk=tool_id)
         if is_summary:
-            profile_manager = ProfileManager.objects.get(
+            profile_manager: ProfileManager = ProfileManager.objects.get(
                 prompt_studio_tool=tool, is_summarize_llm=True
             )
             default_profile = profile_manager
@@ -345,6 +345,12 @@ def index_document(
             reindex=True,
             run_id=run_id,
         )
+        # import time
+
+        # print("Start")
+        # time.sleep(30)  # Pauses the program for 2 seconds
+        # print("End")
+        # doc_id = "123"
 
         logger.info(f"[{tool_id}] Indexing successful for doc: {file_name}")
         PromptStudioHelper._publish_log(

backend/prompt_studio/prompt_studio_core/urls.py

@@ -1,4 +1,6 @@
+from django.db import transaction
 from django.urls import path
+from django.utils.decorators import method_decorator
 from rest_framework.urlpatterns import format_suffix_patterns
 
 from .views import PromptStudioCoreView
@@ -77,7 +79,9 @@
         ),
         path(
             "prompt-studio/index-document/<uuid:pk>",
-            prompt_studio_prompt_index,
+            method_decorator(transaction.non_atomic_requests)(
+                prompt_studio_prompt_index
+            ),
             name="prompt-studio-prompt-index",
         ),
         path(

backend/prompt_studio/prompt_studio_core/views.py

@@ -60,6 +60,7 @@
 logger = logging.getLogger(__name__)
 
 
+# @method_decorator(transaction.non_atomic_requests(using='default'), name='dispatch')
 class PromptStudioCoreView(viewsets.ModelViewSet):
     """Viewset to handle all Custom tool related operations."""
 
@@ -189,7 +190,7 @@ def make_profile_default(self, request: HttpRequest, pk: Any = None) -> Response
             data={"default_profile": profile_manager.profile_id},
         )
 
-    @action(detail=True, methods=["get"])
+    @action(detail=True, methods=["post"])
     def index_document(self, request: HttpRequest, pk: Any = None) -> Response:
         """API Entry point method to index input file.
 

backend/prompt_studio/prompt_studio_index_manager/prompt_studio_index_helper.py

@@ -1,6 +1,7 @@
 import json
 import logging
 
+from django.db import transaction
 from prompt_studio.prompt_profile_manager.models import ProfileManager
 from prompt_studio.prompt_studio_core.exceptions import IndexingAPIError
 from prompt_studio.prompt_studio_document_manager.models import DocumentManager
@@ -18,46 +19,51 @@ def handle_index_manager(
         profile_manager: ProfileManager,
         doc_id: str,
     ) -> IndexManager:
-        document: DocumentManager = DocumentManager.objects.get(pk=document_id)
+        try:
 
-        index_id = "raw_index_id"
-        if is_summary:
-            index_id = "summarize_index_id"
+            with transaction.atomic():
 
-        args: dict[str, str] = dict()
-        args["document_manager"] = document
-        args["profile_manager"] = profile_manager
+                document: DocumentManager = DocumentManager.objects.get(pk=document_id)
 
-        try:
-            # Create or get the existing record for this document and
-            # profile combo
-            index_manager, success = IndexManager.objects.get_or_create(**args)
-
-            if success:
-                logger.info(
-                    f"Index manager doc_id: {doc_id} for "
-                    f"profile {profile_manager.profile_id} created"
-                )
-            else:
-                logger.info(
-                    f"Index manager doc_id: {doc_id} for "
-                    f"profile {profile_manager.profile_id} updated"
-                )
-
-            index_ids = index_manager.index_ids_history
-            index_ids_list = json.loads(index_ids) if index_ids else []
-            if doc_id not in index_ids:
-                index_ids_list.append(doc_id)
-
-            args[index_id] = doc_id
-            args["index_ids_history"] = json.dumps(index_ids_list)
-
-            # Update the record with the index id
-            result: IndexManager = IndexManager.objects.filter(
-                index_manager_id=index_manager.index_manager_id
-            ).update(**args)
+                index_id = "raw_index_id"
+                if is_summary:
+                    index_id = "summarize_index_id"
+
+                args: dict[str, str] = dict()
+                args["document_manager"] = document
+                args["profile_manager"] = profile_manager
+
+                # Create or get the existing record for this document and
+                # profile combo
+                index_manager, success = IndexManager.objects.get_or_create(**args)
+
+                if success:
+                    logger.info(
+                        f"Index manager doc_id: {doc_id} for "
+                        f"profile {profile_manager.profile_id} created"
+                    )
+                else:
+                    logger.info(
+                        f"Index manager doc_id: {doc_id} for "
+                        f"profile {profile_manager.profile_id} updated"
+                    )
 
+                index_ids = index_manager.index_ids_history
+                index_ids_list = json.loads(index_ids) if index_ids else []
+                if doc_id not in index_ids:
+                    index_ids_list.append(doc_id)
+
+                args[index_id] = doc_id
+                args["index_ids_history"] = json.dumps(index_ids_list)
+
+                # Update the record with the index id
+                result: IndexManager = IndexManager.objects.filter(
+                    index_manager_id=index_manager.index_manager_id
+                ).update(**args)
+
+                transaction.commit()
+                return result
         except Exception as e:
+            transaction.rollback()
+            transaction.set_autocommit(True)
             raise IndexingAPIError("Error updating indexing status") from e
-
-        return result