chnages

app/agents/context_agent.py

@@ -4,7 +4,7 @@ from langchain_groq import ChatGroq
 from app.prompts.context_agent_prompt import context_agent_template
 from app.tools.context_agent_tools import context_agent_tools
 from typing import Any
-from app.agent_memory_store import connection
+from app.persistance.memory_store_checkpointer_config import memory_store
 
 context_agent = create_agent(
     model=ChatGroq(
@@ -12,7 +12,7 @@ context_agent = create_agent(
         temperature=0.1,
     ),
     tools=context_agent_tools,
-    store=connection,  
+    store=memory_store,  
     middleware=[
         ToolCallLimitMiddleware[Any,None](
            tool_name="search_memory",

app/agents/memory_manager_agent.py

@@ -2,7 +2,7 @@ import types
 from langchain_groq import ChatGroq
 from langmem import create_memory_store_manager
 from app.schemas.memory_agent_schema import EmailMemory
-from app.agent_memory_store import connection
+from app.persistance.memory_store_checkpointer_config import memory_store
 import os
 from app.core.config import settings
 
@@ -51,7 +51,7 @@ memory_manager_agent = create_memory_store_manager(
     model, 
     schemas=[EmailMemory],
     namespace=namespace,
-    store=connection,
+    store=memory_store,
     instructions="Extract required info from incoming mail and its reply .",
     enable_inserts=True,
     enable_deletes=True,

app/core/config.py

@@ -8,6 +8,9 @@ class Settings(BaseSettings):
     GROQ_API_KEY: str
     DB_URL_FOR_CHECKPOINTER_STORE: str
 
+    GMAIL_CREDENTIALS_PATH: str = "credentials.json"
+    GMAIL_TOKEN_PATH: str = "token.json"
+
     DB_URL_FOR_SQL_AL:str
  
     model_config = SettingsConfigDict(

app/database/connection.py

@@ -35,5 +35,5 @@ pool = ConnectionPool(
     conninfo=DB_URL_FOR_CHECKPOINTER_STORE, 
     min_size=1, 
     max_size=10,
-    kwargs={"autocommit": True} 
+    kwargs={"autocommit": True,"row_factory": dict} 
 )

app/graph.py

@@ -17,7 +17,8 @@ from app.nodes.check_token_count_node import *
 from psycopg import OperationalError # Or sqlalchemy.exc.OperationalError depending on your driver
 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
 
 # Define a standard retry policy for database-heavy nodes
 db_retry_policy = RetryPolicy(
@@ -126,20 +127,29 @@ builder.add_edge("archive_node", END)
 
 
 
-graph=builder.compile()
 
-print("--- Graph compiled successfully ---"
-)
+
+
+
+
+
+toolkit = GmailToolkit()
+
+
+graph=builder.compile(checkpointer=checkpointer, store=memory_store)
+
+
+
 
 display(graph)
 
 
-try:
-    # This creates a PNG and saves it to your project folder
-    graph_png = graph.get_graph().draw_mermaid_png()
-    with open("graph.png", "wb") as f:
-        f.write(graph_png)
-    print("--- Graph image saved as 'graph.png' ---")
-except Exception as e:
-    # This happens if you don't have the 'pypydot' or 'graphviz' dependencies
-    print(f"Could not generate graph image: {e}")
+# try:
+#     # This creates a PNG and saves it to your project folder
+#     graph_png = graph.get_graph().draw_mermaid_png()
+#     with open("graph.png", "wb") as f:
+#         f.write(graph_png)
+#     print("--- Graph image saved as 'graph.png' ---")
+# except Exception as e:
+#     # This happens if you don't have the 'pypydot' or 'graphviz' dependencies
+#     print(f"Could not generate graph image: {e}")

app/main.py

@@ -1,13 +1,33 @@
 from app.graph import graph
-from app.database.connection import pool
-from langgraph.checkpoint.postgres import PostgresSaver
-from langgraph.store.postgres import PostgresStore
-from app.utils.embeddings import remote_embeddings
+from app.state.state import EmailAgentState
+import time
+from psycopg import OperationalError
 
+# Define your thread configuration
+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",
+    "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"""
+}
 
-checkpointer = PostgresSaver(pool)
-memory_store = PostgresStore(pool, embeddings=remote_embeddings)
+
+
+
+for i in range(3): # Try 3 times to account for Neon wake-up
+    try:
+        result = graph.invoke(input_data, config=config)
+        break
+    except OperationalError:
+        print("Waiting for Neon to wake up...")
+        time.sleep(5)
+
+
+print(result)
 
 

app/persistance/memory_store_checkpointer_config.py

@@ -0,0 +1,10 @@
+from app.database.connection import pool
+from langgraph.checkpoint.postgres import PostgresSaver
+from langgraph.store.postgres import PostgresStore
+from app.utils.embeddings import remote_embeddings
+
+
+
+
+checkpointer = PostgresSaver(pool)
+memory_store = PostgresStore(pool, index={"dims": 384, "embed": remote_embeddings})