agent

app.py

@@ -15,7 +15,7 @@ from langchain_groq import ChatGroq
 from langchain_core.messages import HumanMessage
 
 from langgraph.graph import START, END, StateGraph
-from langgraph.prebuilt import ToolNode
+from langgraph.prebuilt import ToolNode, tools_condition
 
 from agent_utils import wiki_ret
 
@@ -38,6 +38,10 @@ class BasicAgent:
         print("BasicAgent initialized.")
         self.graph = self.create_graph()
 
+    def __call___(self, question: str) -> str:
+        response = (self.graph).invoke({"messages": question})
+        return response
+
     def agent(self, question: str) -> str:
 
         # print(f"Agent received question (first 50 chars): {question[:50]}...")
@@ -78,13 +82,10 @@ class BasicAgent:
         builder.add_node("tools", ToolNode(tools = [wiki_ret]))
 
         builder.add_edge(START, "agent")
+        builder.add_conditional_edges("agent", tools_condition)
         builder.add_edge("agent", END)
         return builder.compile()
 
-    def __call___(self, question: str) -> str:
-        response = (self.graph).invoke({"messages": question})
-        return response
-
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """