浏览代码

Feat: `Retrieval` supports internet search. (#5974)

### 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 个月前
父节点
当前提交
dd8779b257
没有帐户链接到提交者的电子邮件
共有 1 个文件被更改,包括 7 次插入0 次删除
  1. 7
    0
      agent/component/retrieval.py

+ 7
- 0
agent/component/retrieval.py 查看文件

@@ -24,6 +24,7 @@ from api.db.services.llm_service import LLMBundle
from api import settings
from agent.component.base import ComponentBase, ComponentParamBase
from rag.app.tag import label_question
from rag.utils.tavily_conn import Tavily


class RetrievalParam(ComponentParamBase):
@@ -40,6 +41,7 @@ class RetrievalParam(ComponentParamBase):
self.kb_ids = []
self.rerank_id = ""
self.empty_response = ""
self.tavily_api_key = ""

def check(self):
self.check_decimal_float(self.similarity_threshold, "[Retrieval] Similarity threshold")
@@ -75,6 +77,11 @@ 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.tavily_api_key:
tav = Tavily(self._param.tavily_api_key)
tav_res = tav.retrieve_chunks(query)
kbinfos["chunks"].extend(tav_res["chunks"])
kbinfos["doc_aggs"].extend(tav_res["doc_aggs"])

if not kbinfos["chunks"]:
df = Retrieval.be_output("")

正在加载...
取消
保存