df_llm.py

from contextlib import redirect_stdout
import pandas as pd
from langchain.chains import LLMChain
from langchain.chat_models import ChatOpenAI
from io import StringIO
from langchain import PromptTemplate

prompt_template = """
This is the head of the dataframe df {df_head}. 
This is the list of df column names {df_columns}.
Use the information given above to create a Python script that can be executed using a dataframe called df to answer the Question.
Don't create any new df from the given information.
Import necessary libraries in the Python script you create.
Add a print function that prints the result at the end of the Python script.
Question {question}:
"""

# Read df
df = ....
df_head = df.head().to_string()
df_columns = list(df.columns)

prompt = PromptTemplate(
    input_variables=["df_head", "df_columns", "question"], template=prompt_template
)
llm = LLMChain(llm=ChatOpenAI(temperature=0, model_name='gpt-3.5-turbo', max_tokens=1500), prompt=prompt)

command = llm.predict(df_head=df_head, df_columns=df_columns, question=query)

try:
    f = StringIO()
    with redirect_stdout(f):
        exec(command, {'df': df})
    output = f.getvalue()
    print(output)
except Exception as ex:
    print(f'Exception {ex}')