fix: background investigator node support more search engine (#75)

Change-Id: I030a2b9218dfbda2dd2383b7a73266dd7de589c7
This commit is contained in:
Zhao Longjie 2025-05-12 20:15:47 +08:00 committed by GitHub
parent 98e0d7cbb6
commit 9266201fe5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 23 additions and 16 deletions

View File

@ -2,9 +2,10 @@
DEBUG=True DEBUG=True
APP_ENV=development APP_ENV=development
# Search Engine # Search Engine, Supported values: tavily (recommended), duckduckgo, brave_search, arxiv
SEARCH_API=tavily SEARCH_API=tavily
TAVILY_API_KEY=tvly-xxx TAVILY_API_KEY=tvly-xxx
# BRAVE_SEARCH_API_KEY=xxx # Required only if SEARCH_API is brave_search
# JINA_API_KEY=jina_xxx # Optional, default is None # JINA_API_KEY=jina_xxx # Optional, default is None
# Optional, volcengine TTS for generating podcast # Optional, volcengine TTS for generating podcast

View File

@ -28,7 +28,7 @@ from src.prompts.template import apply_prompt_template
from src.utils.json_utils import repair_json_output from src.utils.json_utils import repair_json_output
from .types import State from .types import State
from ..config import SEARCH_MAX_RESULTS from ..config import SEARCH_MAX_RESULTS, SELECTED_SEARCH_ENGINE, SearchEngine
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -45,10 +45,11 @@ def handoff_to_planner(
def background_investigation_node(state: State) -> Command[Literal["planner"]]: def background_investigation_node(state: State) -> Command[Literal["planner"]]:
logger.info("background investigation node is running.") logger.info("background investigation node is running.")
query = state["messages"][-1].content
if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY:
searched_content = LoggedTavilySearch(max_results=SEARCH_MAX_RESULTS).invoke( searched_content = LoggedTavilySearch(max_results=SEARCH_MAX_RESULTS).invoke(
{"query": state["messages"][-1].content} {"query": query}
) )
background_investigation_results = None background_investigation_results = None
if isinstance(searched_content, list): if isinstance(searched_content, list):
@ -57,7 +58,11 @@ def background_investigation_node(state: State) -> Command[Literal["planner"]]:
for elem in searched_content for elem in searched_content
] ]
else: else:
logger.error(f"Tavily search returned malformed response: {searched_content}") logger.error(
f"Tavily search returned malformed response: {searched_content}"
)
else:
background_investigation_results = web_search_tool.invoke(query)
return Command( return Command(
update={ update={
"background_investigation_results": json.dumps( "background_investigation_results": json.dumps(

View File

@ -14,11 +14,10 @@ from src.tools.tavily_search.tavily_search_results_with_images import (
TavilySearchResultsWithImages, TavilySearchResultsWithImages,
) )
from .decorators import create_logged_tool from src.tools.decorators import create_logged_tool
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
LoggedTavilySearch = create_logged_tool(TavilySearchResultsWithImages) LoggedTavilySearch = create_logged_tool(TavilySearchResultsWithImages)
tavily_search_tool = LoggedTavilySearch( tavily_search_tool = LoggedTavilySearch(
name="web_search", name="web_search",
@ -53,5 +52,7 @@ arxiv_search_tool = LoggedArxivSearch(
) )
if __name__ == "__main__": if __name__ == "__main__":
results = tavily_search_tool.invoke("cute panda") results = LoggedDuckDuckGoSearch(
name="web_search", max_results=SEARCH_MAX_RESULTS, output_format="list"
).invoke("cute panda")
print(json.dumps(results, indent=2, ensure_ascii=False)) print(json.dumps(results, indent=2, ensure_ascii=False))