Update agent.py

agent.py

@@ -2,7 +2,7 @@ import os
 import pandas as pd
 from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.graph import StateGraph, MessagesState
-from langchain_community.llms import HuggingFaceEndpoint  
+from langchain_community.llms import HuggingFaceHub  
 from tools import TOOLS  
 
 # --- Read local QA data for retriever ---
@@ -14,15 +14,15 @@ qa_dict = {
 }
 
 def build_graph():
-    # Initialize Mistral model - CORRECTED CONFIGURATION
-    llm = HuggingFaceEndpoint(
+    # Initialize Mistral model using HuggingFaceHub
+    llm = HuggingFaceHub(
         repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-        task="text-generation",
-        huggingfacehub_api_token=os.environ["HF_TOKEN"],
-        # model_kwargs={
-        #     "max_new_tokens": 512,
-        #     "temperature": 0.1
-        # }
+        model_kwargs={
+            "max_new_tokens": 512,
+            "temperature": 0.1,
+            "do_sample": False
+        },
+        huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
 
     # Retriever node (unchanged)
@@ -34,11 +34,11 @@ def build_graph():
         print("🔍 No match. Sending to LLM.")
         return {"messages": state["messages"]}
 
-    # Assistant node (LLM) - UPDATED FOR TEXT-GENERATION
+    # Assistant node (LLM) - updated for HuggingFaceHub
     def assistant_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
 
-        # Format system prompt for text-generation
+        # Format system prompt
         system_prompt = (
             "You are a helpful assistant evaluated by the GAIA benchmark. "
             "Only return the final answer, with no explanations. "
@@ -48,8 +48,10 @@ def build_graph():
             "- Never justify or explain"
         )
         
-        # Format prompt for text-generation model
+        # Format prompt for Mistral model
         prompt = f"<s>[INST] {system_prompt}\n\n{query} [/INST]"
+        
+        # Generate response
         response = llm.invoke(prompt).strip()
 
         # Clean up response