AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-07-30 08:02:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

Feat: add use KG to retrieval component. (#5988)

Browse Source 
### What problem does this PR solve?

#5973

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Kevin Hu 2025-03-12 19:10:07 +08:00 committed by GitHub
parent 80389ae61e
commit d660f6b9a5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
agent/component/retrieval.py
View File

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