diff --git a/.env.example b/.env.example index ac984ef..8cc448e 100644 --- a/.env.example +++ b/.env.example @@ -21,6 +21,9 @@ EMBEDDING_DIM=768 DATABASE_URL=postgresql+psycopg://agentforge:agentforge@localhost:5432/agentforge COLLECTION_NAME=agentforge_documents +# Graph checkpointer: "memory" (in-process) or "postgres" (durable HITL/resume). +CHECKPOINT_BACKEND=memory + # RAG retrieval tuning. RETRIEVAL_K=4 # Normalized relevance floor [0,1] (higher = stricter). Calibrate per embedding diff --git a/agentforge/agents/checkpoint.py b/agentforge/agents/checkpoint.py new file mode 100644 index 0000000..9226a39 --- /dev/null +++ b/agentforge/agents/checkpoint.py @@ -0,0 +1,41 @@ +"""Durable LangGraph checkpointer backed by Postgres. + +Used when ``CHECKPOINT_BACKEND=postgres``. A persistent checkpointer makes +human-in-the-loop pauses survive process restarts and resume correctly no matter +which replica handles ``/approve`` (the in-memory ``MemorySaver`` could not). +""" + +from __future__ import annotations + +from functools import lru_cache + +from agentforge.config import get_settings + + +def _conninfo(database_url: str) -> str: + """Turn a SQLAlchemy URL into a libpq conninfo string. + + The app uses ``postgresql+psycopg://…`` for langchain/SQLAlchemy; psycopg's + pool wants a plain ``postgresql://…``. + """ + return database_url.replace("postgresql+psycopg://", "postgresql://", 1) + + +@lru_cache +def get_postgres_checkpointer(): + """Process-wide PostgresSaver over a connection pool (tables created once).""" + from langgraph.checkpoint.postgres import PostgresSaver + from psycopg.rows import dict_row + from psycopg_pool import ConnectionPool + + pool = ConnectionPool( + conninfo=_conninfo(get_settings().database_url), + max_size=20, + # PostgresSaver requires autocommit; unprepared statements avoid clashes + # with PgBouncer-style poolers. + kwargs={"autocommit": True, "prepare_threshold": 0, "row_factory": dict_row}, + open=True, + ) + checkpointer = PostgresSaver(pool) + checkpointer.setup() + return checkpointer diff --git a/agentforge/agents/graph.py b/agentforge/agents/graph.py index c947c1f..824899f 100644 --- a/agentforge/agents/graph.py +++ b/agentforge/agents/graph.py @@ -4,8 +4,9 @@ durable execution (a crash mid-run resumes from the last node) and the pause/ resume mechanics that human-in-the-loop approval relies on. -For production, swap ``MemorySaver`` for ``langgraph.checkpoint.postgres.PostgresSaver`` -so checkpoints survive process restarts — same interface, persistent store. +The backend is configurable: ``MemorySaver`` in-process by default (dev/tests), +or a durable ``PostgresSaver`` when ``CHECKPOINT_BACKEND=postgres`` (see +``checkpoint.py``) so checkpoints survive restarts and work across replicas. """ from __future__ import annotations @@ -50,5 +51,17 @@ def build_graph() -> StateGraph: @lru_cache def get_compiled_graph(): - """Compiled, checkpointed graph (cached as a process-wide singleton).""" - return build_graph().compile(checkpointer=MemorySaver()) + """Compiled, checkpointed graph (cached as a process-wide singleton). + + Uses a durable Postgres checkpointer when ``CHECKPOINT_BACKEND=postgres`` + (set by docker-compose / k8s), otherwise the in-process ``MemorySaver``. + """ + from agentforge.config import get_settings + + if get_settings().checkpoint_backend == "postgres": + from agentforge.agents.checkpoint import get_postgres_checkpointer + + checkpointer = get_postgres_checkpointer() + else: + checkpointer = MemorySaver() + return build_graph().compile(checkpointer=checkpointer) diff --git a/agentforge/config.py b/agentforge/config.py index 550049f..eb4a871 100644 --- a/agentforge/config.py +++ b/agentforge/config.py @@ -45,6 +45,10 @@ class Settings(BaseSettings): default="postgresql+psycopg://agentforge:agentforge@localhost:5432/agentforge" ) collection_name: str = Field(default="agentforge_documents") + # Graph checkpointer backend: "memory" (in-process, dev/tests) or "postgres" + # (durable — HITL state survives restarts and is shared across replicas). + # docker-compose / k8s set this to "postgres". + checkpoint_backend: str = Field(default="memory") # --- RAG retrieval --------------------------------------------------- retrieval_k: int = Field(default=4) diff --git a/deploy/k8s/api.yaml b/deploy/k8s/api.yaml index 3374b31..b49c994 100644 --- a/deploy/k8s/api.yaml +++ b/deploy/k8s/api.yaml @@ -7,9 +7,10 @@ metadata: app.kubernetes.io/name: api app.kubernetes.io/part-of: agentforge spec: - # Single replica: HITL approval state lives in the in-process MemorySaver - # checkpointer, so a paused thread must resume on the same pod. Switch the - # graph to PostgresSaver (see agents/graph.py) before scaling this up. + # Single replica for now. Resume IS durable across pods (PostgresSaver via + # CHECKPOINT_BACKEND=postgres), but the /approvals queue *listing* is still an + # in-process registry, so it would differ per pod under load balancing. Scale + # up once the queue is store-backed. replicas: 1 selector: matchLabels: diff --git a/deploy/k8s/config.yaml b/deploy/k8s/config.yaml index ab2195d..8f2a587 100644 --- a/deploy/k8s/config.yaml +++ b/deploy/k8s/config.yaml @@ -18,5 +18,7 @@ data: MIN_RELEVANCE: "0.2" RETRIEVAL_K: "4" OBSERVABILITY_BACKEND: "none" + # Durable HITL/resume state in Postgres (survives restarts, shared across pods). + CHECKPOINT_BACKEND: "postgres" # Production ingests via a one-shot Job, not at API startup. AUTO_INGEST: "false" diff --git a/docker-compose.yml b/docker-compose.yml index 85127bb..1cfbaaa 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -71,6 +71,7 @@ services: ANTHROPIC_API_KEY: ${ANTHROPIC_API_KEY:-} OPENAI_API_KEY: ${OPENAI_API_KEY:-} OBSERVABILITY_BACKEND: ${OBSERVABILITY_BACKEND:-none} + CHECKPOINT_BACKEND: postgres AUTO_INGEST: "true" ports: - "8000:8000" diff --git a/docs/SMOKE_TEST.md b/docs/SMOKE_TEST.md index c53e1a8..b8e87c3 100644 --- a/docs/SMOKE_TEST.md +++ b/docs/SMOKE_TEST.md @@ -108,6 +108,12 @@ curl -s localhost:8000/approve -H 'content-type: application/json' \ - [ ] `answer` confirms the action ran (contains `SAR drafted`). +Durability (optional): pause on a SAR, run `docker compose restart api`, then +`/approve` that `thread_id`. + +- [ ] The resume still works after the restart — checkpoints persist in Postgres + (`CHECKPOINT_BACKEND=postgres`). + --- ## 6. Streaming (SSE) diff --git a/pyproject.toml b/pyproject.toml index c67b8c7..5b87ca3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,9 @@ dependencies = [ "langchain>=1.0,<2.0", "langchain-text-splitters>=0.3", "langgraph>=1.0,<2.0", + # Durable checkpointer (HITL + resume survive restarts / work across replicas). + "langgraph-checkpoint-postgres>=2.0", + "psycopg-pool>=3.2", "langchain-anthropic>=0.3", "langchain-openai>=0.2", "langchain-ollama>=0.2", diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py new file mode 100644 index 0000000..c4ef6d0 --- /dev/null +++ b/tests/test_checkpoint.py @@ -0,0 +1,22 @@ +"""Checkpointer selection + conninfo handling. No DB connection is opened.""" + +from __future__ import annotations + +from agentforge.agents.checkpoint import _conninfo + + +def test_conninfo_strips_sqlalchemy_driver(): + url = "postgresql+psycopg://user:pass@host:5432/agentforge" + assert _conninfo(url) == "postgresql://user:pass@host:5432/agentforge" + + +def test_conninfo_only_rewrites_scheme(): + # A password that happens to contain the scheme text must be left intact. + url = "postgresql+psycopg://u:postgresql+psycopg@host/db" + assert _conninfo(url) == "postgresql://u:postgresql+psycopg@host/db" + + +def test_default_checkpoint_backend_is_memory(): + from agentforge.config import Settings + + assert Settings().checkpoint_backend == "memory"