TL;DR — Production agents need more than conversation history. They must survive restarts, recall decisions from earlier sessions, and work with current business data. In our Supply Chain Planner Copilot, Lakebase handles all three. The key implementation pattern is a Lakebase Postgres query that finds similar quality incidents and joins them to current inventory and open purchase orders in one step. We use Databricks AI Search for document retrieval, Genie Agents for analytics, LangGraph checkpoints for human approval, and MLflow to trace the complete run.
The problem: agent state lives everywhere
While building the Supply Chain Planner Copilot, we ran into three different kinds of state. The agent needed the current conversation, information remembered from earlier sessions, and up-to-date operational data. Those needs are often handled by separate systems.
- Short-term memory is the context for the active session: the conversation so far, intermediate results, tool outputs, and the latest checkpoint. It is what allows an agent to recover from a restart or answer a follow-up such as "What about that SKU?" without starting over.
- Long-term memory is what the agent carries across sessions: user preferences, previous decisions, learned facts, and recurring patterns or queries. It can include both episodic memory (the record of what happened, such as conversations, tool calls, approvals, and feedback) and semantic memory (the reusable facts, rules distilled from those experiences, or previous conversations). For example: "This planner usually expedites safety-critical parts." The semantic portion is what the agent retrieves when it asks, "What happened the last time we saw a similar issue?"
- Operational data is the live business context the agent needs to act: inventory, purchase orders, customers, transactions, and other current records. It is not memory in the same sense, but without it, the agent cannot make useful operational decisions.
A common architecture uses a different system for each concern: a key-value store for session state, a vector database for semantic recall, and a relational or NoSQL database for tool state and operational records.
That separation works, but it introduces cost and complexity. Each system brings another credential set, another failure mode, and another access-control model to maintain. It also creates a more subtle limitation: the semantic index does not naturally operate over the live operational data.
Consider a manufacturing planner investigating a recurring quality issue:
“Henkel (SUP-001) has recurring adhesive cracking on SKU-1001. Show me similar past cases, join them to current inventory and open purchase orders, and recommend a mitigation.”
Answering that question requires both semantic and relational reasoning. The agent must find similar historical incidents, then join those incidents to current inventory and open purchase orders.
When similarity search lives in one system and operational data lives in another, the application usually performs a vector lookup, pulls the matching IDs back into the agent, and then sends a second query to the operational database. The join and re-filtering happen in application code rather than in the database.
That means multiple round trips, multiple governance surfaces, and more logic for the application to coordinate.
In this Databricks architecture, the historical incidents, vector embeddings, agent memory, and operational tables can all live in Lakebase. Semantic similarity becomes part of the SQL query itself, alongside the joins to current operational records.
That gives the planner one governed path from a loosely worded question to a recommendation grounded in current data.
Agent memory on Lakebase: short-term and long-term
Lakebase is Databricks' managed, autoscaling Postgres database, giving agents a durable foundation for both short-term state and long-term semantic memory. Its Postgres compatibility means teams can use familiar clients, ORMs, and memory libraries, while pgvector enables semantic retrieval alongside the rest of the agent's operational data.
This implementation uses LangGraph as the agent framework. Through the databricks-langchain integration, LangGraph can use Lakebase for checkpointing and long-term memory while the integration handles memory-table creation and updates, connection pooling, short-lived OAuth credentials, credential rotation, and vector-index configuration. With two objects, the agent gets both resumable session state and semantic memory across conversations:
from databricks_langchain import AsyncCheckpointSaver, AsyncDatabricksStore
# Short-term: thread/session checkpoints — resumable runs.
async with AsyncCheckpointSaver(
project=PROJECT,
branch=BRANCH,
schema=SCHEMA,
) as checkpointer,
AsyncDatabricksStore(
# Long-term + semantic: pgvector-backed memory, embedded via a Databricks endpoint.
project=PROJECT,
branch=BRANCH,
schema=SCHEMA,
embedding_endpoint="databricks-gte-large-en",
embedding_dims=1024,
) as store:
yield checkpointer, store
Short-term memory - the session checkpoint. Long-running agents need to be able to pick up where they left off. LangGraph checkpoints the state of each run to Lakebase, so if the app restarts midway through, it can resume from the latest checkpoint instead of repeating retrieval calls and other expensive work.
The conversation history is stored there as well. A recent window is passed to the router and planner, allowing follow-up questions like "What about that SKU?" to resolve against the current session. In this implementation, AsyncCheckpointSaver manages that state, with one thread per planning session.
Long-term memory - the semantic store. Some information needs to persist beyond a single conversation: user preferences, previous approvals, supplier notes, and other context the agent may need later. This memory should also be searchable by meaning (i.e. semantics), not only by exact keywords.
In this implementation, AsyncDatabricksStore stores that memory in Lakebase and uses pgvector for semantic retrieval. You point it to a Databricks embedding endpoint, and the integration manages the vector tables and search configuration. Preferences and approvals are scoped by user, while supplier notes are scoped by supplier. Relevant memory is recalled after the gather phase and before planning, then curated memory is written when the plan is committed.
In practical terms:
The checkpointer answers, “What happened in this thread?”
The store answers, “What should the agent remember across threads, and which past memories are relevant now?”
The integration manages storage, but you still decide what deserves to become memory. That includes which facts to keep, how to distill them from a conversation, and when to update or forget them. These choices are application-specific and worth evaluating explicitly. The Databricks AI Research post on memory scaling is a useful starting point.
Figure 1. Three states, one Lakebase. A stateful agent needs short-term checkpoints, long-term semantic memory, and live operational data on one managed Postgres foundation.
Operational joins with vector similarity: similarity as a SQL predicate
The interesting part is not vector search by itself. It is using similarity inside the same SQL statement that joins matching incidents to current inventory and open purchase orders. The operational tables are stored in Lakebase and synchronized from Delta, while pgvector provides vector search directly in Postgres.
SELECT m.incident_id, m.summary, m.supplier_id, m.sku, m.category,
i.on_hand_qty, po.open_po_qty,
round((1 - (m.embedding <=> %(q)s::vector))::numeric, 3) AS similarity
FROM public.quality_incidents m
JOIN public.inventory_current i ON m.sku = i.sku
JOIN public.open_pos po ON m.supplier_id = po.supplier_id AND m.sku = po.sku
WHERE m.expired_at IS NULL
ORDER BY m.embedding <=> %(q)s::vector
LIMIT 5;
The <=> operator calculates vector distance over the quality_incidents table, which uses 1,024-dimensional embeddings, cosine distance, and an HNSW index. The embeddings are generated through the databricks-gte-large-en endpoint.
The relational joins then add current inventory and open purchase-order data in the same statement. The result is a single round trip: Postgres performs the similarity search, applies the relational joins, and returns the operational context together. The join stays in the data layer rather than being reconstructed in application code.
Figure 2. Two hops versus one governed query. When similarity is one predicate in a broader operational query, Lakebase can perform vector search and relational joins in one governed SQL statement and one round trip.
There are two implementation details worth highlighting.
- Operational data reads run through the app service principal. In this reference implementation, the Lakebase query executes under the app service principal using Postgres roles and grants against tables synchronized from Unity Catalog-governed Delta sources. Authenticated users therefore see the same application-authorized operational data. Per-user row-level filtering is outside the scope of the example. In a production system, that could be added with Postgres row-level security based on current_user() or through an entitlements join.
- The executed SQL is part of the agent's result. The operational tool returns the SQL it ran, which makes the query visible in the MLflow trace and available for evaluation. This provides a way to inspect routing and join behavior directly, while keeping numerical results grounded in the data layer rather than generated by the model.
This pattern works best when semantic similarity is one part of a broader operational query, not the entire retrieval workflow.
When to use Lakebase, Databricks AI Search, and Genie Agents?
Using Lakebase as the agent's state layer does not mean every retrieval task should run through Lakebase. This copilot uses three Databricks retrieval surfaces, each for a different kind of question, and routes requests to the one best suited for the job.
|
Dimension |
Databricks AI Search |
Lakebase: SQL + Memory |
Genie Agents |
|
Best suited for |
Retrieving relevant indexed content using semantic, keyword, or hybrid search |
Agent memory and operational queries that combine state, vector similarity, and relational data |
Natural-language analytical questions over curated business data |
|
Primary data shape |
Documents, chunks, images, embeddings, and associated metadata |
Agent state, memories, transactional records, and operational entities |
Structured tables, metrics, business definitions, and analytical datasets |
|
Query pattern |
ANN, full-text, hybrid search, filters, and reranking |
SQL queries containing relational joins, filters, transactions, and vector predicates |
Natural language translated into analytical SQL |
|
Operational joins |
Not the primary retrieval pattern; downstream joins may require an additional SQL or application step |
Native Postgres joins within the same query |
Joins and aggregations generated as part of analytical SQL over configured datasets |
|
Typical result |
Ranked records or chunks with relevance scores and metadata |
Structured rows combining similarity results with live operational context |
Generated SQL, analytical result tables, and visualizations |
The table is the rule of thumb we use. Databricks AI Search handles managed retrieval over large document collections. Lakebase handles durable state and operational queries where similarity belongs inside a broader SQL statement. Genie Agents handle analytical questions such as "What is total unfulfilled demand by product code for Q4?"
The supervisor can route to one or more of those paths, run the selected gather steps in parallel, and combine their results before planning. The point is not to force every question through one retrieval system; it is to match the retrieval path to the shape of the problem.
Making it a real application: durable human review and observability
A planning agent should not automatically execute a $50,000 expedite without review. For decisions that carry financial or operational risk, the application needs a reliable way to pause, collect a human decision, and continue without losing its place.
Because the agent's short-term state is checkpointed in Lakebase, LangGraph's interrupt() can pause the workflow at the review step and resume it when a user approves or rejects the recommendation. The run continues from the saved checkpoint, even if the application restarts while it is waiting.
def hitl_review_node(state: AgentState) -> dict:
rec = state.get("recommendation")
resume_value = interrupt(
{
"type": "approval_request",
"recommendation": rec.model_dump() if rec else None,
"planned_actions": [
a.model_dump() for a in rec.planned_actions
] if rec else [],
"evidence": _evidence_bundle(state),
"prompt": "Approve or reject this recommendation.",
}
)
# The application resumes the graph with Command(resume=...).
decision = _coerce_decision(
resume_value,
state.get("user_id", "unknown"),
)
return {"hitl_decision": decision, ...}
We made three deliberate choices in the review flow.
- We do not ask a model whether approval is required. Spend, irreversible actions, and other known risk conditions are checked with deterministic rules. Informational responses can continue; action-bearing recommendations pause for review.
- Human review has its own node. The interrupt occurs after the selected gather steps have completed and their results have been combined. Isolating the interrupt from parallel execution prevents unrelated sibling tasks from running again when the workflow resumes.
- The workload runs as a Databricks App. Planning runs execute within the application process, while the React interface streams progress from the app and sends a resume command when the user completes review. This model is better suited to workflows that may pause for human input than a long-running synchronous serving request. Durable checkpoints allow an interrupted run to survive application restarts and resume from the review step.
The same architecture also makes the workflow observable. MLflow captures the LangGraph run as a trace, including supervisor routing, parallel retrieval calls, planning, the approval decision, the pause and resume cycle, and the final memory write-back. Model usage, latency, and cost are recorded at the span level.
The operational tool also returns the SQL it executed, so the query appears alongside the rest of the trace. This makes it possible to evaluate not only the final recommendation, but also whether the request was routed correctly, whether the operational join returned the right evidence, and whether the recommendation was grounded in that evidence.
Figure 3. Durable agent loop. The supervisor routes the question, selected tools gather in parallel, memory is recalled before planning, and deterministic rules pause action-bearing recommendations for human review. Lakebase persists the loop while MLflow traces it end to end.
Try the reference application and make it your own
The repository contains the working Supply Chain Planner Copilot: the multi-agent graph, React interface, Lakebase state and memory layer, operational data, Databricks AI Search pipeline, Genie Agent, approval flow, and MLflow tracing used in this post.
Deploy the complete application
Clone the lakebase-for-ai-developers repository and deploy it to a Databricks workspace with one command:
make deploy PROFILE=<your-profile>The deploy script creates the Databricks App, MLflow experiment, Lakebase project, Genie Agent, and demo-data setup job. It then builds the application, seeds the reference data, and runs verification. The script is idempotent and supports a clean workspace.
Once deployed, follow the demo path:
- Ask an analytical question to see the supervisor route to Genie Agents.
- Ask for similar quality incidents to inspect the Lakebase vector-and-SQL query.
- Run the full planning scenario and approve or reject the recommendation.
- Start a new conversation and verify that the agent recalls the previous decision.
Together, those steps exercise routing, the operational join, durable approval, memory write-back, and cross-session recall.
Customize it with embedded skills
The repository includes project-level skills that give AI coding assistants context about the architecture, conventions, and deployment patterns.
Use them to adapt the reference application to your own domain, data, tools, and workflows without requiring the assistant to infer how the project is structured from the code alone.
Start from the reusable template
For a smaller starting point, use the agent-langgraph-advanced application template. It provides the foundation for running LangGraph on Databricks Apps with Lakebase-backed memory and MLflow tracing.
The reference application builds on that template and adds the manufacturing use case, multi-agent routing, operational vector-and-SQL retrieval, Genie Agents, Databricks AI Search, a React and Vite interface, and the full deployment workflow.
Adapt the pattern to your own data
The manufacturing scenario is only one example. The core architecture transfers to other applications where an agent needs to:
- maintain state during a long-running workflow;
- remember decisions and preferences across sessions;
- retrieve similar historical cases;
- join those cases to current operational records;
- pause safely for human approval; and
- trace how evidence became a recommendation.
A useful test for your own architecture is:
Can the agent find similar historical cases, join them to live operational data, and return the result through a single governed query path?
When semantic similarity is one part of a broader operational query, Lakebase provides a natural place to bring agent memory, vector search, and relational data together. Databricks AI Search remains the better fit for managed retrieval over large knowledge corpora, while Genie Agents handle natural-language analysis over curated business data.
Industry use cases for this pattern
This pattern also appears in warranty operations, where an agent can find similar failures and check coverage, repair history, and parts availability before recommending a repair or replacement. In procurement, the same loop can combine prior contract exceptions with current spend, renewal dates, and supplier performance. Field-service dispatches, account credits, and policy exceptions follow a similar shape: retrieve the relevant history, check the current operational state, and pause before committing an expensive or exceptional action.
Explore the supporting resources
The result is a workflow that can remember what happened, reason over what is happening now, and show how it reached a recommendation.
1 Comment
