agents/dataframe_agent.py

import pandas as pd
from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
from agents.llms import LLM
from langchain.agents.agent_types import AgentType


def get_dataframe_agent(
    df: pd.DataFrame,
    verbose: bool = True,
    allow_dangerous_code: bool = True,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION
):
    """
    Create a pandas DataFrame agent using the custom LLM.
    Args:
        df (pd.DataFrame): The pandas DataFrame to use.
        verbose (bool): Whether to enable verbose output. Default is True.
        allow_dangerous_code (bool): Whether to allow dangerous code execution. Default is True.
        agent_type: The agent type to use. Default is ZERO_SHOT_REACT_DESCRIPTION.
    Returns:
        agent: The created DataFrame agent.
    """
    llm = LLM().chat_model
    agent = create_pandas_dataframe_agent(
        llm,
        df,
        agent_type=agent_type,
        verbose=verbose,
        allow_dangerous_code=allow_dangerous_code
    )
    return agent

# Usage example:
# import pandas as pd
# from agents.dataframe_agent import get_dataframe_agent
# df = pd.read_csv('your_file.csv')
# agent = get_dataframe_agent(df)
# response = agent.invoke('Your question here')
# print(response)