fix: checkpointer setup

app/main.py

@@ -5,6 +5,7 @@
 from langchain.chat_models import init_chat_model
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
 from loguru import logger
+from psycopg.rows import dict_row
 from psycopg_pool import AsyncConnectionPool
 
 from app.agent import ReActAgent
@@ -13,59 +14,65 @@
 from app.agent.tools import BDToolkit
 from app.api.main import api_router
 from app.db.database import engine, init_database
-from app.logging import setup_logger
+from app.log_config import setup_logger
 from app.settings import settings
 
 setup_logger()
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):  # pragma: no cover
-    if settings.AUTH_DEV_MODE and settings.ENVIRONMENT == "development":
-        logger.warning(
-            "AUTH DEV MODE ENABLED: JWT validation is bypassed, "
-            f"all requests will use user_id={settings.AUTH_DEV_USER_ID}"
+    try:
+        if settings.AUTH_DEV_MODE and settings.ENVIRONMENT == "development":
+            logger.warning(
+                "AUTH DEV MODE ENABLED: JWT validation is bypassed, "
+                f"all requests will use user_id={settings.AUTH_DEV_USER_ID}"
+            )
+
+        if settings.AUTH_DEV_MODE and settings.ENVIRONMENT != "development":
+            logger.warning(
+                f"AUTH_DEV_MODE is enabled but ENVIRONMENT is '{settings.ENVIRONMENT}'. "
+                "Auth dev mode will be ignored."
+            )
+
+        await init_database(engine)
+
+        # Connection kwargs defined according to:
+        # https://github.com/langchain-ai/langgraph/issues/2887
+        # https://langchain-ai.github.io/langgraph/how-tos/persistence_postgres
+        conn_kwargs = {
+            "autocommit": True,
+            "prepare_threshold": 0,
+            "row_factory": dict_row,
+        }
+
+        model = init_chat_model(
+            model=settings.MODEL_URI,
+            temperature=settings.MODEL_TEMPERATURE,
+            credentials=settings.GOOGLE_CREDENTIALS,
         )
 
-    if settings.AUTH_DEV_MODE and settings.ENVIRONMENT != "development":
-        logger.warning(
-            f"AUTH_DEV_MODE is enabled but ENVIRONMENT is '{settings.ENVIRONMENT}'. "
-            "Auth dev mode will be ignored."
-        )
-
-    await init_database(engine)
-
-    # Connection kwargs defined according to:
-    # https://github.com/langchain-ai/langgraph/issues/2887
-    # https://langchain-ai.github.io/langgraph/how-tos/persistence_postgres
-    conn_kwargs = {"autocommit": True, "prepare_threshold": 0}
-
-    model = init_chat_model(
-        model=settings.MODEL_URI,
-        temperature=settings.MODEL_TEMPERATURE,
-        credentials=settings.GOOGLE_CREDENTIALS,
-    )
+        async with AsyncConnectionPool(
+            conninfo=settings.DB_URL, max_size=8, kwargs=conn_kwargs
+        ) as pool:
+            checkpointer = AsyncPostgresSaver(pool)
 
-    async with AsyncConnectionPool(
-        conninfo=settings.DB_URL, max_size=8, kwargs=conn_kwargs
-    ) as pool:
-        checkpointer = AsyncPostgresSaver(pool)
-
-        await checkpointer.setup()
-
-        agent = ReActAgent(
-            model=model,
-            tools=BDToolkit.get_tools(),
-            start_hook=trim_messages_before_agent,
-            system_prompt=SYSTEM_PROMPT,
-            checkpointer=checkpointer,
-        )
+            agent = ReActAgent(
+                model=model,
+                tools=BDToolkit.get_tools(),
+                start_hook=trim_messages_before_agent,
+                system_prompt=SYSTEM_PROMPT,
+                checkpointer=checkpointer,
+            )
 
-        app.state.agent = agent
+            app.state.agent = agent
 
-        yield
+            yield
 
-    await engine.dispose()
+        await engine.dispose()
+    except Exception:
+        logger.exception("Lifespan failed:")
+        raise
 
 
 app = FastAPI(lifespan=lifespan)

charts/basedosdados-chatbot/templates/migration.yaml

@@ -19,7 +19,7 @@ spec:
         - name: basedosdados-chatbot-migration
           image: "{{ .Values.chatbot.image.name }}:{{ .Values.chatbot.image.tag }}"
           imagePullPolicy: {{ .Values.chatbot.image.pullPolicy }}
-          command: ["alembic", "upgrade", "head"]
+          command: ["./scripts/migrate.sh"]
           env:
             - name: JWT_ALGORITHM
               valueFrom:

compose.yaml

@@ -24,7 +24,7 @@ services:
 
   migration:
     build: .
-    command: alembic upgrade head
+    command: ./scripts/migrate.sh
     env_file: .env
     depends_on:
       database:

scripts/migrate.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+printf "Running Alembic migrations...\n"
+alembic upgrade head
+
+printf "Setting up LangGraph checkpointer...\n"
+python -m scripts.setup_checkpointer
+
+printf "Migration completed successfully.\n"

scripts/setup_checkpointer.py

@@ -0,0 +1,17 @@
+"""
+Run LangGraph checkpointer setup. `checkpointer.setup()` should be called only once.
+Ref: https://docs.langchain.com/oss/python/langgraph/add-memory#example-using-postgres-checkpointer
+"""
+
+from langgraph.checkpoint.postgres import PostgresSaver
+
+from app.settings import settings
+
+
+def main():
+    with PostgresSaver.from_conn_string(settings.DB_URL) as checkpointer:
+        checkpointer.setup()
+
+
+if __name__ == "__main__":
+    main()