feat(kg): KnowledgeGraphAGE skeleton + graph bootstrap

jphein · claude · jphein · commit a3ee6231f400 · 2026-05-11T08:22:25.000-07:00
First commit of the AGE-backed knowledge graph layer. Skeleton class in `mempalace/knowledge_graph_age.py` that opens a Postgres connection, loads the AGE extension, sets `search_path = ag_catalog, "$user", public` for the session, and creates a graph named `mempalace_kg` in `ag_catalog.ag_graph` if one doesn't already exist. Idempotent bootstrap; safe to instantiate repeatedly. Three `pytest.skipif`-gated integration tests in `tests/test_knowledge_graph_age.py`: - `test_age_kg_instantiates` — class constructs cleanly, closes without exception. - `test_age_graph_created` — `mempalace_kg` is registered in `ag_catalog.ag_graph` with a non-null `graphid` after init. - `test_age_context_manager` — `with KnowledgeGraphAGE(...) as kg:` pattern closes the connection on exit (verifies `_conn.closed` is True after). Verified end-to-end against the live mempalace-db container on disks.jphe.in:5433 (PG16 + pgvector 0.8.2 + AGE 1.6.0): all 3 AGE tests pass; full suite 1854 passed / 1 skipped / 106 deselected with TEST_POSTGRES_DSN set — zero regressions vs the 1851 pre-AGE baseline (+3 AGE tests). Adapted plan code from psycopg3 to psycopg2 for consistency with MemPalace#665's dependency choice (psycopg2-binary is what `pip install -e ".[postgres]"` brings in; introducing psycopg3 alongside would duplicate the driver surface for no real gain at this stage). Future commits in this layer: triple add/query operations, temporal filtering (`as_of` queries), and a `MempalaceConfig.kg_backend` routing flag selecting AGE over the SQLite default. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
diff --git a/mempalace/knowledge_graph_age.py b/mempalace/knowledge_graph_age.py
@@ -0,0 +1,76 @@
+"""AGE-backed implementation of KnowledgeGraph (Apache AGE on Postgres).
+
+Companion to `mempalace.knowledge_graph.KnowledgeGraph` (SQLite). Selectable
+via `MEMPALACE_KG_BACKEND=age` once the config-routing layer is wired up.
+Mirrors the public interface of the SQLite KG so callers can swap backends
+without code changes.
+
+The graph itself is `mempalace_kg` registered in AGE's `ag_catalog`. It is
+created on first init and reused thereafter — initialization is idempotent.
+"""
+
+import psycopg2
+
+
+class KnowledgeGraphAGE:
+    """Cypher-queryable KG using Apache AGE on a Postgres connection.
+
+    Currently a skeleton: instantiation + graph bootstrap only. Triple
+    add/query operations and temporal filtering arrive in subsequent
+    commits; the routing layer that lets `MempalaceConfig.kg_backend`
+    select this backend over the SQLite one is its own change.
+    """
+
+    GRAPH_NAME = "mempalace_kg"
+
+    def __init__(self, dsn: str):
+        """Open a Postgres connection and ensure `mempalace_kg` exists.
+
+        Args:
+            dsn: PostgreSQL DSN. Must point at a database where the AGE
+                extension is installed (CREATE EXTENSION succeeds). The
+                ``apache/age:release_PG16_1.6.0`` image we deploy on the
+                homelab already has the .so files baked in; bare-metal
+                Postgres requires source-build of AGE first.
+        """
+        self._conn = psycopg2.connect(dsn)
+        # KG writes need explicit commit semantics, not autocommit — keep
+        # the same shape as the SQLite KG so the eventual unified write
+        # API can be swapped underneath. The bootstrap below commits its
+        # own changes; subsequent write operations control their own
+        # transactions.
+        self._conn.autocommit = False
+        self._ensure_graph()
+
+    def _ensure_graph(self) -> None:
+        """Idempotent: load AGE, set search_path, create the graph if absent.
+
+        Both `LOAD 'age'` and the `search_path` setting are session-scoped
+        — anything new that takes a fresh cursor on this connection must
+        re-run them before issuing Cypher. The pattern in this module is
+        to always wrap `LOAD 'age'; SET search_path = ag_catalog, "$user",
+        public` around each cypher block; for the bootstrap that's done
+        here, downstream methods will repeat as needed.
+        """
+        with self._conn.cursor() as cur:
+            cur.execute("CREATE EXTENSION IF NOT EXISTS age")
+            cur.execute("LOAD 'age'")
+            cur.execute('SET search_path = ag_catalog, "$user", public')
+            cur.execute(
+                "SELECT graphid FROM ag_catalog.ag_graph WHERE name = %s",
+                (self.GRAPH_NAME,),
+            )
+            if cur.fetchone() is None:
+                cur.execute("SELECT create_graph(%s)", (self.GRAPH_NAME,))
+        self._conn.commit()
+
+    def close(self) -> None:
+        """Close the underlying Postgres connection."""
+        if self._conn and not self._conn.closed:
+            self._conn.close()
+
+    def __enter__(self) -> "KnowledgeGraphAGE":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close()
diff --git a/tests/test_knowledge_graph_age.py b/tests/test_knowledge_graph_age.py
@@ -0,0 +1,60 @@
+"""Integration tests for the AGE-backed KnowledgeGraph implementation.
+
+Requires a live Postgres with Apache AGE installed. Set TEST_POSTGRES_DSN
+to point at one (e.g. the homelab `mempalace-db` container documented in
+`scratch/postgres-preflight-2026-05-10.md`); skipped by default so the
+suite stays green on machines without a postgres at hand.
+
+Pairs with `mempalace/knowledge_graph_age.py`. The classic SQLite-backed
+`KnowledgeGraph` in `mempalace/knowledge_graph.py` stays the default;
+the AGE backend is opt-in via `MEMPALACE_KG_BACKEND=age` once the
+config-routing layer is wired.
+"""
+
+import os
+
+import pytest
+
+POSTGRES_DSN = os.environ.get("TEST_POSTGRES_DSN")
+pytestmark = pytest.mark.skipif(
+    POSTGRES_DSN is None,
+    reason="set TEST_POSTGRES_DSN to run AGE knowledge-graph tests",
+)
+
+
+def test_age_kg_instantiates():
+    """KnowledgeGraphAGE opens a connection and exits cleanly."""
+    from mempalace.knowledge_graph_age import KnowledgeGraphAGE
+
+    kg = KnowledgeGraphAGE(dsn=POSTGRES_DSN)
+    assert kg is not None
+    kg.close()
+
+
+def test_age_graph_created():
+    """`mempalace_kg` graph is registered in AGE's catalog after init."""
+    from mempalace.knowledge_graph_age import KnowledgeGraphAGE
+
+    kg = KnowledgeGraphAGE(dsn=POSTGRES_DSN)
+    try:
+        with kg._conn.cursor() as cur:
+            cur.execute(
+                "SELECT graphid FROM ag_catalog.ag_graph WHERE name = %s",
+                (KnowledgeGraphAGE.GRAPH_NAME,),
+            )
+            row = cur.fetchone()
+            assert row is not None, "mempalace_kg graph should exist after init"
+            assert row[0] is not None, "graph should have a non-null graphid"
+    finally:
+        kg.close()
+
+
+def test_age_context_manager():
+    """`with KnowledgeGraphAGE(...) as kg:` closes the conn on exit."""
+    from mempalace.knowledge_graph_age import KnowledgeGraphAGE
+
+    with KnowledgeGraphAGE(dsn=POSTGRES_DSN) as kg:
+        assert kg._conn is not None
+        assert not kg._conn.closed
+    # After the with block, the connection should be closed.
+    assert kg._conn.closed
