From d660f6b9a5db2eb8d3c62cae3e9ddd6b5d2e6ce2 Mon Sep 17 00:00:00 2001 From: Kevin Hu Date: Wed, 12 Mar 2025 19:10:07 +0800 Subject: [PATCH] Feat: add use KG to retrieval component. (#5988) ### What problem does this PR solve? #5973 ### Type of change - [x] New Feature (non-breaking change which adds functionality) --- agent/component/retrieval.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/agent/component/retrieval.py b/agent/component/retrieval.py index b96bc8bf5..c960725e6 100644 --- a/agent/component/retrieval.py +++ b/agent/component/retrieval.py @@ -42,6 +42,7 @@ class RetrievalParam(ComponentParamBase): self.rerank_id = "" self.empty_response = "" self.tavily_api_key = "" + self.use_kg = False def check(self): self.check_decimal_float(self.similarity_threshold, "[Retrieval] Similarity threshold") @@ -77,6 +78,15 @@ class Retrieval(ComponentBase, ABC): self._param.similarity_threshold, 1 - self._param.keywords_similarity_weight, aggs=False, rerank_mdl=rerank_mdl, rank_feature=label_question(query, kbs)) + if self._param.use_kg: + ck = settings.kg_retrievaler.retrieval(query, + [kbs[0].tenant_id], + self._param.kb_ids, + embd_mdl, + LLMBundle(kbs[0].tenant_id, LLMType.CHAT)) + if ck["content_with_weight"]: + kbinfos["chunks"].insert(0, ck) + if self._param.tavily_api_key: tav = Tavily(self._param.tavily_api_key) tav_res = tav.retrieve_chunks(query)