File size: 1,599 Bytes
82dbc68
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
"""LangGraph checkpointer bootstrap for persistent thread state.

Resolution order:
1. `LANGGRAPH_CHECKPOINT_POSTGRES_URI` -> Postgres checkpointer
2. `LANGGRAPH_CHECKPOINT_SQLITE_PATH`  -> SQLite checkpointer
3. Fallback to `${DEER_FLOW_HOME}/checkpoints.sqlite`
"""

import atexit
import os
from pathlib import Path

from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.sqlite import SqliteSaver

_ctx = None


def _default_sqlite_path() -> str:
    deer_flow_home = os.getenv("DEER_FLOW_HOME")
    if deer_flow_home:
        base = Path(deer_flow_home)
    else:
        base = Path.cwd() / ".deer-flow"
    base.mkdir(parents=True, exist_ok=True)
    return str((base / "checkpoints.sqlite").resolve())


def _init_checkpointer():
    global _ctx

    postgres_uri = os.getenv("LANGGRAPH_CHECKPOINT_POSTGRES_URI", "").strip()
    if postgres_uri:
        _ctx = PostgresSaver.from_conn_string(postgres_uri)
        saver = _ctx.__enter__()
        saver.setup()
        return saver

    sqlite_path = os.getenv("LANGGRAPH_CHECKPOINT_SQLITE_PATH", "").strip() or _default_sqlite_path()
    sqlite_file = Path(sqlite_path).expanduser().resolve()
    sqlite_file.parent.mkdir(parents=True, exist_ok=True)
    _ctx = SqliteSaver.from_conn_string(str(sqlite_file))
    saver = _ctx.__enter__()
    saver.setup()
    return saver


def _close_context() -> None:
    global _ctx
    if _ctx is not None:
        try:
            _ctx.__exit__(None, None, None)
        finally:
            _ctx = None


checkpointer = _init_checkpointer()
atexit.register(_close_context)