From dd8779b257f2aeded6bb160055317b4d21c00115 Mon Sep 17 00:00:00 2001 From: Kevin Hu Date: Wed, 12 Mar 2025 16:51:01 +0800 Subject: [PATCH] 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) --- agent/component/retrieval.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/agent/component/retrieval.py b/agent/component/retrieval.py index 6abfb784c..b96bc8bf5 100644 --- a/agent/component/retrieval.py +++ b/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("")