fix graph flow

agent/graph.py

@@ -26,8 +26,7 @@ def create_graph_builder():
     # Set entry point
     builder.set_entry_point("agent")
 
-    builder.add_edge("agent", "write_memory")
-    builder.add_edge("write_memory", END)
+    # builder.add_edge("agent", "write_memory")
 
     # Add conditional edges from agent
     builder.add_conditional_edges(
@@ -36,12 +35,12 @@ def create_graph_builder():
         {
             "action": "action",
             "write_memory": "write_memory",
-            END: END
         }
     )
     
     # Connect action back to agent
     builder.add_edge("action", "agent")
+    builder.add_edge("write_memory", END)
     
     return builder
 

agent/utils/nodes.py

@@ -80,8 +80,8 @@ def call_model(state: AgentState, config: RunnableConfig, store: BaseStore):
     messages = [
         SystemMessage(content=SYSTEM_PROMPT.format(memory=memory_content))
     ] + state["messages"]
-    
-    response = model.invoke(messages)
+    tool_calling_model = model.bind_tools(tool_belt)
+    response = tool_calling_model.invoke(messages)
     return {"messages": [response]}
 
 def update_memory(state: AgentState, config: RunnableConfig, store: BaseStore):
@@ -100,7 +100,7 @@ def update_memory(state: AgentState, config: RunnableConfig, store: BaseStore):
     store.put(namespace, "user_memory", {"memory": new_memory.content})
     return state
 
-def should_continue(state: AgentState) -> Literal["action", "write_memory", END]:
+def should_continue(state: AgentState) -> Literal["action", "write_memory"]:
     """Determine the next node in the graph."""
     if not state["messages"]:
         return END
@@ -114,15 +114,17 @@ def should_continue(state: AgentState) -> Literal["action", "write_memory", END]
     # Handle tool calls
     if hasattr(last_message, "additional_kwargs") and last_message.additional_kwargs.get("tool_calls"):
         return "action"
+
+    return "write_memory"
     
-    # Handle memory operations for human messages
-    if last_human_message:
-            
-        # Write memory for longer messages that might contain personal information
-        if len(last_human_message.content.split()) > 3:
-            return "write_memory"
+    # # Handle memory operations for human messages
+    # if last_human_message:
+           
+    #     # Write memory for longer messages that might contain personal information
+    #     if len(last_human_message.content.split()) > 3:
+    #         return "write_memory"
             
-    return END
+    # return END
 
 # Define the memory creation prompt
 MEMORY_CREATION_PROMPT = """"You are collecting information about the user to personalize your responses.