Update agent.py

agent.py

@@ -2,7 +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, tools_condition
 from langchain_huggingface import HuggingFaceEndpoint
+from tools import TOOLS  
 
 # --- Load local QA metadata for retriever ---
 QA_PATH = "metadata.jsonl"
@@ -16,10 +18,14 @@ from langchain_huggingface import HuggingFaceEndpoint
 
 def build_graph():
     llm = HuggingFaceEndpoint(
-        repo_id="Qwen/Qwen2.5-32B-Instruct",
+        # repo_id="Qwen/Qwen2.5-32B-Instruct",
+        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
         task="text-generation",
         huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
+    
+    llm_with_tools = llm.bind_tools(TOOLS)
+    
 
     # Node: retriever
     def retriever_node(state: MessagesState):
@@ -30,45 +36,39 @@ def build_graph():
         print("🔍 No match found. Falling back to LLM.")
         return {"messages": state["messages"]}
 
-    # Node: assistant (LLM with fallback prompt logic)
+    # Node: assistant
     def assistant_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
 
         system_prompt = (
-            "You are a helpful AI assistant taking part in the GAIA evaluation benchmark.\n"
-            "You must return only the final answer to the user's question.\n"
-            "- No explanations.\n"
-            "- No formatting like 'Final answer:' or similar.\n"
-            "- If the answer is a list, return comma-separated values.\n"
-            "- If the answer is unknown, return 'Unknown'."
+            "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"
         )
 
-        try:
-            chat_input = [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": query}
-            ]
-            response = llm.invoke(chat_input)
-            print("✅ Used chat-style prompt.")
-        except Exception as e:
-            print(f"⚠️ Chat-style failed: {e}")
-            fallback_prompt = (
-                f"{system_prompt}\n\n"
-                f"Question: {query}\n"
-                f"Answer:"
-            )
-            response = llm.invoke(fallback_prompt)
-            print("🔁 Used fallback prompt format.")
+        # LangChain expects list of messages for tool-call-capable models
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": query},
+        ]
 
+        response = llm_with_tools.invoke(messages)
         return {"messages": [AIMessage(content=response.strip())]}
 
-    # Build the LangGraph
+    # === Build LangGraph ===
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever_node)
     builder.add_node("assistant", assistant_node)
+    builder.add_node("tools", ToolNode(TOOLS))  # LangGraph's tool executor
 
+    # --- Edges ---
     builder.set_entry_point("retriever")
     builder.add_edge("retriever", "assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
     builder.set_finish_point("assistant")
 
     return builder.compile()