184 lines
7.6 KiB
Python
184 lines
7.6 KiB
Python
# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
import os
|
|
from typing import Annotated, Any, Optional, Union
|
|
|
|
from ....doc_utils import export_module
|
|
from ....import_utils import optional_import_block, require_optional_import
|
|
from ....llm_config import LLMConfig
|
|
from ... import Depends, Tool
|
|
from ...dependency_injection import on
|
|
|
|
with optional_import_block():
|
|
from tavily import TavilyClient
|
|
|
|
|
|
@require_optional_import(
|
|
[
|
|
"tavily",
|
|
],
|
|
"tavily",
|
|
)
|
|
def _execute_tavily_query(
|
|
query: str,
|
|
tavily_api_key: str,
|
|
search_depth: str = "basic",
|
|
topic: str = "general",
|
|
include_answer: str = "basic",
|
|
include_raw_content: bool = False,
|
|
include_domains: list[str] = [],
|
|
num_results: int = 5,
|
|
) -> Any:
|
|
"""
|
|
Execute a search query using the Tavily API.
|
|
|
|
Args:
|
|
query (str): The search query string.
|
|
tavily_api_key (str): The API key for Tavily.
|
|
search_depth (str, optional): The depth of the search ('basic' or 'advanced'). Defaults to "basic".
|
|
topic (str, optional): The topic of the search. Defaults to "general".
|
|
include_answer (str, optional): Whether to include an AI-generated answer ('basic' or 'advanced'). Defaults to "basic".
|
|
include_raw_content (bool, optional): Whether to include raw content in the results. Defaults to False.
|
|
include_domains (list[str], optional): A list of domains to include in the search. Defaults to [].
|
|
num_results (int, optional): The maximum number of results to return. Defaults to 5.
|
|
|
|
Returns:
|
|
Any: The raw response object from the Tavily API client.
|
|
"""
|
|
tavily_client = TavilyClient(api_key=tavily_api_key)
|
|
return tavily_client.search(
|
|
query=query,
|
|
search_depth=search_depth,
|
|
topic=topic,
|
|
include_answer=include_answer,
|
|
include_raw_content=include_raw_content,
|
|
include_domains=include_domains,
|
|
max_results=num_results,
|
|
)
|
|
|
|
|
|
def _tavily_search(
|
|
query: str,
|
|
tavily_api_key: str,
|
|
search_depth: str = "basic",
|
|
topic: str = "general",
|
|
include_answer: str = "basic",
|
|
include_raw_content: bool = False,
|
|
include_domains: list[str] = [],
|
|
num_results: int = 5,
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Perform a Tavily search and format the results.
|
|
|
|
This function takes search parameters, executes the query using `_execute_tavily_query`,
|
|
and formats the results into a list of dictionaries containing title, link, and snippet.
|
|
|
|
Args:
|
|
query (str): The search query string.
|
|
tavily_api_key (str): The API key for Tavily.
|
|
search_depth (str, optional): The depth of the search ('basic' or 'advanced'). Defaults to "basic".
|
|
topic (str, optional): The topic of the search. Defaults to "general".
|
|
include_answer (str, optional): Whether to include an AI-generated answer ('basic' or 'advanced'). Defaults to "basic".
|
|
include_raw_content (bool, optional): Whether to include raw content in the results. Defaults to False.
|
|
include_domains (list[str], optional): A list of domains to include in the search. Defaults to [].
|
|
num_results (int, optional): The maximum number of results to return. Defaults to 5.
|
|
|
|
Returns:
|
|
list[dict[str, Any]]: A list of dictionaries, where each dictionary represents a search result
|
|
with keys 'title', 'link', and 'snippet'. Returns an empty list if no results are found.
|
|
"""
|
|
res = _execute_tavily_query(
|
|
query=query,
|
|
tavily_api_key=tavily_api_key,
|
|
search_depth=search_depth,
|
|
topic=topic,
|
|
include_answer=include_answer,
|
|
include_raw_content=include_raw_content,
|
|
include_domains=include_domains,
|
|
num_results=num_results,
|
|
)
|
|
|
|
return [
|
|
{"title": item.get("title", ""), "link": item.get("url", ""), "snippet": item.get("content", "")}
|
|
for item in res.get("results", [])
|
|
]
|
|
|
|
|
|
@export_module("autogen.tools.experimental")
|
|
class TavilySearchTool(Tool):
|
|
"""
|
|
TavilySearchTool is a tool that uses the Tavily Search API to perform a search.
|
|
|
|
This tool allows agents to leverage the Tavily search engine for information retrieval.
|
|
It requires a Tavily API key, which can be provided during initialization or set as
|
|
an environment variable `TAVILY_API_KEY`.
|
|
|
|
Attributes:
|
|
tavily_api_key (str): The API key used for authenticating with the Tavily API.
|
|
"""
|
|
|
|
def __init__(
|
|
self, *, llm_config: Optional[Union[LLMConfig, dict[str, Any]]] = None, tavily_api_key: Optional[str] = None
|
|
):
|
|
"""
|
|
Initializes the TavilySearchTool.
|
|
|
|
Args:
|
|
llm_config (Optional[Union[LLMConfig, dict[str, Any]]]): LLM configuration. (Currently unused but kept for potential future integration).
|
|
tavily_api_key (Optional[str]): The API key for the Tavily Search API. If not provided,
|
|
it attempts to read from the `TAVILY_API_KEY` environment variable.
|
|
|
|
Raises:
|
|
ValueError: If `tavily_api_key` is not provided either directly or via the environment variable.
|
|
"""
|
|
self.tavily_api_key = tavily_api_key or os.getenv("TAVILY_API_KEY")
|
|
|
|
if self.tavily_api_key is None:
|
|
raise ValueError("tavily_api_key must be provided either as an argument or via TAVILY_API_KEY env var")
|
|
|
|
def tavily_search(
|
|
query: Annotated[str, "The search query."],
|
|
tavily_api_key: Annotated[Optional[str], Depends(on(self.tavily_api_key))],
|
|
search_depth: Annotated[Optional[str], "Either 'advanced' or 'basic'"] = "basic",
|
|
include_answer: Annotated[Optional[str], "Either 'advanced' or 'basic'"] = "basic",
|
|
include_raw_content: Annotated[Optional[bool], "Include the raw contents"] = False,
|
|
include_domains: Annotated[Optional[list[str]], "Specific web domains to search"] = [],
|
|
num_results: Annotated[int, "The number of results to return."] = 5,
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Performs a search using the Tavily API and returns formatted results.
|
|
|
|
Args:
|
|
query: The search query string.
|
|
tavily_api_key: The API key for Tavily (injected dependency).
|
|
search_depth: The depth of the search ('basic' or 'advanced'). Defaults to "basic".
|
|
include_answer: Whether to include an AI-generated answer ('basic' or 'advanced'). Defaults to "basic".
|
|
include_raw_content: Whether to include raw content in the results. Defaults to False.
|
|
include_domains: A list of domains to include in the search. Defaults to [].
|
|
num_results: The maximum number of results to return. Defaults to 5.
|
|
|
|
Returns:
|
|
A list of dictionaries, each containing 'title', 'link', and 'snippet' of a search result.
|
|
|
|
Raises:
|
|
ValueError: If the Tavily API key is not available.
|
|
"""
|
|
if tavily_api_key is None:
|
|
raise ValueError("Tavily API key is missing.")
|
|
return _tavily_search(
|
|
query=query,
|
|
tavily_api_key=tavily_api_key,
|
|
search_depth=search_depth or "basic",
|
|
include_answer=include_answer or "basic",
|
|
include_raw_content=include_raw_content or False,
|
|
include_domains=include_domains or [],
|
|
num_results=num_results,
|
|
)
|
|
|
|
super().__init__(
|
|
name="tavily_search",
|
|
description="Use the Tavily Search API to perform a search.",
|
|
func_or_tool=tavily_search,
|
|
)
|