wrappped into class

ai_agent.py

@@ -1,3 +1,4 @@
+import getpass
 import os
 import pandas as pd
 import tabulate
@@ -12,33 +13,35 @@ from langchain_core.prompts import ChatPromptTemplate
 from langchain_experimental.tools import PythonAstREPLTool
 from langchain_experimental.agents import create_pandas_dataframe_agent
 
-# Get the API key
-api_key = os.getenv('OPENAI_API_KEY')
-llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.5)
+class AIAgent:
+    def __init__(self):
+        # Get the API key
+        self.api_key = os.getenv('OPENAI_API_KEY')
+        self.llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.5)
 
-def sql_agent(file, agent_input):
-    # Read the uploaded file into a DataFrame
-    df = pd.read_csv(file)
-    # Create SQLAlchemy engine
-    engine = create_engine("sqlite:///uploaded_data.db")
-    # Write the DataFrame to the SQLite database
-    df.to_sql("uploaded_data", engine, index=False, if_exists='replace')
-    db = SQLDatabase(engine=engine)
-    # Create SQL agent
-    agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)
-    agent_output = agent_executor.invoke(agent_input)
-    return agent_output
-
-def pandas_agent(file, agent_input):
-    # Check if the file extension is CSV
-    if file.name.endswith('.csv'):
+    def sql_agent(self, file, agent_input):
+        # Read the uploaded file into a DataFrame
         df = pd.read_csv(file)
-    # Check if the file extension is XLSX or XLS
-    elif file.name.endswith('.xlsx') or file.name.endswith('.xls'):
-        df = pd.read_excel(file)
-    else:
-        return "Unsupported file format. Only CSV, XLS, or XLSX files are supported."
-    # Proceed with your agent code
-    agent = create_pandas_dataframe_agent(llm, df, agent_type="openai-tools", verbose=True)
-    agent_output = agent.invoke(agent_input)
-    return agent_output
+        # Create SQLAlchemy engine
+        engine = create_engine("sqlite:///uploaded_data.db")
+        # Write the DataFrame to the SQLite database
+        df.to_sql("uploaded_data", engine, index=False, if_exists='replace')
+        db = SQLDatabase(engine=engine)
+        # Create SQL agent
+        agent_executor = create_sql_agent(self.llm, db=db, agent_type="openai-tools", verbose=True)
+        agent_output = agent_executor.invoke(agent_input)
+        return agent_output
+
+    def pandas_agent(self, file, agent_input):
+        # Check if the file extension is CSV
+        if file.name.endswith('.csv'):
+            df = pd.read_csv(file)
+        # Check if the file extension is XLSX or XLS
+        elif file.name.endswith('.xlsx') or file.name.endswith('.xls'):
+            df = pd.read_excel(file)
+        else:
+            return "Unsupported file format. Only CSV, XLS, or XLSX files are supported."
+        # Proceed with your agent code
+        agent = create_pandas_dataframe_agent(self.llm, df, agent_type="openai-tools", verbose=True)
+        agent_output = agent.invoke(agent_input)
+        return agent_output