Repository: BrainBlend-AI/atomic-agents

Issue:
Search tool with tavily
Search tool with tavily

Hints from issue comments:
@KennyVaneetvelde: https://github.com/BrainBlend-AI/atomic-agents/blob/main/atomic-forge/guides/tool_structure.md
☝️  in case anyone wants to give this a go, this is a nice issue for a newbie to pick up - otherwise I'll try to get around to it as soon as I can@lucasjinreal: Hi, thanks for the catch up. I just recently found Tavily is useful for poor users like me, they have free 1000/month, sometimes enough for me to build some agents. Due to not very famillar with atomic-agents, willing to see an official tool for this.

Solution patch:
diff --git a/atomic-forge/tools/tavily_search/.coveragerc b/atomic-forge/tools/tavily_search/.coveragerc
new file mode 100644
index 00000000..c45136ef
--- /dev/null
+++ b/atomic-forge/tools/tavily_search/.coveragerc
@@ -0,0 +1,8 @@
+[run]
+source = tool
+omit = */tests/*
+
+[report]
+exclude_lines =
+    if __name__ == "__main__":
+show_missing = True
diff --git a/atomic-forge/tools/tavily_search/README.md b/atomic-forge/tools/tavily_search/README.md
new file mode 100644
index 00000000..1ae9cf3c
--- /dev/null
+++ b/atomic-forge/tools/tavily_search/README.md
@@ -0,0 +1,87 @@
+# Tavily Search Tool
+
+## Overview
+The Tavily Search Tool is a powerful utility within the Atomic Agents ecosystem that allows you to perform searches using Tavily, a search engine built for AI Agents. This tool enables you to fetch search results from Tavily.
+
+## Prerequisites and Dependencies
+- Python 3.9 or later
+- atomic-agents (See [here](/README.md) for installation instructions)
+- pydantic
+- requests
+
+## Installation
+You can install the tool using any of the following options:
+
+1. Using the CLI tool that comes with Atomic Agents. Simply run `atomic` and select the tool from the list of available tools. After doing so you will be asked for a target directory to download the tool into.
+2. Good old fashioned copy/paste: Just like any other tool inside the Atomic Forge, you can copy the code from this repo directly into your own project, provided you already have atomic-agents installed according to the instructions in the main [README](/README.md).
+
+## Configuration
+
+### Parameters
+
+- `api_url` (str): The api key of the Tavily user.
+- `max_results` (int, optional): The maximum number of search results to return. Defaults to `10`.
+
+### Example
+
+```python
+config = TavilySearchToolConfig(
+    api_key="my-api-key",
+    max_results=5
+)
+```
+
+## Input & Output Structure
+
+### Input Schema
+- `queries` (List[str]): List of search queries.
+
+### Output Schema
+- `results` (List[TavilySearchResultItemSchema]): List of search result items.
+
+Each `TavilySearchResultItemSchema` contains:
+- `title` (str): The title of the search result.
+- `url` (str): The URL of the search result.
+- `content` (Optional[str]): The content snippet of the search result.
+- `score` (float): The score of the search result.
+- `raw_content` (Optional[str]): The raw content of the search result.
+- `query` (str): The query used to obtain this search result.
+- `answer` (Optional[str]): The answer to the query provided by Tavily.
+
+## Usage
+
+Here's an example of how to use the Tavily Search Tool:
+
+
+```python
+import os
+from tool.tavily_search import TavilyTool, TavilySearchToolConfig
+
+# Initialize the tool with your Tavily instance URL
+config = TavilySearchToolConfig(api_key=os.getenv("TAVILY_API_KEY"), max_results=5)
+search_tool = TavilyTool(config=config)
+
+# Define input data
+input_data = TavilyTool.input_schema(
+    queries=["Python programming", "Machine learning"],
+)
+
+# Perform the search
+result = search_tool.run(input_data)
+print(result)
+```
+
+## Contributing
+
+Contributions are welcome! To contribute:
+
+1. Fork the repository.
+2. Create a new feature branch.
+3. Commit your changes with clear messages.
+4. Open a pull request detailing your changes.
+
+Please ensure you follow the project's coding standards and include tests for any new features or bug fixes.
+
+## License
+
+This project is licensed under the same license as the main Atomic Agents project. See the [LICENSE](/LICENSE) file in the repository root for more details.
diff --git a/atomic-forge/tools/tavily_search/pyproject.toml b/atomic-forge/tools/tavily_search/pyproject.toml
new file mode 100644
index 00000000..84a0559b
--- /dev/null
+++ b/atomic-forge/tools/tavily_search/pyproject.toml
@@ -0,0 +1,25 @@
+[tool.poetry]
+name = "tavily-search"
+version = "1.0"
+description = "A tool for performing searches using Tavily"
+authors = ["Kyle Webb <webbjkyle@gmail.com>"]
+readme = "README.md"
+package-mode = false
+
+[tool.poetry.dependencies]
+python = ">=3.12,<4"
+atomic-agents = {path = "../../../", develop = true}
+pydantic = ">=2.8.2,<3"
+aiohttp = ">=3.9.0,<4"
+
+[tool.poetry.group.dev.dependencies]
+coverage = ">=7.6.1,<8.0.0"
+pytest = ">=8.3.3,<9.0.0"
+pytest-asyncio = ">=0.23.5,<1.0.0"
+pytest-cov = ">=5.0.0,<6.0.0"
+python-dotenv = ">=1.0.0,<2.0.0"
+rich = ">=13.7.0,<14.0.0"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
diff --git a/atomic-forge/tools/tavily_search/requirements.txt b/atomic-forge/tools/tavily_search/requirements.txt
new file mode 100644
index 00000000..e7586f3d
--- /dev/null
+++ b/atomic-forge/tools/tavily_search/requirements.txt
@@ -0,0 +1,4 @@
+atomic-agents>=1.0.0,<2.0.0
+pydantic>=2.8.2,<3.0.0
+sympy>=1.12,<2.0.0
+aiohttp>=3.9.0,<4.0.0
\ No newline at end of file
diff --git a/atomic-forge/tools/tavily_search/tool/tavily_search.py b/atomic-forge/tools/tavily_search/tool/tavily_search.py
new file mode 100644
index 00000000..ddd22196
--- /dev/null
+++ b/atomic-forge/tools/tavily_search/tool/tavily_search.py
@@ -0,0 +1,210 @@
+import os
+from typing import List, Optional
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+
+import aiohttp
+from pydantic import Field
+
+from atomic_agents.agents.base_agent import BaseIOSchema
+from atomic_agents.lib.base.base_tool import BaseTool, BaseToolConfig
+
+
+################
+# INPUT SCHEMA #
+################
+class TavilySearchToolInputSchema(BaseIOSchema):
+    """
+    Schema for input to a tool for searching for information, news, references, and other content using Tavily.
+    Returns a list of search results with a short description or content snippet and URLs for further exploration
+    """
+
+    queries: List[str] = Field(..., description="List of search queries.")
+
+
+####################
+# OUTPUT SCHEMA(S) #
+####################
+class TavilySearchResultItemSchema(BaseIOSchema):
+    """This schema represents a single search result item"""
+
+    title: str = Field(..., description="The title of the search result")
+    url: str = Field(..., description="The URL of the search result")
+    content: str = Field(None, description="The content snippet of the search result")
+    score: float = Field(..., description="The score of the search result")
+    raw_content: Optional[str] = Field(None, description="The raw content of the search result")
+    query: Optional[str] = Field(..., description="The query used to obtain this search result")
+    answer: Optional[str] = Field(..., description="The answer to the query provided by Tavily")
+
+
+class TavilySearchToolOutputSchema(BaseIOSchema):
+    """This schema represents the output of the Tavily search tool."""
+
+    results: List[TavilySearchResultItemSchema] = Field(..., description="List of search result items")
+
+
+##############
+# TOOL LOGIC #
+##############
+class TavilySearchToolConfig(BaseToolConfig):
+    api_key: str = ""
+    max_results: Optional[int] = 10
+    search_depth: Optional[str] = "basic"
+    topic: Optional[str] = "general"
+    include_query: Optional[bool] = False
+    include_domains: Optional[List[str]] = []
+    include_answer: Optional[bool] = False
+    include_raw_content: Optional[bool] = False
+
+
+class TavilySearchTool(BaseTool):
+    """
+    Tool for performing searches on Tavily based on the provided queries and topic.
+
+    Attributes:
+        input_schema (TavilySearchToolInputSchema): The schema for the input data.
+        output_schema (TavilySearchToolOutputSchema): The schema for the output data.
+        max_results (int): The maximum number of search results to return.
+        search_depth (string): The depth of the search to perform. (advanced or basic)
+        topic (string): The category that the result is classified under. (general or news)
+        include_domains (List[str]): A list of domains to pull results from.
+        include_answer (bool): Include the answer in the respones from Tavily.
+        include_raw_content (bool): Include the raw content of the search results.
+    """
+
+    input_schema = TavilySearchToolInputSchema
+    output_schema = TavilySearchToolOutputSchema
+
+    def __init__(self, config: TavilySearchToolConfig = TavilySearchToolConfig()):
+        """
+        Initializes the TavilyTool.
+
+        Args:
+            config (TavilySearchToolConfig):
+                Configuration for the tool, including base URL, max results, and optional title and description overrides.
+        """
+        super().__init__(config)
+        self.api_key = config.api_key
+        self.max_results = config.max_results
+        self.search_depth = config.search_depth
+        self.topic = config.topic
+        self.include_query = config.include_query
+        self.include_domains = config.include_domains
+        self.include_answer = config.include_answer
+        self.include_raw_content = config.include_raw_content
+
+    async def _fetch_search_results(self, session: aiohttp.ClientSession, query: str) -> List[dict]:
+        headers = {
+            "accept": "/",
+            "content-type": "application/json",
+            "origin": "https://app.tavily.com/",
+            "referer": "https://app.tavily.com/",
+            "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
+        }
+
+        json_data = {
+            "query": query,
+            "api_key": self.api_key,
+            "search_depth": self.search_depth,
+            "topic": self.topic,
+            "include_query": self.include_query,
+            "include_answer": self.include_answer,
+            "include_raw_content": self.include_raw_content,
+            "include_domains": self.include_domains,
+            "max_results": self.max_results,
+        }
+
+        async with session.post("https://api.tavily.com/search", headers=headers, json=json_data) as response:
+            if response.status != 200:
+                error_message = await response.text()
+                raise Exception(
+                    f"Failed to fetch search results for query '{query}': {response.status} {response.reason}. Details: {error_message}"
+                )
+            data = await response.json()
+            results = data.get("results", [])
+            answer = data.get("answer", "")
+
+            # Add query information to each result
+            for result in results:
+                result["query"] = query
+                if self.include_answer:
+                    result["answer"] = answer
+
+            return results
+
+    async def run_async(
+        self, params: TavilySearchToolInputSchema, max_results: Optional[int] = None
+    ) -> TavilySearchToolOutputSchema:
+        async with aiohttp.ClientSession() as session:
+            # Fetch results for all queries
+            tasks = [self._fetch_search_results(session, query) for query in params.queries]
+            raw_results = await asyncio.gather(*tasks)
+
+        # Process results for each query
+        processed_results = []
+        for query_results in raw_results:
+            query_processed = []
+            for result in query_results:
+                if all(key in result for key in ["title", "url", "content", "score"]):
+                    query_processed.append(
+                        TavilySearchResultItemSchema(
+                            title=result["title"],
+                            url=result["url"],
+                            content=result.get("content", ""),
+                            score=result.get("score", 0),
+                            raw_content=result.get("raw_content"),
+                            query=result.get("query") if self.include_query else None,
+                            answer=result.get("answer") if self.include_answer else None,
+                        )
+                    )
+                else:
+                    print(f"Skipping result due to missing keys: {result}")
+
+            # Limit results per query
+            query_processed = query_processed[: max_results or self.max_results]
+            processed_results.extend(query_processed)
+
+        return TavilySearchToolOutputSchema(results=processed_results)
+
+    def run(self, params: TavilySearchToolInputSchema, max_results: Optional[int] = None) -> TavilySearchToolOutputSchema:
+        """
+        Runs the TavilyTool synchronously with the given parameters.
+
+        This method creates an event loop in a separate thread to run the asynchronous operations.
+
+        Args:
+            params (TavilySearchToolInputSchema): The input parameters for the tool, adhering to the input schema.
+            max_results (Optional[int]): The maximum number of search results to return.
+
+        Returns:
+            TavilySearchToolOutputSchema: The output of the tool, adhering to the output schema.
+
+        Raises:
+            ValueError: If the base URL is not provided.
+            Exception: If the request to Tavily fails.
+        """
+        with ThreadPoolExecutor() as executor:
+            return executor.submit(
+                asyncio.run,
+                self.run_async(
+                    params,
+                    max_results,
+                ),
+            ).result()
+
+
+#################
+# EXAMPLE USAGE #
+#################
+if __name__ == "__main__":
+    from rich.console import Console
+
+    rich_console = Console()
+
+    search_tool_instance = TavilySearchTool(config=TavilySearchToolConfig(api_key=os.getenv("TAVILY_API_KEY"), max_results=2))
+
+    search_input = TavilySearchTool.input_schema(queries=["Python programming", "Machine learning", "Artificial intelligence"])
+
+    output = search_tool_instance.run(search_input)
+
+    rich_console.print(output)


Test patch:
diff --git a/atomic-forge/tools/tavily_search/tests/test_tavily_seach.py b/atomic-forge/tools/tavily_search/tests/test_tavily_seach.py
new file mode 100644
index 00000000..89287f00
--- /dev/null
+++ b/atomic-forge/tools/tavily_search/tests/test_tavily_seach.py
@@ -0,0 +1,230 @@
+import os
+import sys
+import pytest
+from unittest.mock import patch, AsyncMock, MagicMock
+from types import SimpleNamespace
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from tool.tavily_search import (  # noqa: E402
+    TavilySearchTool,
+    TavilySearchToolInputSchema,
+    TavilySearchToolConfig,
+)
+
+
+@pytest.fixture
+def mock_aiohttp_session():
+    # Patch the ClientSession in the module where it's used
+    with patch("tool.tavily_search.aiohttp.ClientSession") as mock_session_class:
+        # Create a mock session
+        mock_session = MagicMock()
+
+        # Mock `post` method with a response
+        mock_response = SimpleNamespace(
+            status=200,
+            json=AsyncMock(
+                return_value={
+                    "results": [
+                        {"title": "Mock Title", "url": "https://example.com", "content": "Mock Content", "score": 1.0},
+                    ]
+                }
+            ),
+        )
+        mock_session.post.return_value.__aenter__.return_value = mock_response
+
+        # Mock the session returned by `ClientSession`
+        mock_session_class.return_value.__aenter__.return_value = mock_session
+        yield mock_session
+
+
+@pytest.mark.asyncio
+async def test_tavily_search_tool_missing_fields(mock_aiohttp_session):
+    mock_api_key = "KEY"
+    mock_query = "query with missing fields"
+    mock_response_data = {
+        "results": [
+            {"title": "Result Missing Content", "url": "https://example.com/1", "score": 0.8},
+            {"content": "Result Missing Title", "url": "https://example.com/2", "score": 0.8},
+            {"title": "Result Missing URL", "content": "Some content", "score": 0.8},
+            {"title": "Result Missing Score", "url": "https://example.com/4", "content": "Some content"},
+            {"title": "Result Missing Query", "url": "https://example.com/4", "content": "Some content", "score": 0.7},
+            {"title": "Valid Result", "url": "https://example.com/2", "content": "Valid content", "score": 1.0},
+        ]
+    }
+
+    # Set up the mocked response for post
+    mock_aiohttp_session.post.return_value.__aenter__.return_value.json.return_value = mock_response_data
+
+    # Initialize the tool
+    tavily_tool = TavilySearchTool(TavilySearchToolConfig(api_key=mock_api_key))
+    input_schema = TavilySearchToolInputSchema(queries=[mock_query])
+
+    # Run the tool
+    result = await tavily_tool.run_async(input_schema)
+
+    # The last two results should be included
+    assert len(result.results) == 2
+    assert result.results[0].title == "Result Missing Query"
+    assert result.results[1].title == "Valid Result"
+
+
+def test_tavily_search_tool_sync_run_method(mock_aiohttp_session):
+    mock_api_key = "KEY"
+    mock_query = "sync query"
+    mock_response_data = {
+        "results": [
+            {
+                "title": "Sync Test Result",
+                "url": "https://example.com/sync",
+                "content": "Sync content",
+                "query": mock_query,
+                "score": 1.0,
+            }
+        ]
+    }
+
+    # Create a mock response object
+    mock_response = AsyncMock()
+    mock_response.status = 200
+    mock_response.json.return_value = mock_response_data
+
+    # Configure the mock
+    mock_aiohttp_session.post.return_value.__aenter__.return_value = mock_response
+
+    # Initialize the tool
+    tavily_tool = TavilySearchTool(TavilySearchToolConfig(api_key=mock_api_key))
+    input_schema = TavilySearchToolInputSchema(queries=[mock_query])
+
+    # Run the tool synchronously
+    result = tavily_tool.run(input_schema)
+
+    # Assertions
+    assert len(result.results) == 1
+    assert result.results[0].title == "Sync Test Result"
+
+
+@pytest.mark.asyncio
+async def test_tavily_search_tool_with_max_results(mock_aiohttp_session):
+    mock_api_key = "KEY"
+    mock_query = "test query with max results"
+    mock_response_data = {
+        "results": [
+            {
+                "title": f"Result {i}",
+                "url": f"https://example.com/{i}",
+                "content": f"Content {i}",
+                "query": mock_query,
+                "score": 0.5,
+            }
+            for i in range(10)
+        ]
+    }
+
+    # Create a mock response object
+    mock_response = AsyncMock()
+    mock_response.status = 200
+    mock_response.json.return_value = mock_response_data
+
+    # Configure the mock
+    mock_aiohttp_session.post.return_value.__aenter__.return_value = mock_response
+
+    # Initialize the tool
+    tavily_tool = TavilySearchTool(TavilySearchToolConfig(api_key=mock_api_key, max_results=10))
+    input_schema = TavilySearchToolInputSchema(queries=[mock_query])
+
+    # Run the tool with max_results=5
+    result = await tavily_tool.run_async(input_schema, max_results=5)
+
+    # Assertions
+    assert len(result.results) == 5
+
+
+@pytest.mark.asyncio
+async def test_tavily_search_tool_no_results(mock_aiohttp_session):
+    # Existing test case remains the same
+    mock_api_key = "KEY"
+    mock_query = "no results query"
+    mock_response_data = {"results": []}
+
+    # Create a mock response object
+    mock_response = AsyncMock()
+    mock_response.status = 200
+    mock_response.json.return_value = mock_response_data
+
+    # Configure the mock
+    mock_aiohttp_session.post.return_value.__aenter__.return_value = mock_response
+
+    # Initialize the tool
+    tavily_tool = TavilySearchTool(TavilySearchToolConfig(api_key=mock_api_key))
+    input_schema = TavilySearchToolInputSchema(queries=[mock_query])
+
+    # Run the tool
+    result = await tavily_tool.run_async(input_schema)
+
+    # Assertions
+    assert len(result.results) == 0
+
+
+@pytest.mark.asyncio
+async def test_tavily_search_tool_error(mock_aiohttp_session):
+    # Existing test case remains the same
+    mock_api_key = "KEY"
+    mock_query = "error query"
+
+    # Create a mock response object with a 500 status
+    mock_response = AsyncMock()
+    mock_response.status = 500
+    mock_response.reason = "Internal Server Error"
+
+    # Configure the mock
+    mock_aiohttp_session.post.return_value.__aenter__.return_value = mock_response
+
+    # Initialize the tool
+    tavily_tool = TavilySearchTool(TavilySearchToolConfig(api_key=mock_api_key))
+    input_schema = TavilySearchToolInputSchema(queries=[mock_query])
+
+    # Run the tool and expect an exception
+    with pytest.raises(Exception) as excinfo:
+        await tavily_tool.run_async(input_schema)
+
+    # Assertion on the exception message
+    assert "Failed to fetch search results" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_tavily_search_tool_include_answer(mock_aiohttp_session):
+    mock_api_key = "KEY"
+    mock_query = "test query with answer"
+    mock_answer = "This is a test answer."
+    mock_response_data = {
+        "results": [
+            {"title": "Result 1", "url": "https://example.com/1", "content": "Content 1", "score": 1.0},
+            {"title": "Result 2", "url": "https://example.com/2", "content": "Content 2", "score": 0.8},
+        ],
+        "answer": mock_answer,
+    }
+
+    # Set up the mocked response for post
+    mock_aiohttp_session.post.return_value.__aenter__.return_value.json.return_value = mock_response_data
+
+    # Initialize the tool with `include_answer=True`
+    tavily_tool = TavilySearchTool(TavilySearchToolConfig(api_key=mock_api_key))
+    tavily_tool.include_answer = True  # Enable the include_answer feature
+    input_schema = TavilySearchToolInputSchema(queries=[mock_query])
+
+    # Run the tool
+    result = await tavily_tool.run_async(input_schema)
+
+    # Assertions
+    assert len(result.results) == 2
+    assert result.results[0].title == "Result 1"
+    assert result.results[1].title == "Result 2"
+
+    # Validate that the answer is included in each result
+    assert result.results[0].answer == mock_answer
+    assert result.results[1].answer == mock_answer
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])