Update agent.py

agent.py

@@ -188,13 +188,12 @@ def build_graph(provider: str = "openai"):
     llm_with_tools = llm.bind_tools(tools)
 
     # Define nodes
-    def assistant(state: MessagesState):
+     def assistant(state: MessagesState):
         """Assistant node - generates responses"""
-        # Get current messages
         messages = state["messages"]
         # Generate response using LLM
         response = llm_with_tools.invoke(messages)
-        # Append new message to state
+        # Return new state with appended message
         return {"messages": messages + [response]}
     
     def retriever(state: MessagesState):
@@ -230,13 +229,23 @@ def build_graph(provider: str = "openai"):
     # Set up edges
     builder.set_entry_point("retriever")
     builder.add_edge("retriever", "assistant")
+    
+    # Conditional tool usage
     builder.add_conditional_edges(
         "assistant",
         tools_condition,
-        {"continue": "tools", "end": END}
+        {
+            # Continue to tools if needed
+            "continue": "tools",
+            # End conversation if no tools needed
+            "end": END
+        }
     )
-    builder.add_edge("tools", "retriever")  # Go back to retriever after tools
     
+    # After tools execute, return to assistant for response generation
+    builder.add_edge("tools", "assistant")
+    
+    builder.add_finish_point(END)  # Explicitly declare END as finish point
     return builder.compile()