Merge remote-tracking branch 'origin/fix/add-thread-join-timeout' into fix/add-thread-join-timeout

.github/workflows/build-api-image.yml

@@ -5,6 +5,7 @@ on:
     branches:
       - 'main'
       - 'deploy/dev'
+      - 'fix/add-thread-join-timeout'
   release:
     types: [ published ]
 

api/commands.py

@@ -481,7 +481,7 @@ def normalization_collections():
                 threads.append(document_format_thread)
                 document_format_thread.start()
             for thread in threads:
-                thread.join()
+                thread.join(timeout=30)
 
     click.echo(click.style('Congratulations! restore {} dataset indexes.'.format(len(normalization_count)), fg='green'))
 

api/core/agent/agent_executor.py

@@ -135,6 +135,7 @@ class AgentExecutor:
 
         try:
             output = agent_executor.run(input=query)
+            logging.warning("agent_executor finished!")
         except InvokeError as ex:
             raise ex
         except Exception as ex:

api/core/indexing_runner.py

@@ -662,7 +662,7 @@ class IndexingRunner:
                     threads.append(document_format_thread)
                     document_format_thread.start()
                 for thread in threads:
-                    thread.join()
+                    thread.join(timeout=30)
             return all_qa_documents
         return all_documents
 

api/core/model_runtime/model_providers/spark/llm/llm.py

@@ -134,7 +134,7 @@ class SparkLargeLanguageModel(LargeLanguageModel):
 
             completion += delta
 
-        thread.join()
+        thread.join(timeout=30)
         # transform assistant message to prompt message
         assistant_prompt_message = AssistantPromptMessage(
             content=completion
@@ -194,7 +194,7 @@ class SparkLargeLanguageModel(LargeLanguageModel):
                 )
             )
 
-        thread.join()
+        thread.join(timeout=30)
 
     def _to_credential_kwargs(self, credentials: dict) -> dict:
         """

api/core/rerank/rerank.py

@@ -1,3 +1,4 @@
+import logging
 from typing import List, Optional
 
 from core.model_manager import ModelInstance
@@ -39,7 +40,7 @@ class RerankRunner:
         )
 
         rerank_documents = []
-
+        logging.warning("Rerank finished!")
         for result in rerank_result.docs:
             # format document
             rerank_document = Document(

api/core/tool/dataset_multi_retriever_tool.py

@@ -1,4 +1,5 @@
 import json
+import logging
 import threading
 from typing import List, Optional, Type
 
@@ -70,7 +71,7 @@ class DatasetMultiRetrieverTool(BaseTool):
             threads.append(retrieval_thread)
             retrieval_thread.start()
         for thread in threads:
-            thread.join()
+            thread.join(timeout=30)
         # do rerank for searched documents
         model_manager = ModelManager()
         rerank_model_instance = model_manager.get_model_instance(
@@ -94,6 +95,7 @@ class DatasetMultiRetrieverTool(BaseTool):
         document_context_list = []
         index_node_ids = [document.metadata['doc_id'] for document in all_documents]
         segments = DocumentSegment.query.filter(
+            DocumentSegment.dataset_id.in_(self.dataset_ids),
             DocumentSegment.completed_at.isnot(None),
             DocumentSegment.status == 'completed',
             DocumentSegment.enabled == True,
@@ -243,6 +245,6 @@ class DatasetMultiRetrieverTool(BaseTool):
                         full_text_index_thread.start()
 
                     for thread in threads:
-                        thread.join()
-
+                        thread.join(timeout=30)
                     all_documents.extend(documents)
+                    logging.warning("DatasetMultiRetrieverTool finished!")

api/core/tool/dataset_retriever_tool.py

@@ -140,7 +140,7 @@ class DatasetRetrieverTool(BaseTool):
                     full_text_index_thread.start()
 
                 for thread in threads:
-                    thread.join()
+                    thread.join(timeout=30)
 
                 # hybrid search: rerank after all documents have been searched
                 if retrieval_model['search_method'] == 'hybrid_search':

api/services/hit_testing_service.py

@@ -93,7 +93,7 @@ class HitTestingService:
             full_text_index_thread.start()
 
         for thread in threads:
-            thread.join()
+            thread.join(timeout=30)
 
         if retrieval_model['search_method'] == 'hybrid_search':
             model_manager = ModelManager()