SQLagent.py

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 .env file
dotenv.load_dotenv()

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

# Get the SQL and file tools from the ComposioToolSet
tools = toolset.get_tools(apps=[App.SQLTOOL, App.FILETOOL])
file_tool = toolset.get_tools(apps=[App.FILETOOL])

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

code_interpreter_tools = ComposioToolSet().get_tools([App.CODEINTERPRETER])
sql_tools = ComposioToolSet().get_tools([App.SQLTOOL])

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=f"""Run a code to get acheive 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,
    )

    code_interpreter_task = Task(
        description=f"""Run Python code to get acheive a task - {main_task}. Exit once image has been created.""",
        expected_output=f"""Python code executed successfully. Return the image path.""",
        agent=code_interpreter_agent,
    )

    sql_agent = Agent(
        role="SQL Agent",
        goal=f"""Run SQL queries to get acheive 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,
        memory=False,
        allow_delegation=True,
    )

    sql_task = Task(
        description=f"""Run SQL queries to get acheive 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)