Spaces:
Sleeping
Sleeping
File size: 2,030 Bytes
5b12cd5 702a58f 5b12cd5 702a58f 5b12cd5 cca7719 5b12cd5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import getpass
import os
import pandas as pd
import tabulate
from sqlalchemy import create_engine
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent
from langchain_openai import ChatOpenAI
from langchain.output_parsers.openai_tools import JsonOutputKeyToolsParser
from langchain_experimental.agents import create_pandas_dataframe_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_experimental.tools import PythonAstREPLTool
from langchain_experimental.agents import create_pandas_dataframe_agent
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(self, 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(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 |