Update app.py

app.py

@@ -26,6 +26,8 @@ from langchain.llms import OpenAI
 from langchain.agents import create_pandas_dataframe_agent
 import pandas as pd
 import docx
+from pandasai import PandasAI
+from pandasai.llm.openai import OpenAI as pai_openai
 
 
 
@@ -42,6 +44,7 @@ stability_api = client.StabilityInference(
 whisper_from_pipeline = pipeline("automatic-speech-recognition",model="openai/whisper-medium")
 EMBEDIDNGS = None
 DATAFRAME_FILE = None
+DATAFRAME = None
 DOCSEARCH = None
 RANDOM_USER = ''.join(chr(random.randint(65,90)) for i in range(8))+f'{random.randint(1,10000000000)}'
 print(f'{RANDOM_USER} chat started')
@@ -156,7 +159,7 @@ def search_document_uploaded(user_query:str) -> str:
 
 
 def ask_dataframes(user_query):
-  return DATAFRAME_FILE.run(user_query)
+  return DATAFRAME_FILE.run(DATAFRAME, prompt = user_query)
 
 ############# GET OPENAI RESPONSE
 def get_open_ai_reponse(user_query:str)->Union[tuple,str]:
@@ -307,18 +310,14 @@ def ask_questions_abt_dataframes(file,file_ext):
   print(file_ext)
   global EMBEDIDNGS
   EMBEDIDNGS = None
-  global DOCSEARCH
-  DOCSEARCH = None
-
-  
 
   reader_function = { '.csv': pd.read_csv, '.xlsx': pd.read_excel }.get(file_ext)
   print(reader_function.__name__)
   global DATAFRAME_FILE
-  DATAFRAME_FILE = create_pandas_dataframe_agent(
-    OpenAI(openai_api_key=os.environ['OPENAI_API_KEY']),
-    reader_function(file.name)
-  )
+  global DATAFRAME
+  DATAFRAME = reader_function(file.name)
+  llm = pai_openai(api_token=os.environ['OPENAI_API_KEY'])
+  DATAFRAME_FILE = PandasAI(llm)