How to build a SQL Agent with CrewAI and Composio?

Introduction

SQL is easily one of the most important languages ​​in the world of computing. It serves as the primary means of communicating with relational databases, where most organizations store crucial data. SQL plays an important role, including analyzing complex data, creating data pipelines, and efficiently managing data warehouses. However, writing optimized SQL queries can often be challenging and cumbersome. But thanks to the rapid progress of ai in recent years, we now have ai agents powered by large language models capable of writing queries on our behalf.

This article demonstrates how to build an ai agent using CrewAI, Composio, and Gemini to access databases and run SQL queries to retrieve data.

Learning objectives

• Understand what CrewAI is.
• Learn more about Composio tools and integrations.
• Understand the ai ​​agent workflow.
• A step-by-step guide to building a SQL agent using Composio and CrewAI.

This article was published as part of the Data science blogathon.

What is CrewAI?

ai Crew it's a open source multi-agent collaborative framework. It allows you to build a team of ai agents with diverse tasks, tools, roles, and motivations similar to a real-world team. CrewAI manages the flow of information from one agent to another, allowing you to create efficient, autonomous agent workflows.

CrewAI mainly consists of five main features: Agents, Tasks, Tools, Processes and Tasks.

• Agents: Agents are the fundamental unit of CrewAI and are responsible for making decisions, performing tasks, and communicating with other agents.
• Tasks:These are the objectives that agents are motivated to meet. A task can be performed by one or more agents.
• Tools:Tools allow agents to interact with the external environment, such as using a web scrapper to retrieve breaking news or a scheduler to schedule calendar events.
• Process: The Process is responsible for managing tasks in CrewAI. It assigns tasks to agents in a defined order. The process can be sequential, hierarchical, or consensual. In a sequential process, one task follows another; in a hierarchical process, a managerial hierarchy dictates the order of tasks; and in a consensual process, agents perform tasks collaboratively.
• Crews: Teams within CrewAI are made up of collaborative agents equipped with tasks and tools, all working together to accomplish complex tasks.

Here is a mind map from CrewAI.

Read also: How to create collaborative ai agents with CrewAI?

What is Composio?

Composition it's a open source platform that provides tool solutions for building trustworthy and useful ai agents. Composio provides over 150 tools and applications with built-in user authentication and authorization to help developers build trusted, secure, production-ready agent workflows. The tools have been designed from the ground up with real-world readiness in mind.

Composio offers several advantages over other tool solutions, including managed user authentication and authorizations, a wide range of tools and integrations, a dashboard to monitor live integrations, and the flexibility to add custom tools.

Composio has four key concepts.

• EntitiesIn Composio, an “entity” is a container for all user or organization accounts and tools, allowing centralized management from a single dashboard.
• Integrations: Integrations link your account with external apps, which involves setting up authentication mechanisms like OAuth and setting access permissions to control app actions. Once established, all users can access the same integration through their accounts.
• Behavior: These are tasks performed by built-in tools, such as sending a Slack message or scheduling a calendar event.
• Triggers: Predefined conditions trigger webhooks to your agents when met and send event details such as entities, message text, and more.

<h2 class="wp-block-heading" id="h-designing-an-ai-powered-sql-agent-workflow”>Design an ai-based SQL Agent workflow

Here, you will create a system agent, which receives a user query for data stored in a database, retrieves the relevant data using a SQL agent, and creates beautiful graphs to visualize the data. For this workflow, we will use CrewAI to organize agents and Composio for tool support.

The magnetic system will have a SQL query writer agent and an encoder agent to write and execute the queries. The SQL agent will have access to the SQL tool from Composio and the coding agent will be able to use the E2B Codeinterpreter through Composio. Codeinterpreter provides an isolated environment for running Python programs.

The SQL agent will connect to a local database and query from an appropriate table. The encoding agent will use the data obtained from the SQL agent to create charts. The computer will return the plot as the final result.

Prerequisites for creating SQL Agent

To run this project successfully, you will need the Composio API key. First, create an account on Composition and later we will explain how to get the API keys. Plus, get a free Gemini account API key Google ai Studio, but remember that the free account has a limited speed. It may take longer to run.

Additionally, you will need a dummy database to run queries. If you don't have an additional database, run the following code to create an employee table with names, departments, and salaries.

import sqlite3

# Connect to the SQLite database
connection = sqlite3.connect('company.db')

# Create a cursor object
cursor = connection.cursor()

# Create the 'employee' table
create_table_query = '''
CREATE TABLE IF NOT EXISTS employee (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    employee_name TEXT NOT NULL,
    department TEXT NOT NULL,
    salary INTEGER NOT NULL
);
'''
cursor.execute(create_table_query)

# Data to insert into the 'employee' table
employees = (
    ("John Doe", "Engineering", 70000),
    ("Jane Smith", "Human Resources", 65000),
    ("Alice Johnson", "Marketing", 72000),
    ("Bob Brown", "Sales", 68000),
    ("Charlie Black", "Engineering", 71000),
    ("Daisy White", "Human Resources", 66000),
    ("Edward Green", "Marketing", 69000),
    ("Fiona Grey", "Sales", 64000),
    ("George Yellow", "Engineering", 73000),
    ("Hannah Blue", "Human Resources", 61000),
    ("Ivan Purple", "Marketing", 75000),
    ("Jessica Cyan", "Sales", 70000),
    ("Kyle Red", "Engineering", 68000),
    ("Lily Orange", "Human Resources", 67000),
    ("Martin Indigo", "Marketing", 72000),
    ("Nina Teal", "Sales", 65000),
    ("Oscar Lime", "Engineering", 73000),
    ("Penny Olive", "Human Resources", 62000),
    ("Quentin Silver", "Marketing", 74000),
    ("Rachel Maroon", "Sales", 69000),
    ("Steve Pink", "Engineering", 71000),
    ("Tina Violet", "Human Resources", 68000),
    ("Ursula Gold", "Marketing", 76000),
    ("Victor Bronze", "Sales", 64000),
    ("Wendy Mauve", "Engineering", 69000),
    ("Xavier Cream", "Human Resources", 65000),
    ("Yolanda Peach", "Marketing", 70000),
    ("Zack Sage", "Sales", 68000),
    ("Abby Coral", "Engineering", 72000),
    ("Bill Moss", "Human Resources", 63000)
)

# Insert data into the 'employee' table
insert_query = 'INSERT INTO employee (employee_name, department, salary) VALUES (?, ?, ?);'
cursor.executemany(insert_query, employees)

# Commit the changes
connection.commit()

# Close the connection
connection.close()

print("Table created and data inserted successfully.")

Step-by-step guide to creating a SQL agent

Now that we've covered the basics, we can start with the coding part. As with any Python project, we'll first set up a virtual environment and environment variables, and install libraries. The project will use Gemini 1.5 Pro as the language model.

Step 1: Install Libraries

Create a virtual environment using Python Venv.

python -m venv sqlagent
cd sqlagent
source bin/active

Install the following libraries using pip install.

composio-core
composio-crewai
langchain-google-genai
dotenv

Step 2: Set environment variables

To use the Composio toolsets, you must authenticate your Composio account. Run the following command to log in to Composio and follow the login flow.

composio login

Now, get your Composio API keys.

Composio whoami

Create a .env file and add the variables COMPOSIO_API_KEY and GOOGLE_API_KEY.

COMPOSIO_API_KEY=your Composio API key
GOOGLE_API_KEY=your Gemini API key

Step 3: Define tools and LLM

Now, create a Python file and import the required libraries.

import os

import dotenv
from composio_langchain import App, ComposioToolSet
from crewai import Agent, Crew, Process, Task
from langchain_google_genai import ChatGoogleGenerativeAI

# Load environment variables from the .env file
dotenv.load_dotenv()

Define Composio tools.

# Initialize the ComposioToolSet
toolset = ComposioToolSet(api_key=os.environ("COMPOSIO_API_KEY"))

code_interpreter_tools = toolset.get_tools((App.CODEINTERPRETER))
sql_tools = toolset.get_tools((App.SQLTOOL))

Now, define the LLM with Gemini 1.5 Pro.

llm = ChatGoogleGenerativeAI(model="gemini-1.5-pro", 
                             api_key=os.environ('GOOGLE_API_KEY')
                          )

Step 4: Define agents and tasks

As mentioned above, here we will define two agents and two tasks: the SQL agent, the encoding agent and their respective tasks.

code_interpreter_agent = Agent(
        role="Python Code Interpreter Agent",
        goal="Run a code to get achieve a task given by the user",
        backstory="You are an agent that helps users run Python code.",
        verbose=True,
        tools=code_interpreter_tools,
        llm=llm,
    )
    
sql_agent = Agent(
        role="SQL Agent",
        goal="Run SQL queries to get achieve a task given by the user",
        backstory=(
            "You are an agent that helps users run SQL queries. "
            "Connect to the local SQLite DB at connection string = company.db"
            "Try to analyze the tables first by listing all the tables and columns "
            "and doing distinct values for each column and once sure, make a query to \
            get the data you need."
        ),
        verbose=True,
        tools=sql_tools,
        llm=llm,
        allow_delegation=True,
    )

In the code snippet above, we defined the agents, each with a defined role, goal, and backstory. Additional information provides additional context to LLMs before generating query responses. A tool equips each agent to perform the actions.

Now, define the tasks.

code_interpreter_task = Task(
        description=f"Run Python code to achieve the task - {main_task}. \
        Exit once the image has been created.",
        expected_output="Python code executed successfully. Return the image path.",
        agent=code_interpreter_agent,
    )
    
   sql_task = Task(
        description=f"Run SQL queries to achieve a task - {main_task}",
        expected_output=f"SQL queries executed successfully. The result of the task \
        is returned - {main_task}",
        agent=sql_agent,
    )

We define the tasks that the agents will perform. Each task has a description, an expected result and the agent responsible for performing it.

Now, define the Crew with the agents and tasks.

crew = Crew(
        agents=(sql_agent, code_interpreter_agent),
        tasks=(sql_task, code_interpreter_task),
    )
result = crew.kickoff()
print(result)

You can put this flow in a while loop to make it more attractive.

while True:
    main_task = input("Enter the task you want to perform (or type 'exit' to quit): ")
    if main_task.lower() == "exit":
        break
        
    code_interpreter_agent = Agent(
        role="Python Code Interpreter Agent",
        goal="Run a code to get achieve a task given by the user",
        backstory="You are an agent that helps users run Python code.",
        verbose=True,
        tools=code_interpreter_tools,
        llm=llm,
    )
    
     sql_agent = Agent(
        role="SQL Agent",
        goal="Run SQL queries to get achieve a task given by the user",
        backstory=(
            "You are an agent that helps users run SQL queries. "
            "Connect to the local SQLite DB at connection string = company.db"
            "Try to analyze the tables first by listing all the tables and columns "
            "and doing distinct values for each column and once sure, make a query to \
            get the data you need."
        ),
        verbose=True,
        tools=sql_tools,
        llm=llm,
        allow_delegation=True,
    )
    code_interpreter_task = Task(
        description=f"Run Python code to achieve the task - {main_task}. \
        Exit once the image has been created.",
        expected_output="Python code executed successfully. Return the image path.",
        agent=code_interpreter_agent,
    )
    
    sql_task = Task(
        description=f"Run SQL queries to achieve a task - {main_task}",
        expected_output=f"SQL queries executed successfully. The result of the task \
        is returned - {main_task}",
        agent=sql_agent,
    )


    crew = Crew(
        agents=(sql_agent, code_interpreter_agent),
        tasks=(sql_task, code_interpreter_task),
    )

    result = crew.kickoff()
    print(result)

This will prompt you to enter a query, which will then be sent to the ai ​​agent team. After execution, you will have the option to ask another question or exit the loop.

Once a query execution is complete, it will generate the file path for the chart image.

I asked him to create a bar chart of the number of employees in each department. This was the result.

GitHub Gist: SQL agent

These steps demonstrated how to create an agent workflow to automate the extraction and visualization of SQL data. However, you can go further and make it more robust and reliable by adding a memory component to the agents and crew. This will help agents remember their past results, allowing them to better direct workflow. You can also add an interface with Streamlit or Gradio to make it more interactive.

Conclusion

The ai ​​landscape is evolving at an unprecedented pace. As the quality of ai models, frameworks, and tools continues to improve, building powerful ai agents becomes increasingly desirable. The future of the workforce is agentic, where humans and ai will complement each other to create even more efficient systems. With frameworks like CrewAI and Composio, you can conveniently create ai workflows to automate many routine tasks. This article demonstrates how to automate data extraction and visualization. You can extend this workflow to handle even more complex scenarios.

Key findings

• CrewAI is an open-source framework for organizing LLM agents to perform complex tasks collaboratively through role assignment, goal sharing, and task delegation.
• Composio is an open-source tool platform that offers production-ready integrations and tools to empower ai agents to reliably perform tasks.
• You can integrate Composio tools with popular platforms such as LangChain, Autogen, CrewAI, and Llamaindex.

Frequent questions

The media shown in this article is not the property of Analytics Vidhya and is used at the author's discretion.