greatmengqi
edf345cd72
- Freeze all config models (AppConfig + 15 sub-configs) with frozen=True - Purify from_file() — remove 9 load_*_from_dict() side-effect calls - Replace mtime/reload/push/pop machinery with single ContextVar + init_app_config() - Delete 10 sub-module globals and their getters/setters/loaders - Migrate 50+ consumers from get_*_config() to get_app_config().xxx - Expand DeerFlowContext: app_config + thread_id + agent_name (frozen dataclass) - Wire into Gateway runtime (worker.py) and DeerFlowClient via context= parameter - Remove sandbox_id from runtime.context — flows through ThreadState.sandbox only - Middleware/tools access runtime.context directly via Runtime[DeerFlowContext] generic - resolve_context() retained at server entry points for LangGraph Server fallback
107 lines
3.4 KiB
Python
107 lines
3.4 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.agents.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.agents.checkpointer.provider`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
from collections.abc import AsyncIterator
|
|
|
|
from langgraph.types import Checkpointer
|
|
|
|
from deerflow.agents.checkpointer.provider import (
|
|
POSTGRES_CONN_REQUIRED,
|
|
POSTGRES_INSTALL,
|
|
SQLITE_INSTALL,
|
|
)
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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":
|
|
try:
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
except ImportError as exc:
|
|
raise ImportError(SQLITE_INSTALL) from exc
|
|
|
|
conn_str = resolve_sqlite_conn_str(config.connection_string or "store.db")
|
|
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
|
|
return
|
|
|
|
if config.type == "postgres":
|
|
try:
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
|
except ImportError as exc:
|
|
raise ImportError(POSTGRES_INSTALL) from exc
|
|
|
|
if not config.connection_string:
|
|
raise ValueError(POSTGRES_CONN_REQUIRED)
|
|
|
|
async with AsyncPostgresSaver.from_conn_string(config.connection_string) as saver:
|
|
await saver.setup()
|
|
yield saver
|
|
return
|
|
|
|
raise ValueError(f"Unknown checkpointer type: {config.type!r}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public async context manager
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def make_checkpointer() -> 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() as checkpointer:
|
|
app.state.checkpointer = checkpointer
|
|
|
|
Yields an ``InMemorySaver`` when no checkpointer is configured in *config.yaml*.
|
|
"""
|
|
|
|
config = AppConfig.current()
|
|
|
|
if config.checkpointer is None:
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
yield InMemorySaver()
|
|
return
|
|
|
|
async with _async_checkpointer(config.checkpointer) as saver:
|
|
yield saver
|