changes

app/database/connection.py

@@ -29,11 +29,4 @@ SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
 def get_session() -> Session:
     return SessionLocal()
 
-DB_URL_FOR_CHECKPOINTER_STORE=settings.DB_URL_FOR_CHECKPOINTER_STORE
 
-pool = ConnectionPool(
-    conninfo=DB_URL_FOR_CHECKPOINTER_STORE, 
-    min_size=1, 
-    max_size=10,
-    kwargs={"autocommit": True,"row_factory": dict} 
-)

app/graph.py

@@ -19,6 +19,10 @@ from app.tools.email_writing_agent_tools import email_writing_agent_tools
 from IPython.display import Image, display
 from app.persistance.memory_store_checkpointer_config import memory_store, checkpointer
 from langchain_google_community import GmailToolkit
+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
 
 # Define a standard retry policy for database-heavy nodes
 db_retry_policy = RetryPolicy(
@@ -136,6 +140,7 @@ builder.add_edge("archive_node", END)
 toolkit = GmailToolkit()
 
 
+
 graph=builder.compile(checkpointer=checkpointer, store=memory_store)
 
 

app/main.py

@@ -1,33 +1,47 @@
-from app.graph import graph
-from app.state.state import EmailAgentState
 import time
 from psycopg import OperationalError
+from app.graph import graph
+from app.state.state import EmailAgentState
+from app.database.connection import pool
 
-# Define your thread configuration
-config = {"configurable": {"thread_id": "user_abc_123"}}
 
+config = {"configurable": {"thread_id": "user_abc_123"}}
 
 input_data: EmailAgentState = {
-    "user_email_id":    "gaykaratharva7@gmail.com",
-    "user_id":          1,
-    "user_name":        "Atharva",
-    "sender_email_id":  "atharvagaykar36@gmail.com",
+    "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*\r\npipeline. We are currently utilizing the fine-tuned DistilBERT model to\r\ngenerate semantic embeddings for incoming messages.\r\n\r\nTo ensure the system is correctly identifying malicious intent, I've\r\nprocessed a suspicious sample using our custom structural tokens: [SSUB],\r\n[SBODY], [LINK], and [PHONE]. This preserves the email's structural context\r\nwhile protecting sensitive data.\r\n[MODEL EVALUATION DATA]\r\n\r\n*1. Semantic Context:* The fine-tuned DistilBERT has mapped the input to a\r\nhigh-dimensional vector space. Initial checks suggest strong clustering\r\nwith known phishing signatures.\r\n\r\n*2. URL Feature Engineering:* Our hybrid pipeline extracted numerical\r\nindicators from the embedded links.\r\n\r\n   -\r\n\r\n   Subdomain count: 4\r\n   -\r\n\r\n   Suspicious keywords: 2\r\n   -\r\n\r\n   Special characters: @, -, .\r\n   -\r\n\r\n   Redirection detected: True\r\n\r\n*3. XGBoost Classification:* The combined feature set (DistilBERT\r\nembeddings + numerical URL features) has been passed to the XGBoost\r\nclassifier.\r\n\r\n   -\r\n\r\n   *Current Test Accuracy:* 99.35%\r\n\r\nFinal Confirmation Needed:\r\n\r\n   1.\r\n\r\n   Should we deploy the current XGBoost weights to the:\r\n   https://huggingface.co/spaces/Gaykar/ClassifyEmail\r\n   2.\r\n\r\n   Do you want to review the Hybrid_model_preparation.ipynb logic before we\r\n   push to the:\r\n   https://github.com/Atharva-Gaykar/AI-Driven-Email-Threat-Detection\r\n   3.\r\n\r\n   Are the [LINK] and [PHONE] placeholders correctly masking the PII\r\n   (Personally Identifiable Information) according to the project spec?\r\n\r\nPlease provide your approval to proceed with the Docker deployment.\r\n\r\nBest regards,\r\n\r\nVinit Security AI Engineer\r\n"""
+    "sender_email_body": """Dear Atharva,\r\n\r\nI have completed the integration of the AI-Driven Email Threat Detection pipeline... [truncated for brevity]"""
 }
 
-
-
-
-for i in range(3): # Try 3 times to account for Neon wake-up
+if __name__ == "__main__":
+    result = None
     try:
-        result = graph.invoke(input_data, config=config)
-        break
-    except OperationalError:
-        print("Waiting for Neon to wake up...")
-        time.sleep(5)
-
-
-print(result)
-
-
+        # 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)
+
+    except Exception as e:
+        print(f"An error occurred during execution: {e}")
+
+    finally:
+        # 5. CRITICAL: Close the pool to prevent "cannot join current thread" error
+        print("Closing connection pool...")
+        pool.close()