Update main.py

main.py

@@ -6,8 +6,6 @@ import numpy as np
 from typing import List
 from pathlib import Path
 
-from langchain_huggingface import HuggingFaceEndpoint
-
 from langchain_openai import ChatOpenAI, OpenAI
 from langchain.schema.runnable.config import RunnableConfig
 from langchain.schema import StrOutputParser
@@ -61,16 +59,9 @@ def create_agent(filename: str):
     """
 
     # Create an OpenAI object.
-    #os.environ['OPENAI_API_KEY'] = os.environ['OPENAI_API_KEY']
-    #llm = ChatOpenAI(temperature=0, model="gpt-4o-2024-05-13")
-    
-    # Create an HuggingFace Mistral object.
-    os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.environ['HUGGINGFACEHUB_API_TOKEN']
+    os.environ['OPENAI_API_KEY'] = os.environ['OPENAI_API_KEY']
+    llm = ChatOpenAI(temperature=0, model="gpt-4o-2024-05-13")
     
-    repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-    llm = HuggingFaceEndpoint(
-        repo_id=repo_id, task="text2text-generation", max_new_tokens=8000, temperature=0.01, streaming=True
-    )
     # Read the CSV file into a Pandas DataFrame.
     if cl.user_session.get("createdb") == None:
         df = pd.read_csv(filename)
@@ -83,7 +74,7 @@ def create_agent(filename: str):
         db = cl.user_session.get("db")
     # Create a SAL agent.
     #e.g agent_executor.invoke({"input": "Quel est le nombre de chargé d'affaires en agencement par entreprise?"})
-    return create_sql_agent(llm, db=db, agent_type="zero-shot-react-description", verbose=False)
+    return create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=False)
 
 def query_agent(agent, query):
     """