Update agent.py

agent.py

@@ -2,9 +2,9 @@ import os
 import pandas as pd
 from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.graph import StateGraph, MessagesState
-from langgraph.prebuilt import ToolNode
-from langchain_huggingface import HuggingFaceEndpoint
-from tools import TOOLS  # your dictionary of tool functions
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_community.chat_models import ChatHuggingFace
+from tools import TOOLS
 
 # --- Load local QA metadata for retriever ---
 QA_PATH = "metadata.jsonl"
@@ -14,11 +14,11 @@ qa_dict = {
     for _, row in qa_pairs.iterrows()
 }
 
-# === LangGraph builder ===
 def build_graph():
-    llm = HuggingFaceEndpoint(
+    # ✅ Use correct conversational wrapper for Hugging Face models
+    llm = ChatHuggingFace(
         repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-        # task="text-generation",
+        task="conversational",
         huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
 
@@ -31,58 +31,45 @@ def build_graph():
         print("🔍 No match. Sending to LLM.")
         return {"messages": state["messages"]}
 
-    # Node 2: Assistant (LLM response parsing)
+    # Node 2: Assistant
     def assistant_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
-
         prompt = (
-            "You are a helpful assistant for GAIA benchmark.\n"
-            "If you can answer directly, output ONLY the answer.\n"
-            "If you need to use a tool, reply in this format:\n"
+            "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"
-            "Never explain anything."
+            "Do not explain anything."
         )
-
-        full_prompt = f"{prompt}\n\nQuestion: {query}\nAnswer:"
-        
-
-        chat_input = [
+        messages = [
             {"role": "system", "content": prompt},
-            {"role": "user", "content": query}
+            {"role": "user", "content": query},
         ]
-        response = llm.invoke(chat_input).strip()
-        
-        # response = llm.invoke(full_prompt).strip()
-        print(f"🧠 LLM said: {response}")
+        response = llm.invoke(messages).strip()
+        print(f"🧠 LLM responded: {response}")
 
-        if response.startswith("use_tool:"):
+        if response.lower().startswith("use_tool:"):
             return {
                 "messages": state["messages"] + [AIMessage(content=response)],
-                "tool_call": response  # carry tool signal
+                "tool_call": response,
             }
-        else:
-            return {"messages": [AIMessage(content=response)]}
+        return {"messages": [AIMessage(content=response)]}
 
-    # Node 3: Tool execution
+    # Node 3: Tool Executor
     def tool_node(state: MessagesState):
         try:
             tool_signal = state.get("tool_call", "")
             _, tool_name, tool_input = tool_signal.split(":", 2)
-            tool_name = tool_name.strip()
-            tool_input = tool_input.strip()
-            tool_fn = TOOLS.get(tool_name)
-
+            tool_fn = TOOLS.get(tool_name.strip())
             if not tool_fn:
                 print(f"❌ Unknown tool: {tool_name}")
                 return {"messages": [AIMessage(content="Unknown")]}
-            
-            print(f"🔧 Using tool: {tool_name} with input: {tool_input}")
-            tool_result = tool_fn(tool_input)
-            return {"messages": [AIMessage(content=str(tool_result))]}
-
+            print(f"🔧 Running tool {tool_name} with input: {tool_input}")
+            result = tool_fn(tool_input.strip())
+            return {"messages": [AIMessage(content=str(result))]}
         except Exception as e:
             print(f"⚠️ Tool error: {e}")
-            return {"messages": [AIMessage(content="Unknown")]}  # fail-safe
+            return {"messages": [AIMessage(content="Unknown")]}
 
     # Build LangGraph
     builder = StateGraph(MessagesState)
@@ -98,15 +85,14 @@ def build_graph():
 
     return builder.compile()
 
-# === BasicAgent wrapper ===
+
+# --- Final Agent Wrapper ---
 class BasicAgent:
     def __init__(self):
-        print("BasicAgent initialized with retriever + LLM + manual tool logic.")
+        print("✅ BasicAgent initialized with retriever + LLM + tools")
         self.graph = build_graph()
 
     def __call__(self, question: str) -> str:
         print(f"📥 Question: {question[:100]}")
         result = self.graph.invoke({"messages": [HumanMessage(content=question)]})
-        answer = result["messages"][-1].content.strip()
-        print(f"📤 Answer: {answer}")
-        return answer
+        return result["messages"][-1].content.strip()