Преглед изворни кода

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)
tags/v0.17.2
Kevin Hu пре 7 месеци
родитељ
комит
d660f6b9a5
No account linked to committer's email address
1 измењених фајлова са 10 додато и 0 уклоњено
  1. 10
    0
      agent/component/retrieval.py

+ 10
- 0
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)

Loading…
Откажи
Сачувај