Update main.py

main.py

@@ -5,7 +5,7 @@ import pandas as pd
 import numpy as np
 from typing import List
 from pathlib import Path
-from langchain_huggingface import HuggingFaceEndpoint
+from langchain_openai import ChatOpenAI
 from langchain.schema.runnable.config import RunnableConfig
 from langchain.schema import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
@@ -48,17 +48,14 @@ def create_agent(filename: str):
     """
 
     # Create an OpenAI object.
-    os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.environ['HUGGINGFACEHUB_API_TOKEN']
-    repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-    llm = HuggingFaceEndpoint(
-        repo_id=repo_id, max_new_tokens=5300, temperature=0.1, task="text2text-generation", streaming=True
-    )
+    os.environ['OPENAI_API_KEY'] = os.environ['OPENAI_API_KEY']
+    llm = ChatOpenAI(temperature=0, model="gpt-4o-2024-05-13")
 
     # Read the CSV file into a Pandas DataFrame.
     df = pd.read_csv(filename)
 
     # Create a Pandas DataFrame agent.
-    return create_pandas_dataframe_agent(llm, df, verbose=False, allow_dangerous_code=True)
+    return create_pandas_dataframe_agent(llm, df, verbose=False, allow_dangerous_code=True, agent_type=AgentType.OPENAI_FUNCTIONS)
 
 def query_agent(agent, query):
     """