youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-14 17:07:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: pass session to dataset retrieval threads (#38391)

Browse Source
This commit is contained in:
zyssyz123
2026-07-03 17:30:31 +08:00
committed by GitHub
parent 822b3e86e5
commit e6fa16f652
2 changed files with 210 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1159,6 +1159,39 @@ class DatasetRetrieval:
all_documents.extend(documents)
def _run_retriever_thread(
self,
*,
flask_app: Flask,
dataset_id: str,
query: str | None,
top_k: int,
all_documents: list[Document],
document_ids_filter: list[str] | None,
metadata_condition: MetadataFilteringCondition | None,
attachment_ids: list[str] | None,
cancel_event: threading.Event | None,
thread_exceptions: list[Exception] | None,
) -> None:
try:
with session_factory.create_session() as session:
self._retriever(
flask_app=flask_app,
session=session,
dataset_id=dataset_id,
query=query or "",
top_k=top_k,
all_documents=all_documents,
document_ids_filter=document_ids_filter,
metadata_condition=metadata_condition,
attachment_ids=attachment_ids,
)
except Exception as e:
if cancel_event:
cancel_event.set()
if thread_exceptions is not None:
thread_exceptions.append(e)
def to_dataset_retriever_tool(
self,
session: Session,
@ -1797,7 +1830,7 @@ class DatasetRetrieval:
else:
continue
retrieval_thread = threading.Thread(
target=self._retriever,
target=self._run_retriever_thread,
kwargs={
"flask_app": flask_app,
"dataset_id": dataset.id,
@ -1807,6 +1840,8 @@ class DatasetRetrieval:
"document_ids_filter": document_ids_filter,
"metadata_condition": metadata_condition,
"attachment_ids": [attachment_id] if attachment_id else None,
"cancel_event": cancel_event,
"thread_exceptions": thread_exceptions,
},
)
threads.append(retrieval_thread)

253
api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py
View File

@ -95,6 +95,16 @@ def _metadata_condition() -> AppMetadataFilteringCondition:
return AppMetadataFilteringCondition(logical_operator="and", conditions=[])
@contextmanager
def _patched_retriever_session():
session = MagicMock()
session_ctx = MagicMock()
session_ctx.__enter__.return_value = session
session_ctx.__exit__.return_value = None
with patch("core.rag.retrieval.dataset_retrieval.session_factory.create_session", return_value=session_ctx):
yield session
def create_side_effect_for_search(documents: list[Document]):
"""
Create a side effect function for mocking search methods.
@ -1682,7 +1692,15 @@ class TestRetrievalService:
# Mock _retriever to return documents
def side_effect_retriever(
flask_app, dataset_id, query, top_k, all_documents, document_ids_filter, metadata_condition, attachment_ids
flask_app,
session,
dataset_id,
query,
top_k,
all_documents,
document_ids_filter,
metadata_condition,
attachment_ids,
):
all_documents.extend([doc1, doc2])
@ -1694,23 +1712,24 @@ class TestRetrievalService:
all_documents = []
# Act - Call with dataset_count = 1
dataset_retrieval._multiple_retrieve_thread(
flask_app=mock_flask_app,
available_datasets=[mock_dataset],
metadata_condition=None,
metadata_filter_document_ids=None,
all_documents=all_documents,
tenant_id=tenant_id,
reranking_enable=True,
reranking_mode="reranking_model",
reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"},
weights=None,
top_k=5,
score_threshold=0.5,
query="test query",
attachment_id=None,
dataset_count=1, # Single dataset - should skip second reranking
)
with _patched_retriever_session():
dataset_retrieval._multiple_retrieve_thread(
flask_app=mock_flask_app,
available_datasets=[mock_dataset],
metadata_condition=None,
metadata_filter_document_ids=None,
all_documents=all_documents,
tenant_id=tenant_id,
reranking_enable=True,
reranking_mode="reranking_model",
reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"},
weights=None,
top_k=5,
score_threshold=0.5,
query="test query",
attachment_id=None,
dataset_count=1, # Single dataset - should skip second reranking
)
# Assert
# DataPostProcessor should NOT be called (second reranking skipped)
@ -1757,7 +1776,15 @@ class TestRetrievalService:
# Mock _retriever to return documents
def side_effect_retriever(
flask_app, dataset_id, query, top_k, all_documents, document_ids_filter, metadata_condition, attachment_ids
flask_app,
session,
dataset_id,
query,
top_k,
all_documents,
document_ids_filter,
metadata_condition,
attachment_ids,
):
all_documents.extend([doc1, doc2])
@ -1793,23 +1820,24 @@ class TestRetrievalService:
mock_dataset2.provider = "dify"
# Act - Call with dataset_count = 2
dataset_retrieval._multiple_retrieve_thread(
flask_app=mock_flask_app,
available_datasets=[mock_dataset, mock_dataset2],
metadata_condition=None,
metadata_filter_document_ids=None,
all_documents=all_documents,
tenant_id=tenant_id,
reranking_enable=True,
reranking_mode="reranking_model",
reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"},
weights=None,
top_k=5,
score_threshold=0.5,
query="test query",
attachment_id=None,
dataset_count=2, # Multiple datasets - should perform second reranking
)
with _patched_retriever_session():
dataset_retrieval._multiple_retrieve_thread(
flask_app=mock_flask_app,
available_datasets=[mock_dataset, mock_dataset2],
metadata_condition=None,
metadata_filter_document_ids=None,
all_documents=all_documents,
tenant_id=tenant_id,
reranking_enable=True,
reranking_mode="reranking_model",
reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"},
weights=None,
top_k=5,
score_threshold=0.5,
query="test query",
attachment_id=None,
dataset_count=2, # Multiple datasets - should perform second reranking
)
# Assert
# DataPostProcessor SHOULD be called (second reranking performed)
@ -1867,7 +1895,15 @@ class TestRetrievalService:
# Mock _retriever to return documents
def side_effect_retriever(
flask_app, dataset_id, query, top_k, all_documents, document_ids_filter, metadata_condition, attachment_ids
flask_app,
session,
dataset_id,
query,
top_k,
all_documents,
document_ids_filter,
metadata_condition,
attachment_ids,
):
all_documents.extend([doc1, doc2])
@ -1889,23 +1925,24 @@ class TestRetrievalService:
all_documents = []
# Act - Call with dataset_count = 1
dataset_retrieval._multiple_retrieve_thread(
flask_app=mock_flask_app,
available_datasets=[mock_dataset],
metadata_condition=None,
metadata_filter_document_ids=None,
all_documents=all_documents,
tenant_id=tenant_id,
reranking_enable=True, # Reranking enabled but should be skipped for single dataset
reranking_mode="reranking_model",
reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"},
weights=None,
top_k=5,
score_threshold=0.5,
query="test query",
attachment_id=None,
dataset_count=1,
)
with _patched_retriever_session():
dataset_retrieval._multiple_retrieve_thread(
flask_app=mock_flask_app,
available_datasets=[mock_dataset],
metadata_condition=None,
metadata_filter_document_ids=None,
all_documents=all_documents,
tenant_id=tenant_id,
reranking_enable=True, # Reranking enabled but should be skipped for single dataset
reranking_mode="reranking_model",
reranking_model={"reranking_provider_name": "cohere", "reranking_model_name": "rerank-v2"},
weights=None,
top_k=5,
score_threshold=0.5,
query="test query",
attachment_id=None,
dataset_count=1,
)
# Assert
# DataPostProcessor should NOT be called
@ -3720,7 +3757,15 @@ class TestKnowledgeRetrievalRegression:
)
def fake_retriever(
flask_app, dataset_id, query, top_k, all_documents, document_ids_filter, metadata_condition, attachment_ids
flask_app,
session,
dataset_id,
query,
top_k,
all_documents,
document_ids_filter,
metadata_condition,
attachment_ids,
):
all_documents.append(document)
@ -3744,32 +3789,35 @@ class TestKnowledgeRetrievalRegression:
thread_exceptions: list[Exception] = []
def target():
with patch.object(dataset_retrieval, "_retriever", side_effect=fake_retriever):
with patch(
with (
patch.object(dataset_retrieval, "_retriever", side_effect=fake_retriever),
patch(
"core.rag.retrieval.dataset_retrieval.DataPostProcessor",
ContextRequiredPostProcessor,
):
dataset_retrieval._multiple_retrieve_thread(
flask_app=flask_app,
available_datasets=[mock_dataset, secondary_dataset],
metadata_condition=None,
metadata_filter_document_ids=None,
all_documents=all_documents,
tenant_id=tenant_id,
reranking_enable=True,
reranking_mode="reranking_model",
reranking_model={
"reranking_provider_name": "cohere",
"reranking_model_name": "rerank-v2",
},
weights=None,
top_k=3,
score_threshold=0.0,
query="test query",
attachment_id=None,
dataset_count=2, # force reranking branch
thread_exceptions=thread_exceptions, # ✅ key
)
),
_patched_retriever_session(),
):
dataset_retrieval._multiple_retrieve_thread(
flask_app=flask_app,
available_datasets=[mock_dataset, secondary_dataset],
metadata_condition=None,
metadata_filter_document_ids=None,
all_documents=all_documents,
tenant_id=tenant_id,
reranking_enable=True,
reranking_mode="reranking_model",
reranking_model={
"reranking_provider_name": "cohere",
"reranking_model_name": "rerank-v2",
},
weights=None,
top_k=3,
score_threshold=0.0,
query="test query",
attachment_id=None,
dataset_count=2, # force reranking branch
thread_exceptions=thread_exceptions,
)
t = threading.Thread(target=target)
t.start()
@ -3781,6 +3829,53 @@ class TestKnowledgeRetrievalRegression:
# Current buggy code should record an exception (not raise it)
assert not thread_exceptions, thread_exceptions
def test_run_retriever_thread_provides_session_to_retriever(self):
dataset_retrieval = DatasetRetrieval()
all_documents: list[Document] = []
with _patched_retriever_session() as session:
with patch.object(dataset_retrieval, "_retriever") as mock_retriever:
dataset_retrieval._run_retriever_thread(
flask_app=_FakeFlaskApp(),
dataset_id="dataset-1",
query="test query",
top_k=3,
all_documents=all_documents,
document_ids_filter=None,
metadata_condition=None,
attachment_ids=None,
cancel_event=None,
thread_exceptions=[],
)
mock_retriever.assert_called_once()
assert mock_retriever.call_args.kwargs["session"] is session
def test_run_retriever_thread_records_retriever_exception(self):
dataset_retrieval = DatasetRetrieval()
all_documents: list[Document] = []
cancel_event = threading.Event()
thread_exceptions: list[Exception] = []
expected_error = RuntimeError("retrieval failed")
with _patched_retriever_session():
with patch.object(dataset_retrieval, "_retriever", side_effect=expected_error):
dataset_retrieval._run_retriever_thread(
flask_app=_FakeFlaskApp(),
dataset_id="dataset-1",
query="test query",
top_k=3,
all_documents=all_documents,
document_ids_filter=None,
metadata_condition=None,
attachment_ids=None,
cancel_event=cancel_event,
thread_exceptions=thread_exceptions,
)
assert cancel_event.is_set()
assert thread_exceptions == [expected_error]
class _FakeFlaskApp:
def app_context(self):