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_huggingface import HuggingFaceEndpoint
+from langchain_community.llms import HuggingFaceEndpoint  
 from tools import TOOLS  
 
 # --- Read local QA data for retriever ---
@@ -14,14 +14,18 @@ qa_dict = {
 }
 
 def build_graph():
-    # Initialize Mistral model
+    # Initialize Mistral model - CORRECTED CONFIGURATION
     llm = HuggingFaceEndpoint(
         repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-        task="conversational",
-        huggingfacehub_api_token=os.environ["HF_TOKEN"]
+        task="text-generation",
+        huggingfacehub_api_token=os.environ["HF_TOKEN"],
+        model_kwargs={
+            "max_new_tokens": 512,
+            "temperature": 0.1
+        }
     )
 
-    # Retriever node
+    # Retriever node (unchanged)
     def retriever_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
         if query in qa_dict:
@@ -30,36 +34,32 @@ def build_graph():
         print("🔍 No match. Sending to LLM.")
         return {"messages": state["messages"]}
 
-    # Assistant node (LLM)
+    # Assistant node (LLM) - UPDATED FOR TEXT-GENERATION
     def assistant_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
 
+        # Format system prompt for text-generation
         system_prompt = (
-            "You are a helpful assistant evaluated by the GAIA benchmark.\n"
-            "Only return the final answer, with no explanations.\n"
-            "- No prefixes like 'Final answer:'\n"
-            "- If it's a list, output comma-separated\n"
-            "- If unknown, say 'Unknown'\n"
+            "You are a helpful assistant evaluated by the GAIA benchmark. "
+            "Only return the final answer, with no explanations. "
+            "- No prefixes like 'Final answer:' "
+            "- If it's a list, output comma-separated "
+            "- If unknown, say 'Unknown' "
             "- Never justify or explain"
         )
-
-        prompt = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": query},
-        ]
         
-        # prompt = f"{system_prompt}\n\nUser: {query}\nAssistant:"
+        # Format prompt for text-generation model
+        prompt = f"<s>[INST] {system_prompt}\n\n{query} [/INST]"
         response = llm.invoke(prompt).strip()
 
-        # response = llm.invoke(messages)
-        
+        # Clean up response
         for tag in ("Final answer:", "Answer:", "assistant:"):
             if response.lower().startswith(tag.lower()):
                 response = response[len(tag):].strip()
         
         return {"messages": [AIMessage(content=response.strip())]}
 
-    # Tool node
+    # Tool node (unchanged)
     def tool_node(state: MessagesState):
         try:
             tool_signal = state.get("tool_call", "")
@@ -80,7 +80,7 @@ def build_graph():
             print(f"⚠️ Tool error: {e}")
             return {"messages": [AIMessage(content="Unknown")]}  # fail-safe
 
-    # Build LangGraph
+    # Build LangGraph (unchanged)
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever_node)
     builder.add_node("assistant", assistant_node)
@@ -94,7 +94,7 @@ def build_graph():
 
     return builder.compile()
 
-# Agent class
+# Agent class (unchanged)
 class BasicAgent:
     def __init__(self):
         print("✅ BasicAgent initialized with retriever + LLM + tools")
@@ -105,4 +105,4 @@ class BasicAgent:
         result = self.graph.invoke({"messages": [HumanMessage(content=question)]})
         answer = result["messages"][-1].content.strip()
         print(f"📤 Answer: {answer}")
-        return answer
+        return answer