|
import os
|
|
from langchain_experimental.agents import create_pandas_dataframe_agent
|
|
from langchain_google_genai import ChatGoogleGenerativeAI
|
|
import pandas as pd
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
def format_agent_output(output):
|
|
if isinstance(output, pd.DataFrame):
|
|
return output
|
|
elif isinstance(output, str):
|
|
if 'DataFrame' in output or 'describe' in output:
|
|
return pd.DataFrame(eval(output.split('Input:')[-1].strip()))
|
|
return output
|
|
return str(output)
|
|
|
|
def readData(path):
|
|
try:
|
|
df = pd.read_csv(path)
|
|
return df
|
|
except Exception as e:
|
|
raise Exception(f"Error reading data: {str(e)}")
|
|
|
|
def getAgent(data):
|
|
try:
|
|
llm = ChatGoogleGenerativeAI(
|
|
model="gemini-pro",
|
|
temperature=0.5,
|
|
google_api_key=os.environ.get("GOOGLE_API_KEY")
|
|
)
|
|
|
|
agent = create_pandas_dataframe_agent(
|
|
llm,
|
|
data,
|
|
verbose=True,
|
|
handle_parsing_errors=True,
|
|
allow_dangerous_code=True
|
|
)
|
|
return agent
|
|
except Exception as e:
|
|
raise Exception(f"Error creating agent: {str(e)}") |
