Skip to content

Commit

Permalink
feat: Tavily search integration (#1039)
Browse files Browse the repository at this point in the history
Co-authored-by: Wendong <w3ndong.fan@gmail.com>
  • Loading branch information
2 people authored and MuggleJinx committed Oct 30, 2024
1 parent c765383 commit 9c1dc20
Show file tree
Hide file tree
Showing 3 changed files with 361 additions and 286 deletions.
52 changes: 51 additions & 1 deletion camel/toolkits/search_toolkit.py
Original file line number Diff line number Diff line change
Expand Up @@ -305,7 +305,56 @@ def query_wolfram_alpha(self, query: str, is_detailed: bool) -> str:
for sub in pod.subpods:
result += (sub.plaintext or "None") + '\n'

return result.rstrip() # Remove trailing whitespace
return result.rstrip()

def tavily_search(
self, query: str, num_results: int = 5, **kwargs
) -> List[Dict[str, Any]]:
r"""Use Tavily Search API to search information for the given query.
Args:
query (str): The query to be searched.
num_results (int): The number of search results to retrieve
(default is `5`).
**kwargs: Additional optional parameters supported by Tavily's API:
- search_depth (str): "basic" or "advanced" search depth.
- topic (str): The search category, e.g., "general" or "news."
- days (int): Time frame in days for news-related searches.
- max_results (int): Max number of results to return
(overrides `num_results`).
See https://docs.tavily.com/docs/python-sdk/tavily-search/
api-reference for details.
Returns:
List[Dict[str, Any]]: A list of dictionaries representing search
results. Each dictionary contains:
- 'result_id' (int): The result's index.
- 'title' (str): The title of the result.
- 'description' (str): A brief description of the result.
- 'long_description' (str): Detailed information, if available.
- 'url' (str): The URL of the result.
- 'content' (str): Relevant content from the search result.
- 'images' (list): A list of related images (if
`include_images` is True).
- 'published_date' (str): Publication date for news topics
(if available).
"""
from tavily import TavilyClient # type: ignore[import-untyped]

Tavily_API_KEY = os.getenv("TAVILY_API_KEY")
if not Tavily_API_KEY:
raise ValueError(
"`TAVILY_API_KEY` not found in environment variables. "
"Get `TAVILY_API_KEY` here: `https://www.tavily.com/api/`."
)

client = TavilyClient(Tavily_API_KEY)

try:
results = client.search(query, max_results=num_results, **kwargs)
return results
except Exception as e:
return [{"error": f"An unexpected error occurred: {e!s}"}]

def get_tools(self) -> List[FunctionTool]:
r"""Returns a list of FunctionTool objects representing the
Expand All @@ -320,6 +369,7 @@ def get_tools(self) -> List[FunctionTool]:
FunctionTool(self.search_google),
FunctionTool(self.search_duckduckgo),
FunctionTool(self.query_wolfram_alpha),
FunctionTool(self.tavily_search),
]


Expand Down
Loading

0 comments on commit 9c1dc20

Please sign in to comment.