changes in structure

app/graph.py

@@ -22,7 +22,7 @@ from app.database.connection import pool
 from app.utils.embeddings import remote_embeddings
 from langgraph.checkpoint.postgres import PostgresSaver
 from langgraph.store.postgres import PostgresStore
-from app.tools.email_writing_agent_tools import email_writing_agent_tools
+
 
 # Define a standard retry policy for database-heavy nodes
 db_retry_policy = RetryPolicy(
@@ -129,6 +129,7 @@ builder.add_conditional_edges(
 )
 
 
+
 builder.add_edge("store_memory_and_data_node", END)
 builder.add_edge("unsafe_emails_node", END)
 builder.add_edge("archive_node", END)

app/nodes/email_writing_node.py

@@ -40,23 +40,10 @@ def email_writing_agent_node(state: EmailAgentState) -> dict:
 
 
 
-def route_after_tools(state: EmailAgentState):
-    # Iterate backwards to find the latest ToolMessage
-    # This handles cases where the LLM might have sent a text follow-up
-    last_tool_msg = next((m for m in reversed(state["messages"]) 
-                         if isinstance(m, ToolMessage)), None)
-
-    if not last_tool_msg:
-        return "email_writing_agent"
-
-    content_upper = last_tool_msg.content.upper()
-    
-    # logic 1: If we just successfully SENT the email, go to Parser -> Memory
-    if last_tool_msg.name == "send_draft_by_id" and "SUCCESS" in content_upper:
-        print("--- ROUTER: Send successful. Moving to Parse/Memory. ---")
+def route_after_tools(state):
+
+    if state.get("sent_message_id"):
+        print("--- ROUTER: Send successful. Moving to Memory. ---")
         return "store_memory_and_data_node"
-    
-    # logic 2: If we just created a DRAFT (or the send failed)
-    # Go back to agent to talk to the user
-    print("--- ROUTER: Draft created or Tool failed. Returning to Agent. ---")
+
     return "email_writing_agent"

app/tools/email_writing_agent_tools.py

@@ -1,3 +1,4 @@
+from os import name
 from langgraph.types import interrupt
 from googleapiclient.errors import HttpError
 from app.schemas.email_writing_agent_tools_schema import CreateDraftSchema, SendDraftSchema
@@ -9,8 +10,8 @@ from langchain.tools import ToolRuntime
 from langgraph.types import Command
 from langgraph.graph import END
 from langchain_core.messages import ToolMessage,SystemMessage, HumanMessage
-
 from app.gmail_auth import gmail_toolkit ,api_resource
+
 @tool(args_schema=CreateDraftSchema)
 def create_gmail_draft(
     to: Union[str, list], 
@@ -49,7 +50,8 @@ def create_gmail_draft(
                     "draft_id": draft_id, 
                     "reply_subject": subject,
                     "reply_email_body": body,
-                    "messages": [ToolMessage(content, tool_call_id=tool_call_id)]
+                    "messages": [ToolMessage(content, tool_call_id=tool_call_id)],
+                    "name": "create_gmail_draft"
                 },goto=END
             )
         except IndexError:
@@ -81,13 +83,14 @@ def send_draft(
         return Command(
             update={
                 "sent_message_id": sent_id,
-                "messages": [ToolMessage(content, tool_call_id=tool_call_id)]
+                "messages": [ToolMessage(content, tool_call_id=tool_call_id)],
+                "name": "send_draft"
             }
         )
     except HttpError as error:
         error_msg = f"ERROR: {error}"
         return Command(
-            update={"messages": [ToolMessage(error_msg, tool_call_id=tool_call_id)]}
+            update={"messages": [ToolMessage(error_msg, tool_call_id=tool_call_id)], "name": "send_draft"}
         )