Update agent.py

agent.py

@@ -4,7 +4,7 @@ import os
 from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import tools_condition
-#from langgraph.prebuilt import ToolNode
+from langgraph.prebuilt import ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
@@ -839,6 +839,37 @@ def process_question(question: str):
 
 
 
+def call_llm(state):
+    messages = state["messages"]
+    response = llm.invoke(messages)
+    return {"messages": messages + [response]}
+
+tool_node = ToolNode([search])
+
+
+builder = StateGraph()
+
+builder.add_node("call_llm", call_llm)
+builder.add_node("call_tool", tool_node)
+
+# Decide what to do next: if tool call → call_tool, else → end
+def should_call_tool(state):
+    last_msg = state["messages"][-1]
+    if isinstance(last_msg, AIMessage) and last_msg.tool_calls:
+        return "call_tool"
+    return "end"
+
+builder.set_entry_point("call_llm")
+builder.add_conditional_edges("call_llm", should_call_tool, {
+    "call_tool": "call_tool",
+    "end": None
+})
+
+# After tool runs, go back to the LLM
+builder.add_edge("call_tool", "call_llm")
+
+
+
 
 # To store previously asked questions and timestamps (simulating state persistence)
 recent_questions = {}