303 lines
10 KiB
Python
303 lines
10 KiB
Python
"""Async checkpointer factory.
|
|
|
|
Provides an **async context manager** for long-running async servers that need
|
|
proper resource cleanup.
|
|
|
|
Supported backends: memory, sqlite, postgres.
|
|
|
|
Usage (e.g. FastAPI lifespan)::
|
|
|
|
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
|
|
|
|
async with make_checkpointer() as checkpointer:
|
|
app.state.checkpointer = checkpointer # InMemorySaver if not configured
|
|
|
|
For sync usage see :mod:`deerflow.runtime.checkpointer.provider`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
from collections.abc import AsyncIterator
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from langgraph.types import Checkpointer
|
|
|
|
from deerflow.config.app_config import AppConfig, get_app_config
|
|
from deerflow.runtime.checkpointer.provider import (
|
|
POSTGRES_CONN_REQUIRED,
|
|
POSTGRES_INSTALL,
|
|
SQLITE_INSTALL,
|
|
)
|
|
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _csv_env(name: str) -> list[str]:
|
|
raw = os.getenv(name, "")
|
|
return [part.strip() for part in raw.split(",") if part.strip()]
|
|
|
|
|
|
def _dual_read_enabled(config) -> bool:
|
|
if getattr(config, "dual_read_enabled", False):
|
|
return True
|
|
return os.getenv("DEER_FLOW_CHECKPOINT_DUAL_READ", "").strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
|
|
def _use_postgres_enabled(config) -> bool:
|
|
raw = os.getenv("DEER_FLOW_USE_PG_CHECKPOINTER")
|
|
if raw:
|
|
return raw.strip().lower() not in {"0", "false", "no", "off"}
|
|
return bool(getattr(config, "use_postgres", True))
|
|
|
|
|
|
def _dual_read_paths(config) -> list[str]:
|
|
paths = list(getattr(config, "dual_read_sqlite_paths", []) or [])
|
|
paths.extend(_csv_env("DEER_FLOW_CHECKPOINT_SQLITE_PATHS"))
|
|
resolved: list[str] = []
|
|
seen: set[str] = set()
|
|
for path in paths:
|
|
conn_str = resolve_sqlite_conn_str(path)
|
|
if conn_str == ":memory:" or conn_str.startswith("file:"):
|
|
logger.warning("Skipping unsupported checkpoint migration SQLite path: %s", path)
|
|
continue
|
|
if conn_str not in seen:
|
|
resolved.append(conn_str)
|
|
seen.add(conn_str)
|
|
return resolved
|
|
|
|
|
|
def _sqlite_conn_str_for_postgres_fallback(config) -> str:
|
|
configured = getattr(config, "sqlite_connection_string", None)
|
|
if configured:
|
|
return configured
|
|
paths = _dual_read_paths(config)
|
|
return paths[0] if paths else "store.db"
|
|
|
|
|
|
def _pool_size_env(name: str, default: int) -> int:
|
|
raw = os.getenv(name, "").strip()
|
|
if not raw:
|
|
return default
|
|
try:
|
|
return max(0, int(raw))
|
|
except ValueError:
|
|
logger.warning("Invalid %s=%r; using default %d", name, raw, default)
|
|
return default
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def _postgres_saver_from_pool(conn_string: str) -> AsyncIterator[Checkpointer]:
|
|
"""Yield an ``AsyncPostgresSaver`` backed by a self-healing connection pool.
|
|
|
|
``AsyncPostgresSaver.from_conn_string`` opens a *single* long-lived
|
|
connection. If that one connection dies (Postgres restart, idle timeout,
|
|
network blip, or a prior broken pipeline), every later checkpointer call
|
|
fails permanently with "the connection is closed" until the process is
|
|
restarted. A pool with ``check=check_connection`` validates each connection
|
|
on checkout and transparently replaces dead ones, so the runtime recovers
|
|
on its own.
|
|
"""
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
|
from psycopg.rows import dict_row
|
|
from psycopg_pool import AsyncConnectionPool
|
|
|
|
pool = AsyncConnectionPool(
|
|
conninfo=conn_string,
|
|
min_size=_pool_size_env("DEER_FLOW_PG_POOL_MIN", 1),
|
|
max_size=_pool_size_env("DEER_FLOW_PG_POOL_MAX", 10),
|
|
max_idle=300.0,
|
|
max_lifetime=3600.0,
|
|
check=AsyncConnectionPool.check_connection,
|
|
kwargs={"autocommit": True, "prepare_threshold": 0, "row_factory": dict_row},
|
|
open=False,
|
|
)
|
|
await pool.open(wait=True)
|
|
try:
|
|
saver = AsyncPostgresSaver(pool)
|
|
await saver.setup()
|
|
yield saver
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def _sqlite_checkpointer_from_conn_str(raw: str) -> AsyncIterator[Checkpointer]:
|
|
try:
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
except ImportError as exc:
|
|
raise ImportError(SQLITE_INSTALL) from exc
|
|
|
|
conn_str = resolve_sqlite_conn_str(raw)
|
|
await asyncio.to_thread(ensure_sqlite_parent_dir, conn_str)
|
|
async with AsyncSqliteSaver.from_conn_string(conn_str) as saver:
|
|
await saver.setup()
|
|
yield saver
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Async factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
|
|
"""Async context manager that constructs and tears down a checkpointer."""
|
|
if config.type == "memory":
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
yield InMemorySaver()
|
|
return
|
|
|
|
if config.type == "sqlite":
|
|
async with _sqlite_checkpointer_from_conn_str(config.connection_string or "store.db") as saver:
|
|
yield saver
|
|
return
|
|
|
|
if config.type == "postgres":
|
|
if not _use_postgres_enabled(config):
|
|
conn_str = _sqlite_conn_str_for_postgres_fallback(config)
|
|
logger.info("Checkpointer: use_postgres=false; using SQLite fallback (%s)", conn_str)
|
|
async with _sqlite_checkpointer_from_conn_str(conn_str) as saver:
|
|
yield saver
|
|
return
|
|
|
|
try:
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver # noqa: F401
|
|
except ImportError as exc:
|
|
raise ImportError(POSTGRES_INSTALL) from exc
|
|
|
|
if not config.connection_string:
|
|
raise ValueError(POSTGRES_CONN_REQUIRED)
|
|
|
|
async with _postgres_saver_from_pool(config.connection_string) as saver:
|
|
if not _dual_read_enabled(config):
|
|
yield saver
|
|
return
|
|
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
|
|
from deerflow.runtime.checkpointer.checkpoint_migration import (
|
|
DualReadCheckpointer,
|
|
SQLiteCheckpointSource,
|
|
)
|
|
|
|
sources: list[SQLiteCheckpointSource] = []
|
|
for sqlite_path in _dual_read_paths(config):
|
|
if not Path(sqlite_path).exists():
|
|
logger.info("Checkpoint dual-read source does not exist, skipping: %s", sqlite_path)
|
|
continue
|
|
ctx = AsyncSqliteSaver.from_conn_string(sqlite_path)
|
|
source_saver = await ctx.__aenter__()
|
|
await source_saver.setup()
|
|
sources.append(
|
|
SQLiteCheckpointSource(
|
|
path=sqlite_path,
|
|
saver=source_saver,
|
|
context=ctx,
|
|
progress_path=Path(sqlite_path).with_suffix(Path(sqlite_path).suffix + ".pg-migration.json"),
|
|
)
|
|
)
|
|
|
|
if not sources:
|
|
logger.info("Checkpoint dual-read enabled but no SQLite sources were opened; using Postgres only")
|
|
yield saver
|
|
return
|
|
|
|
proxy = DualReadCheckpointer(
|
|
saver,
|
|
sources,
|
|
)
|
|
try:
|
|
yield proxy
|
|
finally:
|
|
await proxy.aclose_sources()
|
|
return
|
|
|
|
raise ValueError(f"Unknown checkpointer type: {config.type!r}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public async context manager
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def _async_checkpointer_from_database(db_config) -> AsyncIterator[Checkpointer]:
|
|
"""Async context manager that constructs a checkpointer from unified DatabaseConfig."""
|
|
if db_config.backend == "memory":
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
yield InMemorySaver()
|
|
return
|
|
|
|
if db_config.backend == "sqlite":
|
|
try:
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
except ImportError as exc:
|
|
raise ImportError(SQLITE_INSTALL) from exc
|
|
|
|
conn_str = db_config.checkpointer_sqlite_path
|
|
ensure_sqlite_parent_dir(conn_str)
|
|
async with AsyncSqliteSaver.from_conn_string(conn_str) as saver:
|
|
await saver.setup()
|
|
yield saver
|
|
return
|
|
|
|
if db_config.backend == "postgres":
|
|
try:
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver # noqa: F401
|
|
except ImportError as exc:
|
|
raise ImportError(POSTGRES_INSTALL) from exc
|
|
|
|
if not db_config.postgres_url:
|
|
raise ValueError("database.postgres_url is required for the postgres backend")
|
|
|
|
async with _postgres_saver_from_pool(db_config.postgres_url) as saver:
|
|
yield saver
|
|
return
|
|
|
|
raise ValueError(f"Unknown database backend: {db_config.backend!r}")
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def make_checkpointer(app_config: AppConfig | None = None) -> AsyncIterator[Checkpointer]:
|
|
"""Async context manager that yields a checkpointer for the caller's lifetime.
|
|
Resources are opened on enter and closed on exit -- no global state::
|
|
|
|
async with make_checkpointer(app_config) as checkpointer:
|
|
app.state.checkpointer = checkpointer
|
|
|
|
Yields an ``InMemorySaver`` when no checkpointer is configured in *config.yaml*.
|
|
|
|
Priority:
|
|
1. Legacy ``checkpointer:`` config section (backward compatible)
|
|
2. Unified ``database:`` config section
|
|
3. Default InMemorySaver
|
|
"""
|
|
|
|
if app_config is None:
|
|
app_config = get_app_config()
|
|
|
|
# Legacy: standalone checkpointer config takes precedence
|
|
if app_config.checkpointer is not None:
|
|
async with _async_checkpointer(app_config.checkpointer) as saver:
|
|
yield saver
|
|
return
|
|
|
|
# Unified database config
|
|
db_config = getattr(app_config, "database", None)
|
|
if db_config is not None and db_config.backend != "memory":
|
|
async with _async_checkpointer_from_database(db_config) as saver:
|
|
yield saver
|
|
return
|
|
|
|
# Default: in-memory
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
yield InMemorySaver()
|