Fast, Disposable RAG: Session-Scoped Document Q&A on Databricks

By: Raj Tulluri

Choosing Lakebase over Vector Search for per-chat document upload 

The Problem 

One of the most common requirements for an agentic chatbot is document question answering: upload a file, ask questions about it, and receive answers grounded in that file. This is the same pattern popularized by ChatGPT and Claude, and it has become a baseline expectation for any chat interface. 

This kind of document Q&A is almost always scoped for a single session. A contract gets uploaded, three questions get asked about a termination clause, and the conversation ends. The document has no life beyond that session. Nobody expects it to remain searchable afterward, and nobody wants it surfacing in someone else’s chat. 

The standard Databricks answer to a RAG requirement is Vector Search: parse the documents, chunk them, embed them, sync the result into a managed index, and query it. That pattern suits a durable, shared knowledge base meant to live for months or years. It fits poorly when the content being indexed exists only for the length of one conversation. Provisioning a Vector Search index per chat session does not match how the service is designed to operate, and it does not match the actual requirement: something that can be created in seconds, queried immediately, and left idle at close to no cost. 

That mismatch is the subject of this post. The goal was a document pipeline confined to a session, fast to ingest, fast to query, isolated per chat by construction, and inexpensive when idle. The approach tested here uses Lakebase, Databricks’ managed autoscaling Postgres, as a per-session vector store instead of routing document upload through Vector Search. 

The Starting Architecture 

Before getting into the document pipeline, it is worth defining the two pieces of infrastructure it is built on: the agent itself, and Lakebase. 

The Agent as a Databricks App 

A Databricks App runs custom application code directly on Databricks compute, exposed through the platform’s own authentication and networking. It is a different deployment path from Model Serving, which registers a model and serves it behind a fixed prediction interface. An App is a real, long-running process: it can hold open connections, run background work, and serve more than one route. The agent as a Databricks App rather than a Model Serving endpoint makes it stateful where a multi-route feature like session document upload is straightforward to add. 

@invoke() 
async def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentResponse: 
    # run your agent here and return the full response 
    ... 

@stream() 
async def stream_handler(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]: 
    # run your agent here and yield events as they are produced 
    ... 

agent_server = AgentServer("ResponsesAgent", enable_chat_proxy=True)
app = agent_server.app

In our Disposable RAG example, the agent itself is a LangGraph application, wrapped by MLflow’s agent_server framework. Instead of subclassing MLflow’s ResponsesAgent interface directly, agent_server exposes it through two decorators, @invoke and @stream, applied to plain async functions. Figure 1 shows the resulting request path: a client request lands on the app’s compute, is picked up by the MLflow Agent Server, and reaches the agent code through the ResponsesAgent interface. From there the agent calls out to a model serving endpoint for inference, and to whatever other resources it needs, treated uniformly as tools: Genie spaces, MLflow experiments, Unity Catalog assets, MCP servers, and so on.

This is the layer the document pipeline in this post plugs into. The upload handler and the retrieval tool both run inside the same compute as the conversational agent, rather than as separate services. 

Lakebase for Memory 

Lakebase is Databricks’ managed, autoscaling Postgres offering. It behaves like an ordinary Postgres instance for connection purposes, while the underlying compute suspends when idle and resumes on demand.  

Before session-scoped document RAG existed as a feature, Lakebase already held the agent’s own conversational memory. A LangGraph checkpointer stores short-term thread history, so a conversation can resume mid-thread. A LangGraph store holds long-term facts about a user that persist across separate conversations. Both run against the same autoscaling Postgres instance, through the same connection pooling and credential rotation.

from databricks_langchain import AsyncCheckpointSaver, AsyncDatabricksStore 

lakebase_kwargs = { 
    "project": "my-lakebase-project", 
    "branch": "production", 
    "schema": "agent_memory", 
} 

# Short-term memory: conversation history, keyed by thread id 
checkpointer = AsyncCheckpointSaver(**lakebase_kwargs) 
await checkpointer.__aenter__() 
await checkpointer.setup() 

# Long-term memory: durable facts about a user, with semantic search 
store = AsyncDatabricksStore( 
**lakebase_kwargs, 
embedding_endpoint=”databricks-gte-large-en", 
embedding_dims=1024 
) 
await store.__aenter__() 
await store.setup() 

The RAG Setup 

Figure 2 shows the full path this takes. A chat message and a file upload are two separate entry points from the client, one landing on the agent itself, the other on a dedicated upload handler, but both routes end up reading from and writing to the same Lakebase instance.

Selecting a file starts the upload immediately, independent of whatever is being typed into the message box at that point. 

Ingestion 

The pipeline follows a familiar shape once a file is selected: 

1. The file is parsed and split into chunks. Chunk size and overlap depend on the use case, a legal contract and a support transcript do not need the same chunking strategy, so this step is left flexible rather than fixed to one method. 

2. Each chunk is embedded with a Databricks-hosted embedding model. 

3. The embeddings are written into Lakebase, in a Postgres table with the pgvector extension enabled, alongside the chunk text and the id of the chat session that produced it. 

The table itself is a plain Postgres table, with pgvector supplying the vector column type: 

CREATE EXTENSION IF NOT EXISTS vector;  
CREATE TABLE IF NOT EXISTS session_documents ( 
    id BIGSERIAL PRIMARY KEY, 
    session_id TEXT NOT NULL, 
    filename TEXT NOT NULL, 
    chunk_text TEXT NOT NULL, 
    embedding VECTOR(1024) NOT NULL, 
    created_at TIMESTAMPTZ DEFAULT now() 
); 
CREATE INDEX IF NOT EXISTS session_documents_session_id_idx  
ON session_documents (session_id); 

And the ingestion step itself is a short function: split, embed, insert. 

splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100) 
embeddings = DatabricksEmbeddings(endpoint="databricks-gte-large-en") 

async def ingest_document(session_id: str, filename: str, text: str, conn): 
    chunks = splitter.split_text(text) 
    vectors = await embeddings.aembed_documents(chunks) 

    async with conn.cursor() as cur: 
        for chunk, vector in zip(chunks, vectors): 
            await cur.execute( 
                """ 
               INSERT INTO session_documents (session_id, filename, chunk_text, embedding) 
                VALUES (%s, %s, %s, %s) 
                """, 
                (session_id, filename, chunk, vector), 
            ) 

The session id column is the entire isolation mechanism. Nothing else distinguishes one chat’s documents from another’s, and no per-session setup step is required before a document becomes searchable. 

Retrieval 

Retrieval is a plain LangChain tool the agent can call on its own: 

@tool 
async def search_uploaded_documents(query: str, session_id: str, conn) -> str: 
    """Search documents uploaded earlier in this chat session.""" 
    [query_vector] = await embeddings.aembed_documents([query]) 

    async with conn.cursor() as cur: 
        await cur.execute( 
            """ 
            SELECT chunk_text 
            FROM session_documents 
            WHERE session_id = %s 
            ORDER BY embedding <=> %s 
            LIMIT 5 
            """, 
            (session_id, query_vector), 
        ) 
        rows = await cur.fetchall() 

    if not rows: 
        return "No matching content found for this session." 
    return "\n\n".join(row["chunk_text"] for row in rows) 

The <=> operator is pgvector’s cosine distance, so ordering by it ascending returns the closest matches first. The WHERE session_id = %s clause is doing all the isolation work described above, in plain SQL. 

The model decides on its own when to call this tool: if the question looks like something an uploaded document might answer, the tool runs; if nothing was uploaded in that session, it returns nothing, and the agent answers from the conversation alone. Because pgvector performs the search directly on this Postgres table, a newly inserted chunk is queryable as soon as it is committed, with no separate indexing or sync step between ingestion and retrieval. 

Deletion 

A session-scoped table does not stay small on its own.  Without a purge strategy in place, the session_documents table grows without bound, and at some point a large, ever-growing table on autoscaling Postgres stops being the cheap, fast option this design sets out to be in the first place. Deciding how and when to remove old data is as much a part of the pipeline as ingestion and retrieval. There are two considerations to keep in mind: one about when a chat stops being queryable, and one about what happens to its data once it does.  

The first is retention keyed on last activity rather than upload date. A fixed calendar cutoff purges a chat seven days (or any period of time required) after its first upload regardless of whether anyone has touched it since. Retention keyed on last access instead resets the clock every time the chat is opened, so a chat someone keeps returning to stays queryable indefinitely in practice, while one abandoned after a single question ages out on schedule. 

The second is moving expired chunks into Unity Catalog rather than deleting them. Lakebase is the right home for data that needs to be queried live, but not for an accumulating permanent archive . A Delta table in Unity Catalog, backed by plain object storage, is a much cheaper place for that same data to sit once nobody is actively querying it, and it comfortably scales to a size that would make a Lakebase table increasingly expensive to keep around. 

CREATE TABLE IF NOT EXISTS main.rag_archive.session_documents ( 
    session_id STRING, 
    filename STRING, 
    chunk_text STRING, 
    embedding ARRAY<FLOAT>, 
    archived_at TIMESTAMP 
) USING DELTA; 

Bringing a session back is a small, on-demand lookup rather than a re-run of the ingestion pipeline: the chunks are already parsed and embedded, so restoring them is a copy from Delta back into Lakebase, with the embedding cast back into pgvector’s vector type on the way in. 

def rehydrate_session(session_id: str, pg_conn, uc_connection): 
    rows = uc_connection.execute( 
        """ 
        SELECT filename, chunk_text, embedding 
        FROM main.rag_archive.session_documents 
        WHERE session_id = ? 
        """, 
        (session_id,), 
    ).fetchall() 

    with pg_conn.cursor() as cur: 
        for filename, chunk_text, embedding in rows: 
            cur.execute( 
                """ 
                INSERT INTO session_documents (session_id, filename, chunk_text, embedding) 
                VALUES (%s, %s, %s, %s::vector) 
                """, 
                (session_id, filename, chunk_text, embedding), 
            ) 

Lakebase vs. Databricks Vector Search 

This is the central design decision behind the project, so it is worth addressing directly: why not use Vector Search for this as well? 

Vector Search is a strong product for a large, durable, shared corpus queried by many users over a long period. It provides managed approximate-nearest-neighbor indexing and scales well beyond what a Postgres table with pgvector can reasonably handle. The relevant question here is narrower: does that shape match a document that exists for the length of one conversation? 

Two properties of Vector Search make it a poor match for this case. 

• A Vector Search endpoint runs as always-on serving infrastructure and is billed continuously, whether a conversation is active or not. For a document whose useful life is measured in minutes, that is a poor trade. 

• The natural unit in Vector Search is the index, and creating one per chat session does not fit that model. Building an index and syncing data into it takes real time, typically minutes rather than seconds — the wrong latency budget for a document that needs to be searchable right after upload. 

Lakebase fits this shape of problem for the opposite reasons. It runs as autoscaling Postgres, so it can suspend to zero when idle and resume on the next connection, which means cost tracks actual usage rather than worst-case standing capacity. Partitioning by session requires nothing more than a session id column: no per-session provisioning step exists, and a session’s data is one WHERE clause away from every other session’s. With pgvector running directly on that table, a row becomes queryable the moment it is committed. 

None of this makes Lakebase superior to Vector Search as a general product. It fits this particular job better. A pgvector table on Postgres cannot match purpose-built ANN indexing at very large scale. Vector Search remains the right choice for a corpus that is large and permanent. Lakebase suits a corpus that is small, disposable, and needs to feel instant.