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.

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 →