refactor(langgraph): combine context managers

Author: NicholasGoh

backend/api/core/dependencies.py

@@ -1,13 +1,18 @@
 from contextlib import asynccontextmanager
 from typing import Annotated, AsyncGenerator
 
+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.config import settings
+from api.core.logs import uvicorn
 
 
 def get_llm() -> ChatOpenAI:
@@ -42,3 +47,30 @@ async def mcp_sse_client() -> AsyncGenerator[ClientSession, None]:
         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 mcp_sse_client() as session:
+            tools = await load_mcp_tools(session)
+            yield checkpointer, tools, session

backend/api/routers/llms.py

@@ -3,17 +3,13 @@
 import psycopg.errors
 from fastapi import APIRouter
 from langchain_core.messages import HumanMessage
-from langchain_mcp_adapters.tools import load_mcp_tools
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
-from psycopg_pool import AsyncConnectionPool
 from sse_starlette.sse import EventSourceResponse
 from starlette.responses import Response
 
 from api.core.agent.orchestration import get_config, get_graph
-from api.core.config import settings
-from api.core.dependencies import LLMDep
+from api.core.dependencies import LLMDep, setup_graph
 from api.core.logs import print, uvicorn
-from api.routers.mcps import mcp_sse_client
 
 router = APIRouter(tags=["chat"])
 
@@ -62,23 +58,13 @@ async def stream_graph(
     query: str,
     llm: LLMDep,
 ) -> AsyncGenerator[dict[str, str], None]:
-    # 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 mcp_sse_client() as session:
-            tools = await load_mcp_tools(session)
-            graph = get_graph(llm, tools=tools, checkpointer=checkpointer)
-            config = get_config()
-            events = dict(messages=[HumanMessage(content=query)])
-
-            async for event in graph.astream_events(
-                events, config, version="v2"
-            ):
-                if event.get("event").endswith("end"):
-                    print(event)
-                yield dict(data=event)
+    async with setup_graph() as resources:
+        checkpointer, tools, _ = resources
+        graph = get_graph(llm, tools=tools, checkpointer=checkpointer)
+        config = get_config()
+        events = dict(messages=[HumanMessage(content=query)])
+
+        async for event in graph.astream_events(events, config, version="v2"):
+            if event.get("event").endswith("end"):
+                print(event)
+            yield dict(data=event)