refactor(langgraph): pydantic typehint and validation for context managers

Author: NicholasGoh

backend/api/core/agent/persistence.py

@@ -0,0 +1,40 @@
+from contextlib import asynccontextmanager
+
+import psycopg
+import psycopg.errors
+import uvicorn
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from psycopg_pool import AsyncConnectionPool
+
+from api.core.logs import uvicorn
+
+
+@asynccontextmanager
+async def checkpointer_context(conn_str: str):
+    """
+    Async context manager that sets up and yields a LangGraph checkpointer.
+
+    Uses a psycopg async connection pool to initialize AsyncPostgresSaver.
+    Skips setup if checkpointer is already configured.
+
+    Args:
+        conn_str (str): PostgreSQL connection string.
+
+    Yields:
+        AsyncPostgresSaver: The initialized checkpointer.
+    """
+    # NOTE: LangGraph AsyncPostgresSaver does not support SQLAlchemy ORM Connections.
+    # A compatible psycopg connection is created via the connection pool to connect to the checkpointer.
+    async with AsyncConnectionPool(
+        conninfo=conn_str,
+        kwargs=dict(prepare_threshold=None),
+    ) as pool:
+        checkpointer = AsyncPostgresSaver(pool)
+        try:
+            await checkpointer.setup()
+        except (
+            psycopg.errors.DuplicateColumn,
+            psycopg.errors.ActiveSqlTransaction,
+        ):
+            uvicorn.warning("Skipping checkpointer setup — already configured.")
+        yield checkpointer

backend/api/core/dependencies.py

@@ -1,18 +1,15 @@
 from contextlib import asynccontextmanager
-from typing import Annotated, AsyncGenerator
+from typing import Annotated
 
-import psycopg.errors
 from fastapi import Depends
 from langchain_mcp_adapters.tools import load_mcp_tools
 from langchain_openai import ChatOpenAI
-from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
-from mcp import ClientSession
-from mcp.client.sse import sse_client
-from psycopg_pool import AsyncConnectionPool
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
+from api.core.agent.persistence import checkpointer_context
 from api.core.config import settings
-from api.core.logs import uvicorn
+from api.core.mcps import mcp_sse_client
+from api.core.models import Resource
 
 
 def get_llm() -> ChatOpenAI:
@@ -38,39 +35,15 @@ def get_engine() -> AsyncEngine:
 EngineDep = Annotated[AsyncEngine, Depends(get_engine)]
 
 
-@asynccontextmanager
-async def mcp_sse_client() -> AsyncGenerator[ClientSession, None]:
-    async with sse_client(f"http://mcp:{settings.mcp_server_port}/sse") as (
-        read_stream,
-        write_stream,
-    ):
-        async with ClientSession(read_stream, write_stream) as session:
-            await session.initialize()
-            yield session
-
-
-async def checkpointer_setup(pool):
-    checkpointer = AsyncPostgresSaver(pool)
-    try:
-        await checkpointer.setup()
-    except (
-        psycopg.errors.DuplicateColumn,
-        psycopg.errors.ActiveSqlTransaction,
-    ):
-        uvicorn.warning("Skipping checkpointer setup — already configured.")
-    return checkpointer
-
-
 @asynccontextmanager
 async def setup_graph():
-    # NOTE: LangGraph AsyncPostgresSaver does not support SQLAlchemy ORM Connections.
-    # A compatible psycopg connection is created via the connection pool to connect to the checkpointer.
-    async with AsyncConnectionPool(
-        conninfo=settings.checkpoint_conn_str,
-        kwargs=dict(prepare_threshold=None),
-    ) as pool:
-        checkpointer = await checkpointer_setup(pool)
-
+    async with checkpointer_context(
+        settings.checkpoint_conn_str
+    ) as checkpointer:
         async with mcp_sse_client() as session:
             tools = await load_mcp_tools(session)
-            yield checkpointer, tools, session
+            yield Resource(
+                checkpointer=checkpointer,
+                tools=tools,
+                session=session,
+            )

backend/api/core/mcps.py

@@ -0,0 +1,27 @@
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+
+from mcp import ClientSession
+from mcp.client.sse import sse_client
+
+from api.core.config import settings
+
+
+@asynccontextmanager
+async def mcp_sse_client() -> AsyncGenerator[ClientSession]:
+    """
+    Creates and initializes an MCP client session over SSE.
+
+    Establishes an SSE connection to the MCP server and yields an initialized
+    `ClientSession` for communication.
+
+    Yields:
+        ClientSession: An initialized MCP client session.
+    """
+    async with sse_client(f"http://mcp:{settings.mcp_server_port}/sse") as (
+        read_stream,
+        write_stream,
+    ):
+        async with ClientSession(read_stream, write_stream) as session:
+            await session.initialize()
+            yield session

backend/api/core/models.py

@@ -0,0 +1,13 @@
+from langchain_core.tools import StructuredTool
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from mcp import ClientSession
+from pydantic import BaseModel
+
+
+class Resource(BaseModel):
+    checkpointer: AsyncPostgresSaver
+    tools: list[StructuredTool]
+    session: ClientSession
+
+    class Config:
+        arbitrary_types_allowed = True

backend/api/routers/llms.py

@@ -58,9 +58,12 @@ async def stream_graph(
     query: str,
     llm: LLMDep,
 ) -> AsyncGenerator[dict[str, str], None]:
-    async with setup_graph() as resources:
-        checkpointer, tools, _ = resources
-        graph = get_graph(llm, tools=tools, checkpointer=checkpointer)
+    async with setup_graph() as resource:
+        graph = get_graph(
+            llm,
+            tools=resource.tools,
+            checkpointer=resource.checkpointer,
+        )
         config = get_config()
         events = dict(messages=[HumanMessage(content=query)])
 

backend/api/routers/mcps.py

@@ -3,7 +3,7 @@
 from fastapi import APIRouter
 from mcp import types
 
-from api.core.dependencies import mcp_sse_client
+from api.core.mcps import mcp_sse_client
 from shared_mcp.models import ToolRequest
 
 router = APIRouter(prefix="/mcps", tags=["mcps"])