Update app.py

app.py

@@ -6,7 +6,7 @@ from langgraph.graph.message import add_messages
 from langchain_openai import ChatOpenAI
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_core.messages import HumanMessage, ToolMessage, AIMessage
-from langgraph.prebuilt import ToolNode, tools_condition
+from langgraph.prebuilt import tools_condition
 import os
 
 # Streamlit UI Header
@@ -31,20 +31,18 @@ class State(TypedDict):
 # Initialize LLM and Tools
 llm = ChatOpenAI(model="gpt-4o-mini")
 tool = TavilySearchResults(max_results=2)
-tools = [tool]
-llm_with_tools = llm.bind_tools(tools)
+llm_with_tools = llm.bind_tools([tool])
 
-# Define Agent Function
+# Agent Node
 def Agent(state: State):
     st.sidebar.write("Agent Input State:", state["messages"])
     response = llm_with_tools.invoke(state["messages"])
     st.sidebar.write("Agent Response:", response)
     return {"messages": [response]}
 
-# Define Tool Execution Node
+# Tools Execution Node
 def ExecuteTools(state: State):
-    last_message = state["messages"][-1]
-    tool_calls = last_message.tool_calls
+    tool_calls = state["messages"][-1].tool_calls
     responses = []
 
     if tool_calls:
@@ -54,9 +52,9 @@ def ExecuteTools(state: State):
             st.sidebar.write("Tool Call Detected:", tool_name, args)
 
             if tool_name == "tavily_search_results_json":
-                response = tool.invoke(args)
-                st.sidebar.write("Tool Response:", response)
-                responses.append(ToolMessage(content=str(response), tool_call_id=call["id"]))
+                tool_response = tool.invoke({"query": args["query"]})
+                st.sidebar.write("Tool Response:", tool_response)
+                responses.append(ToolMessage(content=str(tool_response), tool_call_id=call["id"]))
     return {"messages": responses}
 
 # Memory Checkpoint
@@ -67,7 +65,11 @@ graph = StateGraph(State)
 graph.add_node("Agent", Agent)
 graph.add_node("ExecuteTools", ExecuteTools)
 
-graph.add_conditional_edges("Agent", tools_condition, {"True": "ExecuteTools"})
+# Add Conditional Edge to Check for Tools
+def custom_tools_condition(state: State):
+    return "True" if state["messages"][-1].tool_calls else "False"
+
+graph.add_conditional_edges("Agent", custom_tools_condition, {"True": "ExecuteTools", "False": "Agent"})
 graph.add_edge("ExecuteTools", "Agent")
 graph.set_entry_point("Agent")