diff --git a/api/core/rag/datasource/retrieval_service.py b/api/core/rag/datasource/retrieval_service.py index 0c1158a658a..a4ac63004c1 100644 --- a/api/core/rag/datasource/retrieval_service.py +++ b/api/core/rag/datasource/retrieval_service.py @@ -1,4 +1,5 @@ import concurrent.futures +import logging from concurrent.futures import ThreadPoolExecutor from typing import Any @@ -36,6 +37,8 @@ default_retrieval_model = { "score_threshold_enabled": False, } +logger = logging.getLogger(__name__) + class RetrievalService: # Cache precompiled regular expressions to avoid repeated compilation @@ -108,7 +111,7 @@ class RetrievalService: if futures: for future in concurrent.futures.as_completed(futures, timeout=3600): - if future.exception(): + if exceptions: for f in futures: f.cancel() break @@ -215,6 +218,7 @@ class RetrievalService: ) all_documents.extend(documents) except Exception as e: + logger.error(e, exc_info=True) exceptions.append(str(e)) @classmethod @@ -308,6 +312,7 @@ class RetrievalService: else: all_documents.extend(documents) except Exception as e: + logger.error(e, exc_info=True) exceptions.append(str(e)) @classmethod @@ -356,6 +361,7 @@ class RetrievalService: else: all_documents.extend(documents) except Exception as e: + logger.error(e, exc_info=True) exceptions.append(str(e)) @staticmethod diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py index 1880a1b8b70..b32e3e8829d 100644 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ b/api/core/rag/retrieval/dataset_retrieval.py @@ -569,10 +569,14 @@ class DatasetRetrieval: all_threads.append(attachment_thread) attachment_thread.start() - for thread in all_threads: - thread.join(timeout=300) + # Poll threads with short timeout to detect errors quickly (fail-fast) + while any(t.is_alive() for t in all_threads): + for thread in all_threads: + thread.join(timeout=0.1) + if thread_exceptions: + cancel_event.set() + break if thread_exceptions: - cancel_event.set() break if thread_exceptions: @@ -1454,9 +1458,13 @@ class DatasetRetrieval: ) threads.append(retrieval_thread) retrieval_thread.start() - for thread in threads: - thread.join(timeout=300) - # Check for cancellation signal between threads + + # Poll threads with short timeout to respond quickly to cancellation + while any(t.is_alive() for t in threads): + for thread in threads: + thread.join(timeout=0.1) + if cancel_event and cancel_event.is_set(): + break if cancel_event and cancel_event.is_set(): break