from langchain_experimental.agents import create_pandas_dataframe_agent
import pandas as pd
from langchain.agents.agent_types import AgentType
from langchain_community.llms import HuggingFaceEndpoint


def query_agent(data, query):
    # Parse the CSV file and create a Pandas DataFrame from its contents.
    df = pd.read_csv(data)

    llm = HuggingFaceEndpoint(repo_id="HuggingFaceH4/zephyr-7b-beta",
                              temperature = 0.9,
                              stop_sequences =  ["\nQuestion:"])

    # Create a Pandas DataFrame agent with the loaded model and processor
    agent = create_pandas_dataframe_agent(llm, 
                                          df, 
                                          verbose=True, 
                                          handle_parsing_errors=True,
                                          )

    # Modify the query to include details for your new model usage
    modified_query = query + " using tool python_repl_ast or any other relevant tool if needed. First Command you find is most likely to be the most accurate one so use it."

    

    # Run the agent with the modified model and processor instance
    return agent.run(modified_query)