SQL Agent

1
pip install composio-langchain langchain-openai

1
import os
2
import dotenv
3
4
from composio_langchain import App, ComposioToolSet
5
from langchain import hub
6
from langchain.agents import AgentExecutor, create_openai_functions_agent
7
from langchain_openai import ChatOpenAI
8
9
# Load environment variables
10
dotenv.load_dotenv()

1
# Initialize the language model
2
llm = ChatOpenAI(model="gpt-4o")
3
4
# Initialize the Composio ToolSet
5
composio_toolset = ComposioToolSet()
6
7
# Get tools for SQL and File operations
8
sql_file_tools = composio_toolset.get_tools(apps=[App.SQLTOOL, App.FILETOOL])
9
10
# Get tools for code interpreter (for plotting graphs)
11
code_tools = composio_toolset.get_tools(apps=[App.CODEINTERPRETER])
12
13
# Get all tools combined
14
all_tools = composio_toolset.get_tools(apps=[App.SQLTOOL, App.FILETOOL, App.CODEINTERPRETER])
15
16
# Pull the prompt template for the agent
17
prompt = hub.pull("hwchase17/openai-functions-agent")

1
# Create the agent for SQL and File operations
2
sql_agent = create_openai_functions_agent(llm, sql_file_tools, prompt)
3
sql_executor = AgentExecutor(agent=sql_agent, tools=sql_file_tools, verbose=True)
4
5
# Define the SQL query task
6
query_task = (
7
    "Connect to the SQLite database 'companydb' and execute the following query: "
8
    "SELECT * FROM MOCK_DATA ORDER BY salary DESC LIMIT 10. "
9
    "Save the results to a file named 'top_salaries.txt' and return the content."
10
)
11
12
# Execute the SQL query task
13
sql_result = sql_executor.invoke({"input": query_task})
14
print(f"SQL Query Result: {sql_result['output']}")

1
# Create the agent for data visualization
2
visualization_agent = create_openai_functions_agent(llm, code_tools, prompt)
3
visualization_executor = AgentExecutor(agent=visualization_agent, tools=code_tools, verbose=True)
4
5
# Define the visualization task
6
visualization_task = (
7
    f"Using the following data from the top 10 highest salaries: {sql_result['output']}\n"
8
    "Create a bar chart showing first_name on the x-axis and salary on the y-axis. "
9
    "Save the visualization as 'salary_chart.png'."
10
)
11
12
# Execute the visualization task
13
visualization_result = visualization_executor.invoke({"input": visualization_task})
14
print(f"Visualization Result: {visualization_result['output']}")

1
# Execute the workflow
2
def main():
3
    # First, execute SQL query
4
    sql_result = sql_executor.invoke({"input": query_task})
5
    print(f"SQL Query Result: {sql_result['output']}")
6
    
7
    # Then, create visualization
8
    visualization_task = (
9
        f"Using the following data from the top 10 highest salaries: {sql_result['output']}\n"
10
        "Create a bar chart showing first_name on the x-axis and salary on the y-axis. "
11
        "Save the visualization as 'salary_chart.png'."
12
    )
13
    visualization_result = visualization_executor.invoke({"input": visualization_task})
14
    print(f"Visualization Result: {visualization_result['output']}")
15
    
16
    return "Workflow completed successfully!"
17
18
if __name__ == "__main__":
19
    main()