feat: add Tavily as configurable search engine in core search framework#891
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new Tavily search engine integration into the existing search tool framework. This includes defining new data structures for Tavily search requests and results, implementing the Tavily search logic, and integrating it across various parts of the ms_agent search tool infrastructure, such as enums, mappings, and configuration. A new tavily-python dependency is also added. The review comments suggest several improvements: replacing assert with more robust error handling for API key validation, refining type hints for better type safety, and switching from print statements to the logger module for consistent warning and informational messages.
| def __init__(self, api_key: str = None): | ||
|
|
||
| api_key = api_key or os.getenv('TAVILY_API_KEY') | ||
| assert api_key, 'TAVILY_API_KEY must be set either as an argument or as an environment variable' |
There was a problem hiding this comment.
Using assert for validating API keys is not recommended in production code, as assert statements can be optimized out by the Python interpreter, leading to unexpected behavior if the key is missing. A ValueError or RuntimeError provides a more robust and explicit error handling mechanism.
| assert api_key, 'TAVILY_API_KEY must be set either as an argument or as an environment variable' | |
| if not api_key: | |
| raise ValueError('TAVILY_API_KEY must be set either as an argument or as an environment variable') |
| arguments: Dict[str, Any] = field(default_factory=dict) | ||
|
|
||
| # The response from the Tavily search API (dict with 'results' key) | ||
| response: SearchResponse = None |
There was a problem hiding this comment.
The type hint for response should specify the generic type of SearchResponse. Since the TavilySearch.search method populates this with SearchResponse[BaseResult], the type hint here should reflect that for better type safety and clarity.
| response: SearchResponse = None | |
| response: Optional[SearchResponse[BaseResult]] = None |
| Convert the search results to a list of dictionaries. | ||
| """ | ||
| if not self.response or not self.response.results: | ||
| print('***Warning: No search results found.') |
There was a problem hiding this comment.
Using print for warnings is not ideal in a library context. It's better to use the logger module for consistent logging practices, which allows for configurable output and severity levels.
| print('***Warning: No search results found.') | |
| logger.warning('No search results found.') |
| print('***Warning: No search results found.') | ||
| return [] | ||
|
|
||
| if not self.query: |
There was a problem hiding this comment.
Using print for warnings is not ideal in a library context. It's better to use the logger module for consistent logging practices, which allows for configurable output and severity levels.
| if not self.query: | |
| logger.warning('No query provided for search results.') |
| """ | ||
| with open(file_path, 'r', encoding='utf-8') as f: | ||
| data = json.load(f) | ||
| print(f'Search results loaded from {file_path}') |
There was a problem hiding this comment.
Using print for informational messages is not ideal in a library context. It's better to use the logger module for consistent logging practices, which allows for configurable output and severity levels.
1 participant
Summary
BaseResultschema for seamless downstream consumptionFiles Changed
ms_agent/tools/search/search_base.py— AddedTAVILYtoSearchEngineTypeenum and'tavily': 'tavily_search'toENGINE_TOOL_NAMESms_agent/tools/search/tavily/__init__.py— New package init exportingTavilySearch,TavilySearchRequest,TavilySearchResultms_agent/tools/search/tavily/schema.py— New file withTavilySearchRequestandTavilySearchResultdataclassesms_agent/tools/search/tavily/search.py— New file withTavilySearchclass usingtavily-pythonSDKms_agent/tools/search/websearch_tool.py— Added'tavily'toSUPPORTED_ENGINES,get_search_engine_class(),get_search_engine(),_api_keysinit, andconnect()engine instantiationms_agent/tools/search_engine.py— AddedTAVILYbranch inget_web_search_tool()factory andTAVILY_API_KEYsupport in env overridesrequirements/research.txt— Addedtavily-pythondependencyDependency Changes
tavily-pythontorequirements/research.txtEnvironment Variable Changes
TAVILY_API_KEY— required when using the Tavily engineNotes for Reviewers
engine: tavilyin agent YAML config or via theFIN_RESEARCH_SEARCH_ENGINEoverrideTavilySearch.search()method maps Tavily API results (content→summary,url→id) to the sharedBaseResultschema🤖 Generated with Claude Code
Automated Review
get_search_engine_class,get_search_engine,get_web_search_tool,WebSearchTool.connect) are correctly updated. TheTavilySearchclass, schema dataclasses,__init__.py, enum entry, tool-name mapping, andSUPPORTED_ENGINEStuple are all consistent with how Exa and SerpAPI are implemented. Thetavily-pythondependency is added torequirements/research.txt. The Tavily SDK is used correctly (TavilyClient.search(**kwargs)returning a dict with aresultslist). No critical or major issues found; only three minor documentation/style gaps.