Building a Knowledge Assistant Agent on Databricks With LangGraph

A basic custom agent on Databricks can go from a bare LLM call to a knowledge assistant that searches documents and remembers users, and each step is small. The finished agent is built with the langgraph framework. It has a vector search component for looking up text and a memory module, so it can recall what an end user said both within a conversation and across separate ones. Once it runs, evaluation comes next: measuring performance during development and again in production.

The walkthrough follows a set of notebooks, grouped into folders for foundations, the RAG pipeline and memory.

Preparing the workspace

A full agent comes with vector indexes, tools and logs, and those need a tidy home first. On Databricks the governance layer is Unity Catalog (UC), which manages every asset under a three-tier namespace. Running the notebooks in the 00_foundations folder in order creates all the assets the later steps rely on, including a mix of structured and unstructured data that the agent will connect to.

Memory will live in Lakebase, an autoscaling, Postgres-compatible database. So that the tables can be created smoothly later, the setup notebook also creates a Lakebase project and a branch to hold them.

A first, minimal agent

Next come the notebooks under 01_rag_pipeline, starting with 0_your_first_agent_on_databricks.

Every agent is built around a foundation model. Databricks Model Serving lists the available ones, and the list may differ depending on how your workspace is configured. Endpoints are compatible with the OpenAI API format, but for tight langgraph integration the official databricks-langchain library is the better choice.

A brief word on langgraph. A production agent is seldom a single prompt. It usually contains branching logic that depends on the question and on how the user answers, plus loops for consulting external systems and rephrasing queries based on what comes back. Langgraph organizes this by modeling the agent as a graph. Nodes hold functions and LLM calls, and edges decide how context moves between them. The starting graph is very simple.

Invoke the agent and a trace shows up, listing the input question, the model's output and every logic step in between. That comes from MLflow, which integrates natively with langgraph. A single call to mlflow.langchain.autolog() at the top of the notebook turns these logs and diagnostics on. MLflow gets a fuller treatment later.

This first version just passes the user's question to the model. A real agent adds a system prompt to set the tone, shape the output and supply background context, such as what the assistant is for and what it should refuse to answer. So the second iteration inserts a system prompt into the graph.

Adding retrieval

One of the most common agent types is the knowledge retrieval agent, also called a RAG agent. It helps people search a pile of documents and answer questions about them. To give an agent this ability, it needs a vector search index, essentially a database for unstructured text.

The notebook 01_building_a_doc_store_on_vector_search covers building one. Before an index can be attached to the agent, it has to exist. This part sticks to the core steps and skips tuning and optimization.

Building an index takes three actions:

• Parsing means opening a document and reading its content.
• Chunking means using a splitter to break the document into sections.
• Embedding means converting each chunk into a semantic vector stored in the index.

Each has its own subtleties, and mistakes in an early step quietly degrade everything downstream. The tutorial uses defaults throughout. In production, different documents may need different parsers to extract text correctly, and different use cases may call for other chunk sizes for text, tables and figures, along with different embedding models.

Databricks supplies tooling to make this easier. Volumes offer a cloud-native, fully governed object store for raw files. The ai_parse_documents function extracts text, figures and tables with little configuration and scales well. Spark UDFs can split extracted text into chunks at various granularities, ready to be embedded by any of the provided embedding models.

Once embeddings exist, similarity search needs a Vector Search Index. A new data store sounds like extra work, but Delta Sync handles it by embedding and indexing chunks automatically as they're added, changed or deleted. There's no new infrastructure to administer and no sync job to maintain between the chunk table and the index.

In the first two steps of the notebook, the markdown files produced during setup are ingested and split into overlapping 800-character chunks, then loaded into Databricks Vector Search. The databricks_langchain library includes a ready-made VectorSearchRetrieverTool that connects the index to the langgraph agent.

With the index up and tested, step three redefines the agent. It looks much like the earlier ones, with one addition: a ToolNode, which takes a list of tools. Here that list holds only the VectorSearchRetrieverTool.

In the resulting flow, a question enters and passes through the prompt template. The model sees the prompt along with the tools and their descriptions, and decides whether to consult the index. It can do so repeatedly, and it replies to the user only once it has found what it was seeking.

That's a working RAG agent. The graph is still compact, with the tool node as the only addition, which shows how little structure a retrieval loop needs. A knowledge assistant needs one more piece, though: memory.

Giving the agent memory

As built so far, the agent has no memory at all. A follow-up question that refers back to the previous answer simply fails, because the model has no idea what was said a moment ago.

For the agent to account for earlier turns, langgraph needs a checkpointer that stores the chat history. Checkpointers range from plain text files to full databases. For scalability, latency and security a database is the better pick, and on Databricks that means Lakebase.

Memory comes in two kinds. Short-term memory covers only the current active session. Long-term memory retains facts across sessions. Nearly every application needs the first, and whether you need the second depends on the use case.

Short-term memory

The 02_memory folder opens with 01_short_term_memory, which adds a checkpointer. Steps 1 to 3 build a StateGraph identical to the previous one, and step 4 shows how awkward a conversation gets without memory.

Creating a checkpointer tied to Lakebase restores continuity, and databricks_langchain offers a native CheckpointSaver class to instantiate and attach to the graph.

An agent is a service shared by many users, so saved conversations need an owner. The thread_id does that job. In a production app, the front end should generate a unique thread_id per active session, so each user sees only their own chat history.

By step 6 the checkpointer is in place and the conversation flows naturally.

Long-term memory

Short-term memory disappears when the session ends. To remember a user and their preferences, the agent needs long-term memory, and whether to add it depends on the product. A long-lived personal assistant should remember earlier conversations, facts such as the user's name and preferences like short, sharp replies versus detailed explanations. A help service on a customer portal may be better off not keeping those details, since that would mean carrying the burden of securing, encrypting and managing personal data for a long time.

Short-term memory only needed a CheckpointSaver. Long-term memory requires some design. Recording every conversation for posterity usually isn't wanted. What's worth keeping are facts, like a name or a hometown, and preferences.

The notebook 02_long_term_memory goes through the mechanics. Steps 1 to 4 rebuild the RAG agent with short-term memory, and step 5 shows that preferences vanish when the session and thread_id change.

Step 6 introduces a DatabricksStore object that persists data into Lakebase. That part's easy. Deciding what to persist is the real design question. Step 7 takes a basic approach: a handful of hardcoded facts are written into the store.

To use them, step 8 recreates the agent. Earlier, the agent node was a plain function that called invoke on the foundation model. Now it's a more involved function. It takes a user ID (hardcoded for the demo), fetches that user's facts from the DatabricksStore and concatenates them into the context sent to the model. This pattern clearly needs careful design to hold up.

A more flexible approach lets the model itself decide when to save, delete and fetch memories. Step 12 implements a basic version: three new tools interact with the memory module, and the system prompt changes to tell the model about them and encourage their use.

The result is a complete assistant, and each capability can be tested on its own because it arrived through a separate notebook and a separate change to the graph. It searches documents through a vector index, tracks the current conversation with short-term memory and remembers details about users across sessions with long-term memory.

Where it goes next

The path ran from a bot that passed questions straight to an LLM, to one with a prompt template, short-term and long-term memory and Vector Search for context. The follow-on work is productionizing it: evaluating the quality of the agent, connecting more functionality and deploying it as a service for end users.