LangChain SQL - Agent Setup

agent.py

from langchain.chat_models import ChatOpenAI
from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.sql_database import SQLDatabase
from langchain.agents.agent_types import AgentType

# Set up the OpenAI LLM
llm = ChatOpenAI(openai_api_key=OPENAI_API_KEY, model_name = "gpt-3.5-turbo")

# Set up the SQL database connection
db = SQLDatabase.from_uri('sqlite:///shipments.db')

# Set up the SQL database toolkit
toolkit = SQLDatabaseToolkit(db=db, llm=llm)

# Create the SQL agent
agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    return_intermediate_steps=True
)