changes

app/agent_memory_store/memory_store.py

@@ -1,20 +0,0 @@
-from psycopg_pool import ConnectionPool
-from langgraph.store.postgres import PostgresStore
-from langgraph.checkpoint.postgres import PostgresSaver
-from app.agent_memory_store.embeddings import remote_embeddings
-
-from app.core.config import settings
-
-
-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} 
-)
-4
-
-
-memory_store = PostgresStore(pool, index={"dims": 384, "embed": remote_embeddings})

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 memory_store
+from app.agent_memory_store import connection
 
 context_agent = create_agent(
     model=ChatGroq(
@@ -12,7 +12,7 @@ context_agent = create_agent(
         temperature=0.1,
     ),
     tools=context_agent_tools,
-    store=memory_store,  
+    store=connection,  
     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 memory_store
+from app.agent_memory_store import connection
 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=memory_store,
+    store=connection,
     instructions="Extract required info from incoming mail and its reply .",
     enable_inserts=True,
     enable_deletes=True,

app/database/connection.py

@@ -1,6 +1,19 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker, Session
 from app.core.config import settings
+from psycopg_pool import ConnectionPool
+
+
+
+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} 
+)
+
 
 DB_URL_FOR_SQL_AL=settings.DB_URL_FOR_SQL_AL
 
@@ -14,4 +27,13 @@ engine = create_engine(
 SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
 
 def get_session() -> Session:
-    return SessionLocal()
+    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} 
+)

app/graph.py

@@ -2,7 +2,6 @@ from app.state.state import EmailAgentState
 from app.nodes.triage_node import *
 from langgraph.prebuilt import ToolNode ,tools_condition
 from app.nodes.archive_node import  archive_node
-from app.tools.email_writing_agent_tools import *
 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
@@ -16,6 +15,7 @@ from langgraph.types import RetryPolicy
 from app.nodes.summarise_email_body_node import summarise_email_body_node
 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
 
 

app/main.py

@@ -0,0 +1,13 @@
+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
+
+
+
+
+checkpointer = PostgresSaver(pool)
+memory_store = PostgresStore(pool, embeddings=remote_embeddings)
+
+