made changes in tools and graph

app/agents/email_writing_agent.py

@@ -1,8 +1,8 @@
 from langchain_groq import ChatGroq
-from app.tools.email_writing_agent_tools import create_gmail_draft, send_draft_by_id
+from app.tools.email_writing_agent_tools import create_gmail_draft, send_draft
 base_llm = ChatGroq(
     model="qwen/qwen3-32b",
     temperature=0.1,       
 )
-tools = [create_gmail_draft, send_draft_by_id]
+tools = [create_gmail_draft, send_draft]
 email_agent = base_llm.bind_tools(tools)

app/database/connection.py

@@ -2,6 +2,7 @@ from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker, Session
 from app.core.config import settings
 from psycopg_pool import ConnectionPool
+from contextlib import contextmanager
 
 
 
@@ -26,6 +27,7 @@ engine = create_engine(
 
 SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
 
+
 def get_session():
     db = SessionLocal()
     try:

app/graph.py

@@ -4,7 +4,7 @@ from langgraph.prebuilt import ToolNode ,tools_condition
 from app.nodes.archive_node import  archive_node
 from app.nodes.email_writing_node import *
 from langgraph.graph import StateGraph,END,START
-from app.tools.email_writing_agent_tools import create_gmail_draft, send_draft_by_id
+from app.tools.email_writing_agent_tools import create_gmail_draft, send_draft
 from app.nodes.safety_check_node import *
 from app.nodes.parse_node import parse_response_node
 from app.nodes.context_node import prepare_context_node

app/main.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI, HTTPException, Depends
 from pydantic import BaseModel, EmailStr
 from typing import Optional, Dict, Any, TypedDict, Annotated, Sequence
-from langchain_core.messages import BaseMessage
+from langchain_core.messages import BaseMessage, HumanMessage
 from langgraph.graph import add_messages
 from langgraph.types import Command
 import uuid
@@ -11,9 +11,18 @@ from app.state.state import EmailAgentState
 from app.database.connection import get_session
 from app.database.utils import get_or_create_user
 from sqlalchemy.orm import Session
+from app.database.connection import SessionLocal
+
 
 logger = logging.getLogger(__name__)
 
+def get_session():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
 app = FastAPI(title="AI Email Agent API")
 
 # --- Schemas ---
@@ -32,7 +41,10 @@ class ReviewActionRequest(BaseModel):
     status: str  # "approved" or "rejected"
     feedback: Optional[str] = None
 
-
+class SendEmailRequest(BaseModel):
+    thread_id: str
+    user_id: str
+    human_message: str
 # --- Helper Functions ---
 
 def parse_interrupt(final_state: Dict[str, Any]) -> Optional[Dict[str, Any]]:
@@ -90,6 +102,7 @@ def process_email(request: EmailProcessRequest, db: Session = Depends(get_sessio
                     return {
                         "status": "needs_review",
                         "thread_id": thread_id,
+                        "messages": final_state.get("messages", []),
                         "triage_label": final_state.get("triage_label"),
                         "action": parsed_interrupt["action"],
                         "email_draft": {
@@ -133,17 +146,17 @@ def review_action(request: ReviewActionRequest) -> Dict[str, Any]:
         else:
             raise HTTPException(status_code=400, detail="Invalid status")
 
-        final_state = graph.invoke(payload, config=config)
+        intermediate_state = graph.invoke(payload, config=config)
         
         # Still in review phase
-        if "__interrupt__" in final_state and not final_state.get("draft_id"):
-            parsed_interrupt = parse_interrupt(final_state)
+        if "__interrupt__" in intermediate_state and not intermediate_state.get("draft_id"):
+            parsed_interrupt = parse_interrupt(intermediate_state)
             if parsed_interrupt:
                 data = parsed_interrupt["data"]
                 return {
                     "status": "needs_review",
                     "thread_id": request.thread_id,
-                    "triage_label": final_state.get("triage_label"),
+                    "triage_label": intermediate_state.get("triage_label"),
                     "action": parsed_interrupt["action"],
                     "email_draft": {
                         "to": data.get("to"),
@@ -153,12 +166,13 @@ def review_action(request: ReviewActionRequest) -> Dict[str, Any]:
                 }
         
         # Draft created, review complete
-        if final_state.get("draft_id"):
+        if intermediate_state.get("draft_id"):
             return {
                 "thread_id": request.thread_id,
-                "draft_id": final_state["draft_id"],
-                "reply_subject": final_state.get("reply_subject"),
-                "reply_email_body": final_state.get("reply_email_body"),
+                "draft_id": intermediate_state["draft_id"],
+                "messages": intermediate_state.get("messages", []),
+                "reply_subject": intermediate_state.get("reply_subject"),
+                "reply_email_body": intermediate_state.get("reply_email_body"),
             }
 
     except Exception as e:
@@ -167,8 +181,33 @@ def review_action(request: ReviewActionRequest) -> Dict[str, Any]:
 
 
 
+@app.post("/send_email")
+def send_email(request: SendEmailRequest) -> Dict[str, Any]:
+   
+        config = {
+            "configurable": {
+                "thread_id": request.thread_id,
+                "user_id": request.user_id
+            }
+        }
+
+        graph.update_state(
+          config,
+          {"messages": [HumanMessage(content=request.human_message)]},
+       as_node="prepare_context_node" 
+     )  
+        final_state = graph.invoke(None, config=config)
+
+        return {
+            "thread_id": request.thread_id,
+            "messages": final_state.get("messages", []),
+            "sent_message_id": final_state.get("sent_message_id")
+        }
+
+
+
 
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="127.0.0.1", port=8000)

app/nodes/email_writing_node.py

@@ -54,7 +54,7 @@ def route_after_tools(state: EmailAgentState):
     # 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. ---")
-        return "parse_node"
+        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

app/prompts/email_writing_agent_prompt.py

@@ -16,16 +16,16 @@ User: "Reply to client."
 Agent Tool: create_gmail_draft(...)
 Output: "Success! ID: draft_999"
 User: "Send it."
-Agent Tool: send_draft_by_id(draft_id="draft_999")
-Output: "SUCCESS: Sent! Message ID: msg_123"
+Agent Tool: send_draft()
+Output: "SUCCESS: Sent! Message"
 Output: "Stored."
 </one_shot_example>
 
 <rules>
 1. NEW DRAFT: Call `create_gmail_draft` first.
 2. REJECTION: If "DRAFT REJECTED", rewrite and call `create_gmail_draft_tool` again.
-3. SENDING: Only call `send_draft_by_id` when user explicitly orders to send. 
-4. ARCHIVING: After `send_draft_by_id` returns a Message ID.
+3. SENDING: Only call `send_draft` when user explicitly orders to send. 
+4. ARCHIVING: After `send_draft` returns a Message ID.
 important:You are not allowed to send until the user explicitly orders to send.
 5.Use **{sender_email_id}** as the recipient not the name.
 **dont call send_draft_by_id until user  allows you to send the draft**

app/tools/email_writing_agent_tools.py

@@ -5,8 +5,10 @@ from langchain.tools import tool
 from langchain_google_community import GmailToolkit
 from typing import Annotated, Union
 from langchain_core.tools import InjectedToolCallId, tool
+from langchain.tools import ToolRuntime
 from langgraph.types import Command
 from langchain_core.messages import SystemMessage, HumanMessage,ToolMessage,AIMessage,BaseMessage
+from langgraph.graph import END
 
 @tool(args_schema=CreateDraftSchema)
 def create_gmail_draft(
@@ -37,7 +39,8 @@ def create_gmail_draft(
         reply = draft_tool.invoke({"message": body, "to": [to], "subject": subject})
         try:
             draft_id = reply.split(":")[1].strip()
-            content = f"Successfully created draft: <id>{draft_id}</id> <subject>{subject}</subject> <body>{body}</body>"
+
+            content = f"Draft created. User You MUST confirm before calling send_draft()."
             
             # UPDATE STATE: Save draft_id directly
             return Command(
@@ -46,7 +49,7 @@ def create_gmail_draft(
                     "reply_subject": subject,
                     "reply_email_body": body,
                     "messages": [ToolMessage(content, tool_call_id=tool_call_id)]
-                }
+                },goto=END
             )
         except IndexError:
             return f"Draft created, but response parsing failed: {reply}"
@@ -57,20 +60,20 @@ def create_gmail_draft(
 
 #---------------------------------------------------------------------------
 
-@tool(args_schema=SendDraftSchema)
-def send_draft_by_id(
-    draft_id: str,
-    tool_call_id: Annotated[str, InjectedToolCallId] # Injected ID
+
+def send_draft(
+    tool_call_id: Annotated[str, InjectedToolCallId],runtime: ToolRuntime # Injected ID
 ):
     """Sends a finalized Gmail draft by its ID."""
+
     try:
         toolkit = GmailToolkit()
         result = toolkit.api_resource.users().drafts().send(
-            userId="me", body={"id": draft_id}
+            userId="me", body={"id": runtime.state["draft_id"]}
         ).execute()
         
         sent_id = result['id']
-        content = f"SUCCESS: Sent! a Gmail with ID: <id>{sent_id}</id>"
+        content = f"SUCCESS: Sent! a Gmail"
         
         # UPDATE STATE: Save sent_message_id directly
         return Command(
@@ -89,5 +92,5 @@ def send_draft_by_id(
 
 email_writing_agent_tools = [
     create_gmail_draft,
-    send_draft_by_id
+    send_draft
 ]