Coding Your First AutoGen Tool: Tavily Search Walkthrough

Lennex Zinyando

• 7 min read

AI agents are software entities that perform tasks, powered by LLMs. They take in inputs, reason, make decisions, and take action. However, their true potential is unlocked when we extend their capabilities beyond their inherent knowledge. In this blog post, we'll be learning how to extend the capabilities of AI agents built using AutoGen by creating our very own tool using the Tavily Search API.

AutoGen is a powerful framework for building multi-agent AI systems. It allows developers to create conversational AI agents that can leverage various tools and APIs to enhance their capabilities. By integrating Tavily Search API into our AutoGen agent, we're implementing a form of Retrieval-Augmented Generation (RAG). This allows our AI to access up-to-date information from the internet, significantly enhancing its knowledge and capabilities. Don't worry if you're new to this - we'll break everything down into easy-to-follow steps. By the end of this tutorial, you'll have the skills to create custom tools that can significantly enhance your AI projects.

If you are in a rush and just want to see the code head over to GitHub.

Setting Up Our Environment

Before we start building the Tavily Search tool, we need to set up our development environment.

Python Installation

This tutorial assumes you're using Python 3.8 or later. Ensure your Python installation is up to date before proceeding.

If you don't have Python installed, visit the official Python website and follow the installation instructions for your operating system.

Installing the dependencies

We'll the Python packages for AutoGen, tavily-python, and other dependencies that will make the development of our chatbot easier:

pip install autogen tavily-python groq python-dotenv pydantic

Getting the TAVILY_API_KEY

To use the Tavily search API, you'll need to obtain an API key. Follow these steps:

  1. Go to the Tavily website at https://tavily.com/
  2. If you don't have an account, you'll need to sign up first. Look for a "Dashboard" button on the homepage.
  3. Once you're logged in, navigate to your account dashboard.
  4. Look for an option to generate or view your API key. It might be labeled as "API Keys".
  5. Generate a new API key if one isn't already available.
  6. Copy the API key and store it securely. You'll need to use this key in your code to authenticate requests to the Tavily search API.
  7. Be aware that Tavily may offer different pricing tiers or a free trial. Check their pricing page for current offers and any usage limits.

Remember to keep your API key confidential and never share it publicly. You may want to set it as an environment variable in your development environment for added security.

Getting the GROQ_API_KEY

Groq will provide our language model. Follow these steps:

With these steps completed, your environment should be ready for building our RAG chatbot!

Building our search tool

Let's create our tool step by step. Create a project directory named basic-autogen-tavily-tool. Inside this directory, create a .env file to store our API keys

TAVILY_API_KEY=tvly-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
GROQ_API_KEY=gsk_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx

Import Required Libraries

Create a file named chatbot.py in your project directory. We'll build our chatbot in this file.

First, let's import the necessary libraries:

import os
from typing import Annotated
from tavily import TavilyClient
from pydantic import BaseModel, Field
from autogen import register_function, ConversableAgent
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

This section imports the necessary libraries and loads environment variables. We're using:

  • os for environment variable access
  • typing and pydantic for type hinting and data validation
  • tavily for the Tavily Search API
  • autogen for creating conversational agents
  • dotenv for loading environment variables from a .env file

Configuring the Language Model

llm_config = {
    "config_list": [
        {
            "model": "gemma2-9b-it",
            "api_key": os.getenv("GROQ_API_KEY"),
            "api_type": "groq",
        }
    ]
}

Here, we configure the language model to be used by our agents. In this case, we're using the Gemma 2 9B model through the Groq API. The API key is loaded from an environment variable.

Creating Conversational Agents

assistant = ConversableAgent(
    name="Assistant",
    system_message="You are a helpful AI assistant with access to internet search capabilities.",
    llm_config=llm_config,
)

user_proxy = ConversableAgent(name="User", human_input_mode="NEVER", llm_config=False)

Here, we create two agents:

  1. An assistant agent that can use the language model and respond to queries.
  2. A user proxy agent that can execute the Tavily search tool but doesn't use the language model itself.

Defining the Tavily Search Input Model

class TavilySearchInput(BaseModel):
    query: Annotated[str, Field(description="The search query string")]
    max_results: Annotated[
        int, Field(description="Maximum number of results to return", ge=1, le=10)
    ] = 5
    search_depth: Annotated[
        str,
        Field(
            description="Search depth: 'basic' or 'advanced'",
            choices=["basic", "advanced"],
        ),
    ] = "basic"

This section defines the structure and validation rules for the Tavily search input using Pydantic. Pydantic is a powerful library for data validation and settings management using Python type annotations. It's particularly useful in this context because it allows us to:

  1. Define a clear structure for our input data
  2. Automatically validate incoming data
  3. Provide clear error messages if the input doesn't meet our requirements
  4. Generate OpenAPI (Swagger) schema for our input model

Let's break down the TavilySearchInput class:

  • It inherits from pydantic.BaseModel, which provides the core functionality for data validation.
  • Each field is defined with type annotations and additional metadata:
    • query: A required string field for the search query.
    • max_results: An optional integer field with a default value of 5. The ge=1, le=10 arguments ensure the value is between 1 and 10.
    • search_depth: An optional string field with a default value of "basic". The choices argument restricts the allowed values to "basic" or "advanced".

The Annotated type is used to attach metadata to the fields without affecting their runtime behavior. This is particularly useful for generating documentation or API schemas.

By using Pydantic, we ensure that our function receives valid input data, making our code more robust and easier to maintain.

Implementing the Tavily Search Function

def tavily_search(
    input: Annotated[TavilySearchInput, "Input for Tavily search"]
) -> str:
    # Initialize the Tavily client with your API key
    client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))

    # Perform the search
    response = client.search(
        query=input.query,
        max_results=input.max_results,
        search_depth=input.search_depth,
    )

    # Format the results
    formatted_results = []
    for result in response.get("results", []):
        formatted_results.append(
            f"Title: {result['title']}\\nURL: {result['url']}\\nContent: {result['content']}\\n"
        )

    return "\\n".join(formatted_results)

This function performs the actual Tavily search. It initializes the Tavily client, sends the search request, and formats the results into a string. Here's a detailed breakdown:

  1. The function takes a TavilySearchInput object as its input. This ensures that the input has already been validated by Pydantic before reaching this function.
  2. We initialize the Tavily client using the API key stored in our environment variables. This is a secure way to handle API keys without hardcoding them in our script.
  3. We perform the search using the client.search() method, passing in the parameters from our input object. This allows us to easily customize the search based on the user's requirements.
  4. The search results are then formatted into a more readable string format. We iterate through each result, extracting the title, URL, and content, and formatting them into a consistent structure.
  5. Finally, we join all the formatted results into a single string and return it.

This approach allows us to encapsulate all the search logic in a single function, making it easy to use and maintain.

In a production environment, you'd want to add error handling to deal with potential API failures or network issues. For simplicity, we've omitted this in our example.

Registering the Tavily Search Function

register_function(
    tavily_search,
    caller=assistant,
    executor=user_proxy,
    name="tavily_search",
    description="A tool to search the internet using the Tavily API",
)

This step is crucial for integrating our Tavily search function with the AutoGen framework. Here's what's happening and why:

  1. We use AutoGen's register_function() to make our tavily_search function available to our agents.
  2. The caller parameter is set to assistant. This means that the assistant agent is allowed to call (or request) this function when it needs to perform a search.
  3. The executor parameter is set to user_proxy. This specifies that the actual execution of the search function will be handled by the user proxy agent. This separation of concerns allows for more flexible and secure function execution.
  4. We provide a name for the function. This is the name that the assistant will use when it wants to call the function.
  5. The description provides a brief explanation of what the function does. This can be used by the assistant to understand when and how to use this function.

By registering the function this way, we're effectively giving our AI assistant the ability to perform internet searches when needed to answer user queries. The assistant can decide when a search is necessary and request it, while the user proxy handles the actual execution of the search.

This setup allows for a more dynamic and capable AI assistant that can gather up-to-date information from the internet to supplement its responses.

Main Chat Loop

def main():
    while True:
        user_input = input("User: ")

        if user_input.lower() in ["exit", "quit", "bye"]:
            print("Chatbot: Goodbye! Have a great day!")
            break

        # Initiate a chat between the user proxy and the assistant
        chat_result = user_proxy.initiate_chat(
            assistant,
            message=user_input,
            max_turns=2,
        )

        # Extract the assistant's reply from the chat history
        reply = next(
            (
                msg["content"]
                for msg in chat_result.chat_history
                if msg.get("name") == "Assistant"
            ),
            "I apologize, but I couldn't generate a response.",
        )

        print(f"Chatbot: {reply}")

The main function implements a chat loop that:

  1. Takes user input
  2. Initiates a chat between the user proxy and the assistant
  3. Extracts the assistant's reply
  4. Prints the response

Starting the chatbot

To start the chatbot, navigate to your project directory in the terminal and run:

python chatbot.py

You can ask the chatbot questions and it the terminal you will get confirmation that the tavily_search tool was used to get an answer from the Tavily Search API.

Conclusion

In this walkthrough, we've explored how to create a powerful and intelligent chatbot using AutoGen and the Tavily Search API. This combination allows us to build AI assistants that can not only engage in conversation but also access up-to-date information from the internet to provide more accurate and informative responses.

Let's recap the key components of our AutoGen tool:

  1. We used Pydantic to create a robust input model for our Tavily search, ensuring data validation and clear structure.
  2. We implemented a Tavily search function that interfaces with the API and formats the results for easy consumption.
  3. We created two AutoGen agents: an assistant capable of natural language interaction and a user proxy that can execute the search function.
  4. We registered the Tavily search function with AutoGen, allowing our assistant to seamlessly incorporate internet search capabilities into its responses.

This project demonstrates the power of combining large language models with external tools and APIs. By giving our AI assistant the ability to search the internet, we've significantly expanded its knowledge base and usefulness. This approach can be extended to incorporate other APIs and tools, creating even more capable AI assistants.

Remember, the code is available on GitHub. If you have any questions, you can reach out to me on X (formerly Twitter) or LinkedIn.

AI should drive results, not complexity. AgentemAI helps businesses build scalable, efficient, and secure AI solutions. See how we can help.