Refactor code to use new AsyncPostgresSaver. Delete old PostgresSaver…

… class.

backend/app/core/config.py

@@ -2,6 +2,7 @@
 import warnings
 from typing import Annotated, Any, Literal
 
+from psycopg.rows import dict_row
 from pydantic import (
     AnyUrl,
     BeforeValidator,
@@ -66,6 +67,13 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn:
             path=self.POSTGRES_DB,
         )
 
+    # For checkpointer
+    SQLALCHEMY_CONNECTION_KWARGS: dict[str, Any] = {
+        "autocommit": True,
+        "prepare_threshold": 0,
+        "row_factory": dict_row,
+    }
+
     @computed_field  # type: ignore[misc]
     @property
     def PG_DATABASE_URI(self) -> str:

backend/app/core/graph/build.py

@@ -13,7 +13,8 @@
 )
 from langchain_core.runnables import RunnableLambda
 from langchain_core.runnables.config import RunnableConfig
-from langgraph.checkpoint import BaseCheckpointSaver
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
 from langgraph.graph import END, StateGraph
 from langgraph.graph.graph import CompiledGraph
 from langgraph.prebuilt import (
@@ -22,7 +23,6 @@
 from psycopg import AsyncConnection
 
 from app.core.config import settings
-from app.core.graph.checkpoint.postgres import PostgresSaver
 from app.core.graph.members import (
     GraphLeader,
     GraphMember,
@@ -471,8 +471,11 @@ async def generator(
     ]
 
     try:
-        async with await AsyncConnection.connect(settings.PG_DATABASE_URI) as conn:
-            checkpointer = PostgresSaver(async_connection=conn)
+        async with await AsyncConnection.connect(
+            settings.PG_DATABASE_URI,
+            **settings.SQLALCHEMY_CONNECTION_KWARGS,
+        ) as conn:
+            checkpointer = AsyncPostgresSaver(conn=conn)
             if team.workflow == "hierarchical":
                 teams = convert_hierarchical_team_to_dict(team, members)
                 team_leader = list(teams.keys())[0]