opensource
/
dify
mirror of https://github.com/langgenius/dify.git
2
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

chore: resolve review issue

This commit is contained in:
fatelei 2025-12-28 11:24:37 +08:00
parent e03b371c43
commit 8818acba31
No known key found for this signature in database
GPG Key ID: 2F91DA05646F4EED
2 changed files with 21 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
api/core/rag/datasource/retrieval_service.py
View File

@ -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

20
api/core/rag/retrieval/dataset_retrieval.py
View File

@ -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