File size: 1,254 Bytes
da4182b |
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 |
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 # Enable code execution
)
return agent
except Exception as e:
raise Exception(f"Error creating agent: {str(e)}") |