Chat with SQL database (MySQL)

sqlchain.py

from langchain.llms import OpenAI
from langchain import OpenAI, SQLDatabase, SQLDatabaseChain
from langchain.prompts.prompt import PromptTemplate
import os

def chat_with_mysql_database(prompt: str):
    """
    Args:
        prompt (str): Question to ask e.g. "Generate an SQL query to retrieve all users"
    Returns:
        str: Answer e.g. "SELECT * from user;"
    """
    db = SQLDatabase.from_uri("mysql+pymysql://your_db_user:[email protected]/db_name")
    llm = OpenAI(temperature=0, verbose=True, openai_api_key=os.environ.get('OPENAI_API_KEY'))
    db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True, use_query_checker=True, top_k=1)
    result = db_chain.run(prompt)
    return result