Update agent.py

agent.py

@@ -2,11 +2,10 @@ import os
 import pandas as pd
 from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.graph import StateGraph, MessagesState
-from langgraph.prebuilt import ToolNode, tools_condition
-from langchain_community.chat_models import ChatHuggingFace
-from tools import TOOLS
+from langchain_huggingface import HuggingFaceEndpoint
+from tools import TOOLS  
 
-# --- Load local QA metadata for retriever ---
+# --- Read local QA data for retriever ---
 QA_PATH = "metadata.jsonl"
 qa_pairs = pd.read_json(QA_PATH, lines=True)
 qa_dict = {
@@ -15,14 +14,13 @@ qa_dict = {
 }
 
 def build_graph():
-    # ✅ Use correct conversational wrapper for Hugging Face models
-    llm = ChatHuggingFace(
+    # Initialize Mistral model
+    llm = HuggingFaceEndpoint(
         repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-        task="conversational",
         huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
 
-    # Node 1: Retriever
+    # Retriever node
     def retriever_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
         if query in qa_dict:
@@ -31,45 +29,47 @@ def build_graph():
         print("🔍 No match. Sending to LLM.")
         return {"messages": state["messages"]}
 
-    # Node 2: Assistant
+    # Assistant node (LLM)
     def assistant_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
-        prompt = (
-            "You are a helpful assistant in the GAIA benchmark.\n"
-            "If you know the answer, reply directly with only the answer.\n"
-            "If tool use is needed, reply in the format:\n"
-            "use_tool: <tool_name>: <tool_input>\n"
-            "Do not explain anything."
+
+        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"
+            "- Never justify or explain"
         )
+
         messages = [
-            {"role": "system", "content": prompt},
+            {"role": "system", "content": system_prompt},
             {"role": "user", "content": query},
         ]
-        response = llm.invoke(messages).strip()
-        print(f"🧠 LLM responded: {response}")
 
-        if response.lower().startswith("use_tool:"):
-            return {
-                "messages": state["messages"] + [AIMessage(content=response)],
-                "tool_call": response,
-            }
-        return {"messages": [AIMessage(content=response)]}
+        response = llm.invoke(messages)
+        return {"messages": [AIMessage(content=response.strip())]}
 
-    # Node 3: Tool Executor
+    # Tool node
     def tool_node(state: MessagesState):
         try:
             tool_signal = state.get("tool_call", "")
             _, tool_name, tool_input = tool_signal.split(":", 2)
-            tool_fn = TOOLS.get(tool_name.strip())
+            tool_name = tool_name.strip()
+            tool_input = tool_input.strip()
+            tool_fn = TOOLS.get(tool_name)
+
             if not tool_fn:
                 print(f"❌ Unknown tool: {tool_name}")
                 return {"messages": [AIMessage(content="Unknown")]}
-            print(f"🔧 Running tool {tool_name} with input: {tool_input}")
-            result = tool_fn(tool_input.strip())
-            return {"messages": [AIMessage(content=str(result))]}
+            
+            print(f"🔧 Using tool: {tool_name} with input: {tool_input}")
+            tool_result = tool_fn(tool_input)
+            return {"messages": [AIMessage(content=str(tool_result))]}
+
         except Exception as e:
             print(f"⚠️ Tool error: {e}")
-            return {"messages": [AIMessage(content="Unknown")]}
+            return {"messages": [AIMessage(content="Unknown")]}  # fail-safe
 
     # Build LangGraph
     builder = StateGraph(MessagesState)
@@ -85,8 +85,7 @@ def build_graph():
 
     return builder.compile()
 
-
-# --- Final Agent Wrapper ---
+# Agent class
 class BasicAgent:
     def __init__(self):
         print("✅ BasicAgent initialized with retriever + LLM + tools")
@@ -95,4 +94,6 @@ class BasicAgent:
     def __call__(self, question: str) -> str:
         print(f"📥 Question: {question[:100]}")
         result = self.graph.invoke({"messages": [HumanMessage(content=question)]})
-        return result["messages"][-1].content.strip()
+        answer = result["messages"][-1].content.strip()
+        print(f"📤 Answer: {answer}")
+        return answer