chnages

app/database/connection.py

@@ -26,7 +26,11 @@ engine = create_engine(
 
 SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
 
-def get_session() -> Session:
-    return SessionLocal()
+def get_session():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
 
 

app/graph.py

@@ -41,6 +41,9 @@ tool_node_retry_policy = RetryPolicy(
 
 builder = StateGraph(EmailAgentState)
 
+# Nodes
+builder = StateGraph(EmailAgentState)
+
 # Nodes
 builder.add_node("safety_check_node", safety_classifier_node)  
 builder.add_node("check_previous_email_exist_node", check_previous_email_exist_node)
@@ -63,9 +66,10 @@ builder.add_node(
 )
 
 builder.add_node("archive_node", archive_node,retry=db_retry_policy)   
-builder.add_node("parse_node", parse_response_node)
-builder.add_node("tools", ToolNode(email_writing_agent_tools), retry_policy=tool_node_retry_policy)
 
+builder.add_node("tools", ToolNode(tools), retry=tool_node_retry_policy)
+
+# Edges (Same as your original logic)
 builder.add_edge(START, "safety_check_node")
 
 builder.add_conditional_edges(
@@ -119,24 +123,16 @@ builder.add_conditional_edges(
     "tools",
     route_after_tools,
     {
-        "parse_node": "parse_node",
+        "store_memory_and_data_node": "store_memory_and_data_node",
         "email_writing_agent": "email_writing_agent"
     }
 )
 
-builder.add_edge("parse_node", "store_memory_and_data_node")
+
 builder.add_edge("store_memory_and_data_node", END)
 builder.add_edge("unsafe_emails_node", END)
 builder.add_edge("archive_node", END)
 
-
-
-
-
-
-
-
-
 toolkit = GmailToolkit()
 
 

app/main.py

@@ -1,48 +1,171 @@
-import time
-from psycopg import OperationalError
+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 langgraph.graph import add_messages
+from langgraph.types import Command
+import uuid
+import logging
 from app.graph import graph
 from app.state.state import EmailAgentState
-from app.database.connection import pool
-import uuid
+from app.database.connection import get_session
+from app.database.utils import get_or_create_user
+from sqlalchemy.orm import Session
 
+logger = logging.getLogger(__name__)
 
-config = {"configurable": {"thread_id": str(uuid.uuid4()), "user_id": "1"}}
+app = FastAPI(title="AI Email Agent API")
 
-input_data: EmailAgentState = {
-    "user_email_id": "gaykaratharva7@gmail.com",
-    "user_id": 1,
-    "user_name": "Atharva",
-    "sender_email_id": "atharvagaykar36@gmail.com",
-    "sender_subject": "URGENT: Validation of Hybrid Phishing Detection Model & XGBoost Integration",
-    "sender_email_body": """Dear Atharva,\r\n\r\nI have completed the integration of the AI-Driven Email Threat Detection pipeline... [truncated for brevity]"""
-}
+# --- Schemas ---
 
-if __name__ == "__main__":
-    result = None
+class EmailProcessRequest(BaseModel):
+    thread_id: str
+    user_email: EmailStr
+    sender_email_id: EmailStr
+    sender_subject: str
+    sender_email_body: str
+
+
+class ReviewActionRequest(BaseModel):
+    thread_id: str
+    user_id: str
+    status: str  # "approved" or "rejected"
+    feedback: Optional[str] = None
+
+
+# --- Helper Functions ---
+
+def parse_interrupt(final_state: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """Parse interrupt from graph state."""
+    if "__interrupt__" not in final_state:
+        return None
+
+    interrupt_state = final_state.get("__interrupt__")
+    if not interrupt_state:
+        return None
+
+    interrupt = interrupt_state[0]
+    value = getattr(interrupt, "value", {}) or {}
+
+    return {
+        "action": value.get("action"),
+        "data": value.get("data", {})
+    }
+
+
+# --- Endpoints ---
+
+@app.post("/process-email")
+def process_email(request: EmailProcessRequest, db: Session = Depends(get_session)) -> Dict[str, Any]:
+    """Process email through the graph pipeline."""
+    
     try:
-        # 3. Retry loop for Neon wake-up/stability
-        for i in range(3): 
-            try:
-                print(f"Attempt {i+1}: Invoking graph...")
-                result = graph.invoke(input_data, config=config)
-                break  # Success! Exit the loop.
-            except OperationalError as e:
-                if i < 2:
-                    print("Waiting for Neon database to wake up...")
-                    time.sleep(10) # Increased sleep slightly for Neon cold starts
-                else:
-                    print("Max retries reached. Database connection failed.")
-                    raise e
-
-        # 4. Output the result
-        if result:
-            print("\n--- Graph Execution Result ---")
-            print(result)
+        user = get_or_create_user(db, request.user_email)
+        
+        thread_id = request.thread_id
+        config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "user_id": str(user.id)
+            }
+        }
+        
+        input_data = {
+            "user_email_id": request.user_email,
+            "user_id": user.id,
+            "user_name": "Atharva",
+            "sender_email_id": request.sender_email_id,
+            "sender_subject": request.sender_subject,
+            "sender_email_body": request.sender_email_body,
+        }
+
+        final_state = graph.invoke(input_data, config=config)
+
+        if final_state.get('triage_label') == "FOLLOW_UP_REQUIRED":
+            if "__interrupt__" in final_state and not final_state.get("draft_id"):
+                parsed_interrupt = parse_interrupt(final_state)
+                if parsed_interrupt:
+                    data = parsed_interrupt["data"]
+
+                    return {
+                        "status": "needs_review",
+                        "thread_id": thread_id,
+                        "triage_label": final_state.get("triage_label"),
+                        "action": parsed_interrupt["action"],
+                        "email_draft": {
+                            "to": data.get("to"),
+                            "subject": data.get("subject"),
+                            "body": data.get("body"),
+                        }
+                    }
+
+        return {
+            "thread_id": thread_id,
+            "triage_label": final_state.get("triage_label"),
+        }
 
     except Exception as e:
-        print(f"An error occurred during execution: {e}")
+        logger.error(f"Error processing email: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/review-action")
+def review_action(request: ReviewActionRequest) -> Dict[str, Any]:
+    """Resume graph execution based on user review."""
+    
+    try:
+        config = {
+            "configurable": {
+                "thread_id": request.thread_id,
+                "user_id": request.user_id
+            }
+        }
 
-    finally:
-        # 5. CRITICAL: Close the pool to prevent "cannot join current thread" error
-        print("Closing connection pool...")
-        pool.close()
+        if request.status == "rejected":
+            payload = Command(resume={
+                "status": "rejected",
+                "feedback": request.feedback
+            })
+        elif request.status == "approved":
+            payload = Command(resume={
+                "status": "approved"
+            })
+        else:
+            raise HTTPException(status_code=400, detail="Invalid status")
+
+        final_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 parsed_interrupt:
+                data = parsed_interrupt["data"]
+                return {
+                    "status": "needs_review",
+                    "thread_id": request.thread_id,
+                    "triage_label": final_state.get("triage_label"),
+                    "action": parsed_interrupt["action"],
+                    "email_draft": {
+                        "to": data.get("to"),
+                        "subject": data.get("subject"),
+                        "body": data.get("body"),
+                    }
+                }
+        
+        # Draft created, review complete
+        if final_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"),
+            }
+
+    except Exception as e:
+        logger.error(f"Error in review action: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

app/tools/email_writing_agent_tools.py

@@ -3,19 +3,27 @@ from googleapiclient.errors import HttpError
 from app.schemas.email_writing_agent_tools_schema import CreateDraftSchema, SendDraftSchema
 from langchain.tools import tool
 from langchain_google_community import GmailToolkit
-
+from typing import Annotated, Union
+from langchain_core.tools import InjectedToolCallId, tool
+from langgraph.types import Command
+from langchain_core.messages import SystemMessage, HumanMessage,ToolMessage,AIMessage,BaseMessage
 
 @tool(args_schema=CreateDraftSchema)
-def create_gmail_draft(to: str, subject: str, body: str):
+def create_gmail_draft(
+    to: Union[str, list], 
+    subject: str, 
+    body: str,
+    tool_call_id: Annotated[str, InjectedToolCallId] # Injected ID
+):
     """Creates a new Gmail draft after human approval."""
-
+    
     if isinstance(to, list):
-        if len(to) > 0:
-            to = str(to[0])
-        else:
-            return "ERROR: 'to' parameter is an empty list. Please provide a valid email string."
+        to = str(to[0]) if len(to) > 0 else "ERROR"
     
-    # 1. Pause and ask for review
+    if to == "ERROR":
+        return "ERROR: 'to' parameter is empty."
+
+    # 1. Human-in-the-loop Interrupt
     response = interrupt({
         "action": "review_draft",
         "data": {"to": to, "subject": subject, "body": body}
@@ -24,40 +32,59 @@ def create_gmail_draft(to: str, subject: str, body: str):
     toolkit = GmailToolkit() 
     draft_tool = [t for t in toolkit.get_tools() if t.name == "create_gmail_draft"][0]
 
-    # 2. Handle the response
+    # 2. Handle Logic
     if response.get("status") == "approved":
-        reply = draft_tool.invoke({
-            "message": body,
-            "to": [to],
-            "subject": subject
-        })
-
-        draft_id=reply.split(":")[1].strip()
-        return f"Successfully created draft : <id>{draft_id}</id> <subject>{subject}</subject> <body>{body}</body> take user permission before submitting"
-    
+        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>"
+            
+            # UPDATE STATE: Save draft_id directly
+            return Command(
+                update={
+                    "draft_id": draft_id, 
+                    "reply_subject": subject,
+                    "reply_email_body": body,
+                    "messages": [ToolMessage(content, tool_call_id=tool_call_id)]
+                }
+            )
+        except IndexError:
+            return f"Draft created, but response parsing failed: {reply}"
     else:
-        # Get the feedback from the user response
-        feedback = response.get("feedback", "User rejected without specific notes.")
-        
-        # We return this to the AGENT so it can read it and rewrite the draft
-        return f"DRAFT REJECTED BY USER. Feedback: {feedback}. Please rewrite the draft based on this feedback and try again."
+        feedback = response.get("feedback", "User rejected.")
+        return f"DRAFT REJECTED BY USER. Feedback: {feedback}. Please rewrite."
     
 
-
+#---------------------------------------------------------------------------
 
 @tool(args_schema=SendDraftSchema)
-def send_draft_by_id(draft_id: str):
+def send_draft_by_id(
+    draft_id: str,
+    tool_call_id: Annotated[str, InjectedToolCallId] # 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}
         ).execute()
-        return f"SUCCESS: Sent! a Gmail with  ID: <id>{result['id']}</id>"
+        
+        sent_id = result['id']
+        content = f"SUCCESS: Sent! a Gmail with ID: <id>{sent_id}</id>"
+        
+        # UPDATE STATE: Save sent_message_id directly
+        return Command(
+            update={
+                "sent_message_id": sent_id,
+                "messages": [ToolMessage(content, tool_call_id=tool_call_id)]
+            }
+        )
     except HttpError as error:
-        if error.resp.status == 404:
-            return f"ERROR: Draft ID {draft_id} was not found. Please verify the ID or check if it was already sent."
-        return f"ERROR: An unexpected error occurred: {error}"
+        error_msg = f"ERROR: {error}"
+        return Command(
+            update={"messages": [ToolMessage(error_msg, tool_call_id=tool_call_id)]}
+        )
+
 
 
 email_writing_agent_tools = [

app/utils/interrupt_utils.py

@@ -0,0 +1,18 @@
+
+
+
+def parse_interrupt(final_state: dict):
+    if "__interrupt__" not in final_state:
+        return None
+
+    interrupt_state = final_state["__interrupt__"]
+    if not interrupt_state:
+        return None
+
+    interrupt = interrupt_state[0]
+    value = getattr(interrupt, "value", {}) or {}
+
+    return {
+        "action": value.get("action"),
+        "data": value.get("data", {})
+    }