LangGraph II

uv add langgraph langgraph-checkpoint-sqlite

# script_id: day_054_langgraph_durability/checkpoint_min
# Save a graph's state to SQLite, then read it back from a brand-new graph object.
# Install: uv add langgraph langgraph-checkpoint-sqlite · no API key, no services
import operator, sqlite3
from pathlib import Path
from typing import Annotated, TypedDict
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, START, StateGraph

class State(TypedDict):
    steps: Annotated[list[str], operator.add]   # reducer: node outputs are appended

def build(path: str):
    g = StateGraph(State)
    g.add_node("a", lambda s: {"steps": ["a"]})
    g.add_node("b", lambda s: {"steps": ["b"]})
    g.add_edge(START, "a"); g.add_edge("a", "b"); g.add_edge("b", END)
    saver = SqliteSaver(sqlite3.connect(path, check_same_thread=False))
    return g.compile(checkpointer=saver)

for old in Path(".").glob("ckpt.db*"):   # start clean: reruns would CONTINUE the old thread
    old.unlink()

cfg = {"configurable": {"thread_id": "t1"}, "recursion_limit": 25}
build("ckpt.db").invoke({"steps": []}, cfg)      # first "process" runs the graph
fresh = build("ckpt.db")                         # second "process": new objects, same file
print("t1", fresh.get_state(cfg).values, "· next:", fresh.get_state(cfg).next)

t1 {'steps': ['a', 'b']} · next: ()

# script_id: day_054_langgraph_durability/break_it
# Three ways to get checkpointing wrong, each caught and printed. No API key, no services.
import operator, sqlite3
from pathlib import Path
from typing import Annotated, TypedDict
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, START, StateGraph

class State(TypedDict):
    steps: Annotated[list[str], operator.add]

for old in Path(".").glob("break.db*"):   # start clean: reruns would CONTINUE the old thread
    old.unlink()

g = StateGraph(State)
g.add_node("a", lambda s: {"steps": ["a"]})
g.add_edge(START, "a"); g.add_edge("a", END)

# 1. from_conn_string returns a context manager, not a saver.
try:
    g.compile(checkpointer=SqliteSaver.from_conn_string("break.db"))
except TypeError as e:
    print("1.", str(e)[:70])
with SqliteSaver.from_conn_string("break.db") as saver:   # the right way to use it
    app = g.compile(checkpointer=saver)
    # 2. A checkpointer with no thread_id has nowhere to file the save.
    try:
        app.invoke({"steps": []})
    except ValueError as e:
        print("2.", str(e)[:70])
    # 3. Two users sharing one thread_id read and write the SAME state.
    shared = {"configurable": {"thread_id": "support"}, "recursion_limit": 25}
    app.invoke({"steps": ["alice-secret"]}, shared)
    print("3. bob sees:", app.invoke({"steps": ["bob"]}, shared)["steps"])

1. Invalid checkpointer provided. Expected an instance of `BaseCheckpoint
2. Checkpointer requires one or more of the following 'configurable' keys
3. bob sees: ['alice-secret', 'a', 'bob', 'a']

# script_id: day_054_langgraph_durability/durable_graph
# Day 54 additions to atlas/graph.py: a durable Issue Investigator you can crash, resume and fork.
# Install: uv add langgraph langgraph-checkpoint-sqlite langchain-core (live mode also: langchain langchain-openai)
# Env: ORBIT_DATA (default datasets/orbit) · ATLAS_DIR (default .atlas) · ATLAS_LEDGER (optional)
#      ATLAS_CRASH_AT (demo only) · ATLAS_LIVE=1 + OPENAI_API_KEY for a real model (optional)
# Services: none; everything is SQLite files under ATLAS_DIR.
import hashlib, json, os, sqlite3
from pathlib import Path
from typing import Any, Callable
from langgraph.checkpoint.sqlite import SqliteSaver

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
ATLAS_DIR = Path(os.environ.get("ATLAS_DIR", ".atlas"))
MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")


def make_checkpointer(path: str | Path) -> SqliteSaver:
    """A SqliteSaver that lives as long as the process, so no `with` block is needed."""
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    # check_same_thread=False because LangGraph may use the connection from worker threads.
    return SqliteSaver(sqlite3.connect(str(path), check_same_thread=False))


def thread_config(thread_id: str, recursion_limit: int = 25) -> dict:
    """The config for every invoke/get_state call. The default limit is 10007, so set your own."""
    return {"configurable": {"thread_id": thread_id}, "recursion_limit": recursion_limit}

# script_id: day_054_langgraph_durability/durable_graph
# built on Day 37 (MockTracker) and Day 35 (tools): compact copies, not imports.
class MockTracker:
    """A fake issue tracker. Each side effect is one JSONL line, so duplicates are easy to see."""
    def __init__(self, path: str | Path):
        self.path = Path(path); self.path.parent.mkdir(parents=True, exist_ok=True)

    def post_comment(self, issue_id: int, body: str, idempotency_key: str | None = None) -> dict:
        # Like most real APIs without key support, this one ignores idempotency_key.
        return self._append({"action": "post_comment", "issue_id": issue_id, "body": body})

    def _append(self, rec: dict) -> dict:
        with self.path.open("a") as f:
            f.write(json.dumps(rec) + "\n")
        return {"ok": True, "comment_no": len(self.actions())}

    def actions(self) -> list[dict]:
        lines = self.path.read_text().splitlines() if self.path.exists() else []
        return [json.loads(line) for line in lines if line]


def load_issue(issue_id: int) -> dict:
    for line in (DATA_DIR / "issues" / "issues.jsonl").read_text().splitlines():
        if json.loads(line)["id"] == issue_id:
            return json.loads(line)
    raise KeyError(issue_id)


def search_issues(query: str, k: int = 3) -> list[dict]:
    """Keyword stand-in for Day 35's search: the issues sharing the most words with the query."""
    words = set(query.lower().split())
    rows = [json.loads(l) for l in (DATA_DIR / "issues" / "issues.jsonl").read_text().splitlines()]
    scored = sorted(rows, key=lambda r: -len(words & set(r["title"].lower().split())))
    return [{"id": r["id"], "title": r["title"]} for r in scored[:k]]


def get_doc_section(doc_id: str, max_chars: int = 400) -> str:
    """Read the start of one Orbit doc page by id, e.g. 'configuration-reference'."""
    path = DATA_DIR / "docs" / f"{Path(doc_id).name}.md"   # .name blocks "../" escapes
    return path.read_text()[:max_chars] if path.exists() else f"ERROR: no doc '{doc_id}'"

# script_id: day_054_langgraph_durability/durable_graph
# built on Day 53: the same three-node graph, now with a checkpointer and a pluggable model.
from typing import Annotated, TypedDict
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages


class AgentState(TypedDict, total=False):
    messages: Annotated[list, add_messages]   # reducer: appends, or replaces a message by id
    issue_id: int
    answer: str


class ScriptedModel:
    """Keyless stand-in for the LLM: search, read one doc, comment, answer. Same .invoke() shape."""
    def __init__(self, doc_id: str = "configuration-reference"):
        self.doc_id = doc_id

    def invoke(self, messages: list) -> AIMessage:
        n = sum(isinstance(m, ToolMessage) for m in messages)
        issue_id = int(messages[0].content.split("#")[1].split(":")[0])
        last_doc = next((c["args"]["doc_id"] for m in reversed(messages)
                         if isinstance(m, AIMessage) for c in m.tool_calls
                         if c["name"] == "get_doc_section"), self.doc_id)
        plan = [("search_issues", {"query": messages[0].content}),
                ("get_doc_section", {"doc_id": self.doc_id}),
                ("post_comment", {"issue_id": issue_id, "body": f"Likely config precedence; see {last_doc}."})]
        if n < len(plan):
            name, args = plan[n]
            return AIMessage(content="", tool_calls=[{"name": name, "args": args, "id": f"call_{n}"}])
        return AIMessage(content=f"Diagnosis for #{issue_id}: see {last_doc}.")

# script_id: day_054_langgraph_durability/durable_graph
def build_agent_graph(tools: dict[str, Callable], *, model: Any = None, checkpointer: Any = None):
    model = model or ScriptedModel()

    def call_model(state: AgentState) -> dict:
        return {"messages": [model.invoke(state["messages"])]}

    def run_tools(state: AgentState) -> dict:
        out = []
        for call in state["messages"][-1].tool_calls:
            fn = tools.get(call["name"])
            result = fn(**call["args"]) if fn else f"ERROR: unknown tool {call['name']}"
            out.append(ToolMessage(json.dumps(result), tool_call_id=call["id"], name=call["name"]))
        return {"messages": out}

    def finalize(state: AgentState) -> dict:
        # .text, not .content: a live Responses API reply's content is a list of blocks.
        return {"answer": str(state["messages"][-1].text)}

    g = StateGraph(AgentState)
    g.add_node("call_model", call_model)
    g.add_node("run_tools", run_tools)
    g.add_node("finalize", finalize)
    g.add_edge(START, "call_model")
    g.add_conditional_edges("call_model",
                            lambda s: "run_tools" if s["messages"][-1].tool_calls else "finalize",
                            ["run_tools", "finalize"])
    g.add_edge("run_tools", "call_model")
    g.add_edge("finalize", END)
    return g.compile(checkpointer=checkpointer)


def start_input(issue_id: int) -> dict:
    issue = load_issue(issue_id)
    return {"issue_id": issue_id, "messages": [HumanMessage(f"Investigate issue #{issue_id}: {issue['title']}")]}

Diagnosis for #96: see configuration-reference. · checkpoints: 10 · tracker lines: 1

# script_id: day_054_langgraph_durability/durable_graph
def effect_key(action: str, args: dict, scope: str = "") -> str:
    """Same action + same args (+ same scope, e.g. a thread_id) gives the same key."""
    blob = json.dumps({"action": action, "args": args, "scope": scope}, sort_keys=True)
    return hashlib.sha256(blob.encode()).hexdigest()[:16]


def run_once(key: str, fn: Callable[[], Any]) -> Any:
    """Run fn at most once per key. A repeat returns the stored result and skips fn."""
    ledger = Path(os.environ.get("ATLAS_LEDGER", ATLAS_DIR / "effects.sqlite"))
    ledger.parent.mkdir(parents=True, exist_ok=True)
    with sqlite3.connect(ledger) as db:
        db.execute("CREATE TABLE IF NOT EXISTS effects (key TEXT PRIMARY KEY, result TEXT)")
        row = db.execute("SELECT result FROM effects WHERE key = ?", (key,)).fetchone()
    if row:
        return json.loads(row[0])
    result = fn()   # a crash after fn() but before the INSERT below still repeats the effect
    with sqlite3.connect(ledger) as db:
        db.execute("INSERT OR REPLACE INTO effects VALUES (?, ?)", (key, json.dumps(result)))
    return result


def crash_if(point: str) -> None:
    """Demo switch: simulate a redeploy at a named point (os._exit skips all cleanup)."""
    if os.environ.get("ATLAS_CRASH_AT") == point:
        os._exit(1)

# script_id: day_054_langgraph_durability/durable_graph
class KeyedTracker(MockTracker):
    """A tracker that honours idempotency keys, like an API with an Idempotency-Key header."""
    def post_comment(self, issue_id: int, body: str, idempotency_key: str | None = None) -> dict:
        if idempotency_key and any(a.get("key") == idempotency_key for a in self.actions()):
            return {"ok": True, "deduplicated": True}
        return self._append({"action": "post_comment", "issue_id": issue_id, "body": body,
                             "key": idempotency_key})


def make_tools(tracker: MockTracker, *, safe: bool = True, scope: str = "") -> dict[str, Callable]:
    def post_comment(issue_id: int, body: str) -> dict:
        """Post a comment on an Orbit issue (a side effect: safe to repeat only via run_once)."""
        key = effect_key("post_comment", {"issue_id": issue_id, "body": body}, scope)
        def effect() -> dict:
            crash_if("before_effect")
            out = tracker.post_comment(issue_id, body, idempotency_key=key)
            crash_if("mid_effect")          # effect done, ledger not written yet
            return out
        result = run_once(key, effect) if safe else effect()
        crash_if("after_tool")              # tool finished, checkpoint not written yet
        return result
    return {"search_issues": search_issues, "get_doc_section": get_doc_section,
            "post_comment": post_comment}

# script_id: day_054_langgraph_durability/durable_graph
def replay_from(cfg: dict, checkpoint_id: str) -> dict:
    """A config that makes app.invoke(None, it) re-run every step AFTER that checkpoint."""
    return {**cfg, "configurable": {**cfg["configurable"], "checkpoint_id": checkpoint_id}}


def make_live_model(tools: dict[str, Callable]) -> Any:
    """Optional real model (needs OPENAI_API_KEY). Imported lazily so the demo stays keyless."""
    from langchain.chat_models import init_chat_model
    return init_chat_model(f"openai:{MODEL}", use_responses_api=True).bind_tools(list(tools.values()))


def build_durable_agent(run_dir: str | Path, *, safe: bool = True, thread_id: str = "issue-96",
                        tracker: MockTracker | None = None):
    run_dir = Path(run_dir)
    tracker = tracker or MockTracker(run_dir / "tracker.jsonl")
    tools = make_tools(tracker, safe=safe, scope=thread_id)
    model = make_live_model(tools) if os.environ.get("ATLAS_LIVE") == "1" else None
    app = build_agent_graph(tools, model=model, checkpointer=make_checkpointer(run_dir / "ckpt.sqlite"))
    return app, thread_config(thread_id), tracker

# script_id: day_054_langgraph_durability/durable_graph
def fresh_run_dir(name: str) -> Path:
    """An empty demo folder with its own effects ledger, so every rerun starts clean."""
    import shutil
    run_dir = ATLAS_DIR / name
    shutil.rmtree(run_dir, ignore_errors=True)
    run_dir.mkdir(parents=True)
    os.environ["ATLAS_LEDGER"] = str(run_dir / "effects.sqlite")   # children inherit it
    return run_dir


if __name__ == "__main__":
    app, cfg, tracker = build_durable_agent(fresh_run_dir("smoke"))
    out = app.invoke(start_input(96), cfg, durability="sync")
    print(out["answer"], "· checkpoints:", len(list(app.get_state_history(cfg))),
          "· tracker lines:", len(tracker.actions()))

# script_id: day_054_langgraph_durability/crash_resume_demo
# Kill the agent mid-run in a child process, restart it, and count the comments it posted.
# Usage: python crash_resume_demo.py            (all four rows)
# Needs durable_graph.py from this lesson next to it. No API key, no services.
import os, subprocess, sys
from durable_graph import (KeyedTracker, MockTracker, build_durable_agent, fresh_run_dir,
                           start_input)

ROWS = [  # name, safe (run_once), crash point, tracker class
    ("naive", False, "after_tool", "MockTracker"),
    ("run_once", True, "after_tool", "MockTracker"),
    ("run_once, crash mid-effect", True, "mid_effect", "MockTracker"),
    ("run_once + keyed API", True, "mid_effect", "KeyedTracker"),
]


def child(run_dir: str, safe: str, tracker_cls: str, phase: str) -> None:
    tracker = {"MockTracker": MockTracker, "KeyedTracker": KeyedTracker}[tracker_cls](
        f"{run_dir}/tracker.jsonl")
    app, cfg, _ = build_durable_agent(run_dir, safe=safe == "1", tracker=tracker)
    if phase == "start":
        app.invoke(start_input(96), cfg, durability="sync")   # sync: each checkpoint hits disk first
    else:
        snap = app.get_state(cfg)
        print(f"   resuming at step {snap.metadata['step']}, next={snap.next}")
        app.invoke(None, cfg, durability="sync")              # None = continue, don't restart


def main(rows: list[tuple] = ROWS) -> None:
    for i, (name, safe, point, tracker_cls) in enumerate(rows):
        run_dir = str(fresh_run_dir(f"crash_{i}"))
        args = [sys.executable, os.path.abspath(__file__), "child", run_dir,
                "1" if safe else "0", tracker_cls]
        died = subprocess.run(args + ["start"], env={**os.environ, "ATLAS_CRASH_AT": point}).returncode
        subprocess.run(args + ["resume"], check=True)
        comments = len(MockTracker(f"{run_dir}/tracker.jsonl").actions())
        print(f"{name:28} crash at {point:10} exit={died} · comments={comments}")


if __name__ == "__main__":
    child(*sys.argv[2:]) if sys.argv[1:2] == ["child"] else main()

   resuming at step 5, next=('run_tools',)
naive                        crash at after_tool exit=1 · comments=2
   resuming at step 5, next=('run_tools',)
run_once                     crash at after_tool exit=1 · comments=1
   resuming at step 5, next=('run_tools',)
run_once, crash mid-effect   crash at mid_effect exit=1 · comments=2
   resuming at step 5, next=('run_tools',)
run_once + keyed API         crash at mid_effect exit=1 · comments=1

# script_id: day_054_langgraph_durability/history_replay
# List a finished run's checkpoints, then replay from the one just before the comment.
# Needs durable_graph.py from this lesson next to it. No API key, no services.
from durable_graph import build_durable_agent, fresh_run_dir, replay_from, start_input

app, cfg, tracker = build_durable_agent(fresh_run_dir("history"))
app.invoke(start_input(96), cfg)

history = list(app.get_state_history(cfg))           # newest first
print(f"{'step':>4}  {'next':<15} msgs  checkpoint_id (tail)")
for snap in reversed(history):
    print(f"{snap.metadata['step']:>4}  {str(snap.next):<15} {len(snap.values.get('messages', [])):>4}"
          f"  {snap.config['configurable']['checkpoint_id'][-8:]}")

# The checkpoint whose next step runs the post_comment tool call.
before_comment = next(s for s in history if s.next == ("run_tools",)
                      and s.values["messages"][-1].tool_calls[0]["name"] == "post_comment")
cid = before_comment.config["configurable"]["checkpoint_id"]
app.invoke(None, replay_from(cfg, cid))                # re-runs run_tools → call_model → finalize
print("replayed from step", before_comment.metadata["step"],
      "· checkpoints now:", len(list(app.get_state_history(cfg))),
      "· tracker lines:", len(tracker.actions()))      # still 1: run_once caught the repeat

step  next            msgs  checkpoint_id (tail)
  -1  ('__start__',)     0  d1f80671
   0  ('call_model',)    1  b2691925
   1  ('run_tools',)     2  f9c93268
   2  ('call_model',)    3  4b23aed8
   3  ('run_tools',)     4  771d2239
   4  ('call_model',)    5  ec0d579b
   5  ('run_tools',)     6  60799de7
   6  ('call_model',)    7  2343edc6
   7  ('finalize',)      8  81f2f407
   8  ()                 8  7b9a7fb7
replayed from step 5 · checkpoints now: 14 · tracker lines: 1

# script_id: day_054_langgraph_durability/time_travel_drill
# A recorded bad run read a 2.x doc for a 3.x issue. Find the step, fork with a fix, compare.
# Needs durable_graph.py from this lesson next to it. No API key, no services.
from langchain_core.messages import AIMessage
from durable_graph import (MockTracker, ScriptedModel, build_agent_graph, fresh_run_dir,
                           make_checkpointer, make_tools, start_input, thread_config)

run_dir = fresh_run_dir("drill")
tracker = MockTracker(run_dir / "tracker.jsonl")
app = build_agent_graph(make_tools(tracker, scope="issue-96"),
                        model=ScriptedModel(doc_id="environment-variables-2x"),   # the bug
                        checkpointer=make_checkpointer(run_dir / "ckpt.sqlite"))
cfg = thread_config("issue-96")
bad = app.invoke(start_input(96), cfg)
print("bad run:  ", bad["answer"])

# 1. Find the checkpoint where the model had just asked for the wrong doc.
history = list(app.get_state_history(cfg))
snap = next(s for s in history if s.next == ("run_tools",)
            and s.values["messages"][-1].tool_calls[0]["name"] == "get_doc_section")
wrong = snap.values["messages"][-1]
print("culprit:   step", snap.metadata["step"], wrong.tool_calls[0]["args"])

# 2. Fork: same message id, so add_messages REPLACES it instead of appending.
fixed = AIMessage(id=wrong.id, content="", tool_calls=[
    {**wrong.tool_calls[0], "args": {"doc_id": "configuration-reference"}}])
fork_cfg = app.update_state(snap.config, {"messages": [fixed]}, as_node="call_model")
good = app.invoke(None, fork_cfg)
print("fork run: ", good["answer"])
print("checkpoints:", len(history), "→", len(list(app.get_state_history(cfg))),
      "· comments:", [a["body"] for a in tracker.actions()])

bad run:   Diagnosis for #96: see environment-variables-2x.
culprit:   step 3 {'doc_id': 'environment-variables-2x'}
fork run:  Diagnosis for #96: see configuration-reference.
checkpoints: 10 → 16 · comments: ['Likely config precedence; see environment-variables-2x.', 'Likely config precedence; see configuration-reference.']

# script_id: day_054_langgraph_durability/postgres_fragment
# fragment
# uv add langgraph-checkpoint-postgres "psycopg[binary]" · needs Postgres (compose.yaml, Day 26)
# Env: DATABASE_URL, e.g. postgresql://user:password@localhost:5432/atlas
import os
from langgraph.checkpoint.postgres import PostgresSaver
from durable_graph import MockTracker, build_agent_graph, make_tools, start_input, thread_config

with PostgresSaver.from_conn_string(os.environ["DATABASE_URL"]) as saver:
    saver.setup()   # once per database: creates the checkpoint tables (safe to call again)
    tools = make_tools(MockTracker(".atlas/pg/tracker.jsonl"), scope="issue-96")
    app = build_agent_graph(tools, checkpointer=saver)
    app.invoke(start_input(96), thread_config("issue-96"), durability="sync")