Update agent.py

agent.py

@@ -6,7 +6,8 @@ from langchain_core.messages import HumanMessage, AIMessage
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.runnables import Runnable
 # from llama_index.core.agent.workflow import AgentWorkflow, ReActAgent
-from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
+# from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
+from langchain_huggingface import HuggingFaceEndpoint
 from tools import TOOLS
 import pandas as pd
 
@@ -21,7 +22,11 @@ def build_graph():
     # Initialize the LLM (e.g., Gemini Flash, zero temperature)
     # llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
     # llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-    llm = HuggingFaceInferenceAPI(model_name="Qwen/Qwen2.5-Coder-32B-Instruct")
+    # llm = HuggingFaceInferenceAPI(model_name="Qwen/Qwen2.5-Coder-32B-Instruct")
+    llm = HuggingFaceEndpoint(
+        url="https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1",
+        huggingfacehub_api_token=os.environ["HF_TOKEN"]
+    )
     llm_with_tools = llm.bind_tools(TOOLS)
 
     # Step 1: Retriever node