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, tools_condition
+from langgraph.prebuilt import ToolNode
 from langchain_huggingface import HuggingFaceEndpoint
-from tools import TOOLS  
+from tools import TOOLS  # your dictionary of tool functions
 
 # --- Load local QA metadata for retriever ---
 QA_PATH = "metadata.jsonl"
@@ -14,74 +14,90 @@ qa_dict = {
     for _, row in qa_pairs.iterrows()
 }
 
-from langchain_huggingface import HuggingFaceEndpoint
-
+# === LangGraph builder ===
 def build_graph():
     llm = HuggingFaceEndpoint(
-        # 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
+    # Node 1: Retriever
     def retriever_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
         if query in qa_dict:
             print("✅ Exact match found in retriever.")
             return {"messages": [AIMessage(content=qa_dict[query])]}
-        print("🔍 No match found. Falling back to LLM.")
+        print("🔍 No match. Sending to LLM.")
         return {"messages": state["messages"]}
 
-    # Node: assistant
+    # Node 2: Assistant (LLM response parsing)
     def assistant_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
 
-        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"
+        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"
+            "use_tool: <tool_name>: <tool_input>\n"
+            "Never explain anything."
         )
 
-        # LangChain expects list of messages for tool-call-capable models
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": query},
-        ]
+        full_prompt = f"{prompt}\n\nQuestion: {query}\nAnswer:"
+        response = llm.invoke(full_prompt).strip()
+        print(f"🧠 LLM said: {response}")
+
+        if response.startswith("use_tool:"):
+            return {
+                "messages": state["messages"] + [AIMessage(content=response)],
+                "tool_call": response  # carry tool signal
+            }
+        else:
+            return {"messages": [AIMessage(content=response)]}
 
-        response = llm_with_tools.invoke(messages)
-        return {"messages": [AIMessage(content=response.strip())]}
+    # Node 3: Tool execution
+    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)
 
-    # === Build LangGraph ===
+            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))]}
+
+        except Exception as e:
+            print(f"⚠️ Tool error: {e}")
+            return {"messages": [AIMessage(content="Unknown")]}  # fail-safe
+
+    # 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
+    builder.add_node("tool", tool_node)
 
-    # --- Edges ---
     builder.set_entry_point("retriever")
     builder.add_edge("retriever", "assistant")
-    builder.add_conditional_edges("assistant", tools_condition)
-    builder.add_edge("tools", "assistant")
+    builder.add_edge("assistant", "tool")
+    builder.add_edge("tool", "assistant")
     builder.set_finish_point("assistant")
 
     return builder.compile()
 
-
-# --- Agent class wrapper for app.py ---
+# === BasicAgent wrapper ===
 class BasicAgent:
     def __init__(self):
-        print("BasicAgent initialized with retriever + fallback LLM.")
+        print("BasicAgent initialized with retriever + LLM + manual tool logic.")
         self.graph = build_graph()
 
     def __call__(self, question: str) -> str:
-        print(f"📥 Received question: {question[:80]}")
+        print(f"📥 Question: {question[:100]}")
         result = self.graph.invoke({"messages": [HumanMessage(content=question)]})
         answer = result["messages"][-1].content.strip()
         print(f"📤 Answer: {answer}")