95 lines
3 KiB
Python
95 lines
3 KiB
Python
"""PostgreSQL checkpointer utilities for LangGraph memory."""
|
|
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from langgraph.checkpoint.postgres import PostgresSaver
|
|
from psycopg_pool import AsyncConnectionPool
|
|
from modules.shared.configuration import APP_CONFIG
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Global checkpointer instance
|
|
_checkpointer_instance: Optional[PostgresSaver] = None
|
|
_connection_pool: Optional[AsyncConnectionPool] = None
|
|
|
|
|
|
async def initialize_checkpointer() -> None:
|
|
"""Initialize the PostgreSQL checkpointer for LangGraph.
|
|
|
|
This should be called during application startup.
|
|
Creates a connection pool and PostgresSaver instance.
|
|
"""
|
|
global _checkpointer_instance, _connection_pool
|
|
|
|
if _checkpointer_instance is not None:
|
|
logger.info("Checkpointer already initialized")
|
|
return
|
|
|
|
try:
|
|
# Get database configuration from environment
|
|
host = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_HOST", "localhost")
|
|
database = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_DATABASE", "poweron_chat")
|
|
user = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_USER", "poweron_dev")
|
|
password = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_PASSWORD_SECRET")
|
|
port = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_PORT", "5432")
|
|
|
|
# Build connection string
|
|
connection_string = f"postgresql://{user}:{password}@{host}:{port}/{database}"
|
|
|
|
# Create async connection pool
|
|
_connection_pool = AsyncConnectionPool(
|
|
conninfo=connection_string,
|
|
min_size=2,
|
|
max_size=10,
|
|
)
|
|
|
|
# Initialize the connection pool
|
|
await _connection_pool.open()
|
|
|
|
# Create PostgresSaver with the pool
|
|
_checkpointer_instance = PostgresSaver(_connection_pool)
|
|
|
|
# Setup the checkpointer (creates tables if needed)
|
|
await _checkpointer_instance.setup()
|
|
|
|
logger.info("PostgreSQL checkpointer initialized successfully")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize PostgreSQL checkpointer: {str(e)}")
|
|
raise
|
|
|
|
|
|
async def close_checkpointer() -> None:
|
|
"""Close the checkpointer and connection pool.
|
|
|
|
This should be called during application shutdown.
|
|
"""
|
|
global _checkpointer_instance, _connection_pool
|
|
|
|
if _connection_pool is not None:
|
|
try:
|
|
await _connection_pool.close()
|
|
logger.info("PostgreSQL checkpointer connection pool closed")
|
|
except Exception as e:
|
|
logger.error(f"Error closing checkpointer connection pool: {str(e)}")
|
|
|
|
_checkpointer_instance = None
|
|
_connection_pool = None
|
|
|
|
|
|
def get_checkpointer() -> PostgresSaver:
|
|
"""Get the global PostgreSQL checkpointer instance.
|
|
|
|
Returns:
|
|
The initialized PostgresSaver instance
|
|
|
|
Raises:
|
|
RuntimeError: If checkpointer is not initialized
|
|
"""
|
|
if _checkpointer_instance is None:
|
|
raise RuntimeError(
|
|
"PostgreSQL checkpointer not initialized. "
|
|
"Call initialize_checkpointer() during application startup."
|
|
)
|
|
return _checkpointer_instance
|