run_chain.py

from sqlalchemy import create_engine
import openai
import pandas as pd
from langchain import OpenAI, SQLDatabase, SQLDatabaseChain

# load data in memeory to sql lite database
def load_data(df):
  engine = create_engine('sqlite:///:memory:')

  # Write the data to the SQLite database
  df.to_sql('weather', engine, if_exists='replace', index=False)

  # Check if the data was loaded correctly
  df_loaded = pd.read_sql('SELECT * FROM weather', engine)
  db = SQLDatabase(engine)
  return db

# run langchain sql database chain 
def run_query(query, db):
  openai_token = YOUR_OPENAI_KEY
  llm = OpenAI(temperature=0, verbose=True, openai_api_key=openai_token)
  db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True, use_query_checker=True)
  response = db_chain.run(query)
  return response