feat: add Tavily as configurable search engine in core search framework#891
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces Tavily as a new, configurable search engine within the existing search framework. The change enhances the system's ability to perform AI-optimized web searches by providing an additional, high-relevance search option. The implementation adheres to current architectural patterns, ensuring that the new engine integrates smoothly and is easily configurable without impacting the functionality of other search providers. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![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.
| print(f'Search results loaded from {file_path}') | |
| logger.info(f'Search results loaded from {file_path}') |
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.