wanwu/deer-flow
1
0
Fork 0
mirror of https://gitee.com/wanwujie/deer-flow synced 2026-04-19 12:24:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: Refactors the retriever function to use async/await (#821)

Browse Source 
* refactor: Refactors the retriever function to use async/await
This commit is contained in:
Xun
2026-01-20 19:56:26 +08:00
committed by GitHub
parent 2ed0eeb107
commit 0e64c52975
10 changed files with 196 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/unit/tools/test_tools_retriever.py
View File

@@ -66,18 +66,20 @@ async def test_retriever_tool_arun():
mock_retriever = Mock(spec=Retriever)
chunk = Chunk(content="async content", similarity=0.8)
doc = Document(id="doc2", chunks=[chunk])
mock_retriever.query_relevant_documents.return_value = [doc]
# Mock the async method
async def mock_async_query(*args, **kwargs):
return [doc]
mock_retriever.query_relevant_documents_async = mock_async_query
resources = [Resource(uri="test://uri", title="Test")]
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
mock_run_manager = Mock(spec=AsyncCallbackManagerForToolRun)
mock_sync_manager = Mock(spec=CallbackManagerForToolRun)
mock_run_manager.get_sync.return_value = mock_sync_manager
result = await tool._arun("async keywords", mock_run_manager)
mock_run_manager.get_sync.assert_called_once()
assert isinstance(result, list)
assert len(result) == 1
assert result[0] == doc.to_dict()