from langchain_core.tools.base import BaseTool
from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
import pandas as pd
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents.agent_types import AgentType

class AnswerExcelTool(BaseTool):
    name : str = "answer_excel_tool"
    description: str = "Given the path to a file containing an excel file and a query, this tool tries to get an answer by querying the excel file."

    def _run(query: str, file_path: str) -> str:
        df = pd.read_excel(file_path)

        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0),

        agent_executor = create_pandas_dataframe_agent(
            llm,
            df,
            agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=True,
            allow_dangerous_code=True # IMPORTANT: Understand the risks
        )

        return agent_executor(query)