CoACT initialize (#292)
This commit is contained in:
Linxin Song
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from .tavily_search import TavilySearchTool
|
||||
|
||||
__all__ = ["TavilySearchTool"]
|
||||
@@ -0,0 +1,183 @@
|
||||
# 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,
|
||||
)
|
||||
Reference in New Issue
Block a user