Update agent.py

agent.py

@@ -1,11 +1,10 @@
 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 tools import TOOLS  
+from langgraph.graph import StateGraph, END, MessagesState
+from langchain_huggingface import HuggingFaceEndpoint
+from tools import TOOLS
 
-# --- Read local QA data for retriever ---
 QA_PATH = "metadata.jsonl"
 qa_pairs = pd.read_json(QA_PATH, lines=True)
 qa_dict = {
@@ -14,7 +13,6 @@ qa_dict = {
 }
 
 def build_graph():
-    # Initialize Mistral model with proper configuration
     llm = HuggingFaceEndpoint(
         endpoint_url="https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3",
         task="text-generation",
@@ -25,50 +23,43 @@ def build_graph():
         huggingfacehub_api_token=os.environ["HF_TOKEN"]
     )
 
-    # Retriever node
     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. Sending to LLM.")
-        return {"messages": state["messages"]}
+            answer = qa_dict[query]
+            return {"messages": [AIMessage(content=answer)]}
+        print("🔍 No match. Passing to LLM.")
+        return 
 
-    # Assistant node (LLM)
     def assistant_node(state: MessagesState):
         query = state["messages"][-1].content.strip()
 
-        # Format system prompt
         system_prompt = (
-            "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"
+            "You are a helpful assistant. To answer the user's question, you can use tools. "
+            "To use a tool, respond with a single line: 'tool:tool_name:input'. "
+            "For example: 'tool:wiki_search:Apple Inc.' "
+            "If you have the final answer, provide it directly without any prefixes. "
+            "Never justify or explain your final answer."
         )
         
-        # Format prompt for Mistral model
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": query}
         ]
         
-        # Generate response
         response = llm.invoke(messages).strip()
 
-        # Clean up response
-        for tag in ("Final answer:", "Answer:", "assistant:", "assistant:"):
+        for tag in ("", "Answer:", "assistant:"):
             if response.lower().startswith(tag.lower()):
                 response = response[len(tag):].strip()
         
-        return {"messages": [AIMessage(content=response.strip())]}
+        return {"messages": [AIMessage(content=response)]}
 
-    # Tool node
     def tool_node(state: MessagesState):
+        last_message = state["messages"][-1].content.strip()
         try:
-            tool_signal = state.get("tool_call", "")
-            _, tool_name, tool_input = tool_signal.split(":", 2)
+            _, tool_name, tool_input = last_message.split(":", 2)
             tool_name = tool_name.strip()
             tool_input = tool_input.strip()
             tool_fn = TOOLS.get(tool_name)
@@ -83,23 +74,40 @@ def build_graph():
 
         except Exception as e:
             print(f"⚠️ Tool error: {e}")
-            return {"messages": [AIMessage(content="Unknown")]}  # fail-safe
+            return {"messages": [AIMessage(content="Unknown")]}
+
+    def route_after_retriever(state: MessagesState):
+        if isinstance(state["messages"][-1], AIMessage):
+            return END
+        return "assistant"
+
+    def route_after_assistant(state: MessagesState):
+        last_message = state["messages"][-1].content.strip().lower()
+        if last_message.startswith("tool:"):
+            return "tool"
+        return END
 
-    # Build LangGraph
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever_node)
     builder.add_node("assistant", assistant_node)
     builder.add_node("tool", tool_node)
 
     builder.set_entry_point("retriever")
-    builder.add_edge("retriever", "assistant")
-    builder.add_edge("assistant", "tool")
+
+    builder.add_conditional_edges(
+        "retriever",
+        route_after_retriever,
+        {"assistant": "assistant", END: END}
+    )
+    builder.add_conditional_edges(
+        "assistant",
+        route_after_assistant,
+        {"tool": "tool", END: END}
+    )
     builder.add_edge("tool", "assistant")
-    builder.set_finish_point("assistant")
 
     return builder.compile()
 
-# Agent class
 class BasicAgent:
     def __init__(self):
         print("✅ BasicAgent initialized with retriever + LLM + tools")
@@ -107,7 +115,15 @@ class BasicAgent:
 
     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
+        config = {"recursion_limit": 50}
+        try:
+            result = self.graph.invoke(
+                {"messages": [HumanMessage(content=question)]},
+                config=config
+            )
+            answer = result["messages"][-1].content.strip()
+            print(f"📤 Answer: {answer}")
+            return answer
+        except Exception as e:
+            print(f"Error during graph invocation: {e}")
+            return f"AGENT ERROR: {e}"