From da112e3db03d80e49c912e0805d97a4d17879e9d Mon Sep 17 00:00:00 2001 From: Stephen Hu <812791840@qq.com> Date: Thu, 21 May 2026 14:07:15 +0800 Subject: [PATCH] Refactor:improve dify retrieval logic (#15036) ### What problem does this PR solve? improve dify retrieval logic for o(n) io to o(1) ### Type of change - [x] Refactoring --- api/apps/sdk/dify_retrieval.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/api/apps/sdk/dify_retrieval.py b/api/apps/sdk/dify_retrieval.py index b4520685c..59ea26845 100644 --- a/api/apps/sdk/dify_retrieval.py +++ b/api/apps/sdk/dify_retrieval.py @@ -294,9 +294,15 @@ async def retrieval(tenant_id): if ck["content_with_weight"]: ranks["chunks"].insert(0, ck) + doc_ids = list(set([c["doc_id"] for c in ranks["chunks"]])) + docs = DocumentService.get_by_ids(doc_ids) + doc_map = {doc.id: doc for doc in docs} + records = [] for c in ranks["chunks"]: - e, doc = DocumentService.get_by_id(c["doc_id"]) + doc = doc_map.get(c["doc_id"]) + if not doc: + continue c.pop("vector", None) meta = getattr(doc, 'meta_fields', {}) meta["doc_id"] = c["doc_id"]