AI Memory — durable agent memory on Lakebase

Give an agent persistent, queryable memory backed by a single governed Lakebase Postgres store. Three layers share one database: short-term conversation history, long-term facts, and semantic recall over those facts with pgvector. A minimal Chainlit chat app demonstrates all three; the reusable code lives in agent_memory/.

Features

Layer	What it stores	How it’s used
Short-term	Conversation threads (Chainlit’s SQLAlchemy data layer)	Resumable thread sidebar + cross-session history, backed by Lakebase
Long-term facts	Durable user facts (“prefers metric units”)	Survive across sessions (`remember` / `list_memories`)
Semantic recall	A `pgvector` embedding of each fact	Cosine-similarity lookup of the most relevant memories (`recall`)

App owns its schema. The app service principal creates (and therefore owns) its tables on startup, so schema access survives redeploys without re-granting table privileges after each bundle deploy.
On-platform embeddings. A Databricks Foundation Model endpoint (databricks-bge-large-en by default) produces embeddings — no third-party API key.
Fully bundle-driven. Every workspace-specific value is a DAB variable.

Architecture

User --> Chainlit App (Databricks App)
             |                           +-----------------------------------+
             |  threads/steps (history)  |  Lakebase Postgres (schema: aimem)|
             |<------------------------->|  users/threads/steps/...          |
             |  recall(query) /          |         (Chainlit data layer)     |
             |  remember(fact)           |  memories + pgvector              |
             |-------------------------->|         (long-term + recall)      |
             |  chat + embed             +-----------------------------------+
             v
   Databricks Foundation Models

Short-term history uses Chainlit’s SQLAlchemy data layer (the users/threads/steps tables) for the resumable-thread sidebar; long-term facts + recall use the memories table with a pgvector column. Everything lives in a dedicated aimem schema, created and owned by the app service principal, which mints short-lived OAuth tokens (auto-refreshed) as the Postgres password. Chat and embeddings use Databricks Foundation Model serving endpoints.

The memories.embedding column is indexed with HNSW (vector_cosine_ops, m = 16, ef_construction = 64) — the same index type and parameters as the GraphRAG example. HNSW rather than IVFFlat because this schema is created at application startup, against an empty table: IVFFlat derives its lists from a k-means step over the rows present at build time, so building it empty clusters poorly and degrades recall as the table grows. HNSW has no training step and can be built empty.

Because CREATE INDEX IF NOT EXISTS matches on name only, the HNSW index uses a new name and the legacy memories_embedding_idx is dropped by name — otherwise an existing database would keep its IVFFlat index forever. That is the same migration the Genie caching stores use, so an existing database converges on HNSW at the next start with no operator step. On a large existing table the rebuild is worth doing out of band first; the README has the CONCURRENTLY recipe and the measured cost.

The index contract is checked offline by smoketest/aimem_index_smoketest.py (no Lakebase or model endpoint needed): the index type and opclass, that the switch actually takes effect on an existing database, parameter parity with the GraphRAG schema, re-runnable DDL, and statement ordering.

Deploy

cd agents/ai_memory
databricks bundle validate -t demo
databricks bundle deploy -t demo \
  --var lakebase_branch="projects/<project>/branches/<branch>" \
  --var lakebase_database="projects/<project>/branches/<branch>/databases/<id>" \
  --var lakebase_instance="<your-lakebase-instance-name>"

Chat with the app; say remember: I prefer metric units to store a long-term fact, and later questions will recall it automatically.

Configuration

Variable	What it does	Default
`lakebase_branch`	Lakebase project/branch path	`projects/CHANGE_ME/branches/production`
`lakebase_database`	Full Lakebase database resource path	`.../databases/CHANGE_ME`
`lakebase_instance`	Lakebase database instance name (mints OAuth credentials)	`CHANGE_ME`
`lakebase_catalog` / `lakebase_schema`	UC catalog / Postgres schema	`default` / `public`
`embedding_endpoint`	Foundation Model embedding endpoint	`databricks-bge-large-en`

Edit this page →