Author(s): Ishwarya Modika

Originally published on Towards AI.

Inside Databricks Knowledge Assistant: The Architecture Behind Smarter, Faster Enterprise Search

Every retrieval system can find documents. Very few can follow instructions while doing it. Here’s the story of how one product learned to do both — and then learned to do it fast.

Tell a chat bot “find me the revenue, but not from that product line,” and something odd happens: it hears “revenue” and “product line,” and just… forgets the word “not.” That’s not a one-off bug. It’s a structural blind spot in how retrieval has traditionally worked, and it’s exactly the problem Databricks set out to fix in Knowledge Assistant.

Here’s the interesting part — they didn’t fix it once. They fixed it twice, and each fix tackled a completely different failure mode:

• The first redesign taught retrieval to actually follow instructions, instead of just matching keywords.
• The second redesign made that instruction-following fast, by swapping slow sequential retries for parallel search.

The architecture moved from something called Instructed Retriever to its successor, Instructed-Retriever-1 — the design running Knowledge Assistant today. To see why each redesign happened, it helps to put all three through the same real question and watch what each one does with it.

Understanding Things With an Example

User query: “What is XYZ Pharma’s total revenue in Q3 2026, excluding its weight-loss drug?”

System instructions: prefer official filings over analyst chatter, keep the answer under 150 words, and don’t touch anything older than 12 months.

Picture this happening inside ABC BFSI, a bank/advisory firm whose internal knowledge base holds documents for hundreds of client companies — filings, credit memos, earnings summaries, analyst notes, call transcripts, all sitting in shared vector indexes. An analyst fires off the question above and needs a number they can actually put in front of a client.

It reads like one question. It’s really four separate demands, stacked on top of each other:

• Identify the right company — one of hundreds sharing the index
• Identify the right quarter — Q3 2026, and nothing else
• Perform a hard exclusion — exclude the docs tied to the weight-loss product
• Obey formatting rules — sourcing, recency, and a strict word limit

Normal RAG fails at this. And that failure is exactly the gap the retrieval system inside Knowledge Assistant was built to close.

Traditional RAG

The Question That Exposes the Gap

“What is XYZ Pharma’s total revenue in Q3 2026, excluding its weight-loss drug?”

In plain language: XYZ Pharma sells a bunch of products. One of them — a weight-loss drug — sells extremely well. The analyst wants the revenue excluding the weight-loss drug’s revenue, so they can see how the rest of the business is doing on its own.

That word excluding is the whole problem. Traditional RAG matches on meaning and topic similarity, not on logical constraints. A vector or embedding search only ever asks one question:

“How similar in meaning is this chunk of text to the user query?”

It sees “weight-loss drug” and thinks topically relevant — but it will not exclude the documents having it. To the embedding model, both phrases are just more text to match against the vector DB.

Since Normal RAG runs one embedding search on the whole query text, what it’s effectively looking for is something like “revenue Q3 2026.” It has no mechanism to act on the exclusion at all, and no way to enforce something like “only the most recent 10-Q” as a hard rule either. Recency and exclusion are just words floating in the query, indistinguishable from any other word.

What Actually Comes Back

[101] Q1 2026 10-Q — total revenue $4.6B
[105] Q3 2026 10-Q — total revenue $4.82B
[109] Analyst note: "Slimzol likely drove ~$600M this quarter"

It doesn’t exclude the old document, and it doesn’t filter out the irrelevant one either — it’s the LLM that has to catch [109] at response generation time and subtract the revenue. Retrieval just hands over a mixed pile, quietly passing all the retrieved docs downstream to the LLM.

Where it fails: plain RAG retrieves everything related. It has no way to also retrieve only what’s allowed.

Generation One: Instructed Retriever

This is where the Knowledge Assistant story really begins, and honestly, the fix is almost simple in hindsight: stop hiding the instructions from retrieval.

Instead of the rules only showing up at the response-generation step, the system reads the query and the instructions together, up front, and turns both into one structured search plan.

There’s a third input worth calling out on its own: the index schema — the actual metadata field names attached to chunks in the vector database, things like company_name, fiscal_period, or product_line. It feeds directly into the query generation model, since that's the stage doing the real work of turning instructions into filters, and it needs to know which field names it's even allowed to write filters against. Unlike instructions and examples, which travel with the query all the way through reranking and response generation, schema is only consulted once, right where filters get built.

To be precise about where everything lands:

• System specs (instructions + examples) → query generation model, reranking model, and the response generator — the same rules ride along with the query all the way to the final answer
• Index schema → query generation model only — so filters reference real, existing metadata fields instead of ones the model is just guessing at.

Step 1 — Query Generation

Search queries: “total revenue,” “revenue by product line”

Filters:

{
 "company_name": "XYZ Pharma",
 "fiscal_period": "Q3 2026",
 "doc_type": ["10-Q", "earnings_release", "press_release"],
 "doc_timestamp": ">= 2025-08-15",
 "source_priority": "official_filing",
 "product_line": "!= weight loss drug"
}

Every field name in that filter — company_name, fiscal_period, product_line — only works because the index schema, fed to the query generation model up front, told it those fields exist and could actually be filtered on. Without that schema input, the model could propose a filter on a field the index doesn't support, and the whole exclusion would silently fail once it hit retrieval.

That last line — product_line != weight loss drug — is quietly doing three jobs at once:

• It prevents a document that’s tagged as the weight-loss product line from ever entering the segment-revenue retrieval pool for the “total” side of the calculation. If segment-level revenue documents are chunked/indexed per product line, this filter cleanly separates “everything except the drug” from “the drug itself” — rather than retrieving all segment documents and hoping the generator later remembers to subtract the right one.
• It narrows the reranker’s job. Instead of reranking across all segments and relying on the model to figure out at generation time which one to drop, the reranker only ever sees documents that are already correctly scoped — so relevance scoring is spent on things like recency and source-priority, not on redoing the exclusion logic.
• It protects the final number. The generator receives pre-filtered inputs where the subtraction is either already done or trivially computable from clean, correctly-scoped segment data — reducing the chance it hallucinates the wrong number or forgets the exclusion instruction entirely once several other constraints (word limit, source priority) are competing for its attention.

Step 2 — Retrieval

Using the structured query and filters from Step 1, the system runs a single retrieval pass against the vector index. This returns the top-k chunks that pass the filter and rank highest for similarity to the query.

Step 3 — Reranking

Unlike traditional RAG reranking, which scores documents purely against the user query, this reranker is context-aware. Guided by system specs — instructions and examples — it evaluates documents for contextual relevance rather than surface-level query similarity, then returns the top-n results as the final, narrowed set.

User query + system specs (instructions + examples) + chunks → Reranker model

Step 4 — Response Generation

The top-n results, user query, and system instructions are fed together into the LLM, which generates the response — the culmination of every context-aware step before it.

Top-n results + user query + system instructions → LLM → Response

The measured impact: On Databricks’ StaRK-Instruct benchmark, Instructed Retriever delivered a 35–50% improvement in retrieval recall over a raw-query baseline. Against a basic RAG setup, it performed up to 70% better. Even compared to more advanced DIY solutions that add a reranking stage on top of RAG, it still came out 15% ahead. And when plugged into a multi-step search agent as a tool, it boosted performance by over 30% compared to using plain RAG as the tool — while cutting task completion time by 8%.

Where it wins: it removes the mess before the LLM ever sees it, instead of asking the LLM to clean up after.

The Gap That Led to Generation Two

Instructed Retriever wasn’t broken. It just had a blind spot.

Here’s the issue: for every question, it writes exactly one search plan. Just one. And real-world documents don’t all describe the same thing the same way. Take our example — the system might search for “revenue by product line,” but the actual filing could label that same data “product revenue by therapeutic area” instead. Same information, different words. Since the system generates search queries with only one phrasing, it might fail to fetch — or might score poorly — chunks that express the same information in different wording.

This single-search design shows up as two distinct retrieval problems, both tied to fetching a fixed top-k:

• Low recall — because only one phrasing gets searched, some genuinely relevant documents never make it into the top-k at all. If a document happens to use different wording than the query, it simply doesn’t surface, no matter how relevant it actually is.
• Low precision — the flip side of the same problem. Since the single query might not match well, some of the documents that do make it into the top-k may not be all that relevant either — the system fills its quota of k documents, but not all of them are genuinely useful.

So it’s not just that the system might miss the right document — it’s that a fixed-size, single-shot search can simultaneously miss good documents (hurting recall) and let weaker ones take their place (hurting precision).

The natural fix sounds simple: if the LLM decides the retrieved information isn’t enough, it loops back and searches again — and again — until it finds what it needs. But that fix has a hidden cost: each loop is sequential, meaning the system has to wait for one attempt to finish before trying the next. Every extra loop adds more delay, and that’s a real problem for a product that’s supposed to feel instant, not like it’s thinking it over.

So this was Instructed Retriever’s real weakness: not wrong answers, but a system that could quietly miss good documents and pull in weak ones instead — just because of how the query was worded. The only way to fix it was to repeat the search, and repeating it meant waiting longer, since each attempt had to finish before the next one could start.

This is exactly the gap the next generation was designed to close. And the fix wasn’t “loop again to get results, but smarter.” It was “don’t rely on a single retrieval attempt in the first place.”

Generation Two: Instructed-Retriever-1

This is the architecture running Knowledge Assistant today. The core idea hasn’t changed — instructions still shape retrieval, filters still keep out the wrong documents, reranking still weighs relevance against the rules. What’s different is how it resolves Instructed Retriever’s tradeoffs: instead of looping to fix a missed search, it adds parallelism, running multiple searches and comparisons at once instead of one after another.

Two things make this different from a more typical pipeline. First, query generation and reranking come from the same model — not two separate components passed back and forth, which turns out to matter more than it seems. Second, there are two distinct parallel stages, each doing its own job: one expands the search, the other refines the ranking.

Step 1 — Ask the Question Several Ways at Once

Rather than settling on a single phrasing, the model writes multiple search plans in the same pass.

When product_line exists in the index schema: Query 1 and Query 2 work as intended. The exclusion filter — product_line != weight loss drug — keeps the drug's chunk out of the candidate pool during retrieval. If it somehow passes through, reranking's instruction-aware scoring removes it. Either way, the drug's chunk gets excluded before it reaches the final answer.

When product_line is not part of the index schema, two things go wrong with Query 1 and Query 2:

• There’s no field to write the exclusion filter against.
• Nothing structurally stops the drug’s chunk from appearing in their results, since both queries now rely on topical similarity alone.

Query 3 is built for this exact gap:

• It doesn’t depend on a filter that can’t be constructed.
• It searches for the drug’s revenue directly — “weight-loss drug sales Q3 2026” — using intent instead of a schema-based filter.
• It brings back a clear, well-defined chunk for that number, rather than getting it by accident somewhere else.
• It makes sure the number is actually found, instead of hoping a missing filter doesn’t cause problems.

This is also why the subtraction ends up happening at the response stage instead of earlier. When schema exists, the drug’s revenue gets excluded before the LLM ever sees the results — it just gets the already-correct total. When schema doesn’t exist, the system can’t do that exclusion earlier, so it hands the LLM both numbers — the total, and the drug’s own figure that Query 3 found — and lets the LLM do the subtraction itself. So exclusion happens as a filter when schema exists, and as a calculation at generation time when it doesn’t — and Query 3 is what makes that second path possible at all.

This is the system’s first lever: asking the question in more than one way improves recall, and it doesn’t slow anything down, since all the searches run at the same time instead of one after another.

Step 2 — Retrieve Without Standing in Line

All three searches hit the vector index at the same moment:

Query 1 → top-k: chunk 1 (filing revenue table), chunk 3 (analyst note, total revenue)
Query 2 → top-k: chunk 3, chunk 6 (segment breakdown)
Query 3 → top-k: chunk 7 (drug-specific revenue), chunk 9 (call transcript excerpt)

Chunk 3 turns up in both Query 1 and Query 2 — that’s not wasted effort, it’s confirmation that chunk 3 genuinely matters across multiple angles of the “total revenue” side of the question. Chunk 7, the drug-specific figure, is retrieved cleanly by Query 3 precisely because it was built to find it — not because it slipped past a filter meant to keep it out. Metadata filters and vector similarity aren’t applied as two separate steps here, either: wherever a filter exists, it shapes the search while it happens, so a chunk from an unrelated company never even makes it into the candidate pool.

At this point, the system already has more relevant chunks in hand than a single search ever could — chunk 1 and chunk 3 for the total, chunk 6 for the segment breakdown, chunk 7 for the drug-specific figure — all surfaced at once instead of depending on one query to catch everything.

That’s recall in action. Recall measures how much of the truly relevant information actually gets found, not just how accurate the results look. Three differently-worded queries running at the same time cover more of that ground than one query ever could, since each phrasing might catch documents the others miss — a filing might say “segment revenue,” another “product line performance,” and no single query can hedge against that variation alone. And because all three run simultaneously rather than one after another, this broader coverage doesn’t cost any extra time — no retry, no waiting for one search to finish before trying the next.

Step 3 — Merge Into One Candidate Pool

Before any ranking begins, everything gets pooled together and deduplicated:

Merged pool = {1, 3, 6, 7, 9, ...}

This step simply wouldn’t exist in a system that only ever ran one search — it exists precisely because Step 1 introduced parallelism. Multiple simultaneous searches naturally produce overlapping results, and something needs to reconcile that before reranking can make sense of the full picture. Notice the pool now contains both the “total” side (chunks 1, 3, 6) and the drug-specific figure (chunk 7) side by side — deliberately, not accidentally.

Step 4 — Multi-Pivot Groupwise Reranking

This is where the system does its most understated but important work, and it’s worth slowing down for. Scoring every merged chunk on its own is fast but shallow — there’s no sense of relative quality. Comparing every chunk against every other chunk gives plenty of context, but the cost grows too quickly to stay practical: each chunk would need to be measured against all n-1 others, and that comparison count grows fast as the pool grows.

Instructed-Retriever-1 strikes a balance: it selects a handful of chunks to act as pivots — reference points — and builds a small group around each one. Pivot grouping is what actually brings the cost back down — each chunk only needs to be compared against its group’s single pivot, not against everyone else. Parallelism then adds a second, separate benefit on top of that: since the groups don’t depend on each other, all of them can be ranked at the same time, so having multiple groups doesn’t cost any extra time compared to having just one.

How Pivots Actually Get Chosen

Databricks hasn’t published exactly how pivots get chosen, so what follows is reasoned inference, not confirmed fact. That said, three mechanisms fit the rest of the architecture better than any alternative would.

1. Instruction-aware relevance scoring

Rather than ranking candidates purely by their raw similarity to the query, the model combines two signals for each one — how closely it matches the query, and how well it satisfies the instructions — and picks the highest scorers as pivots.

relevance_score = similarity_score × instruction_fit_score

The similarity part works the usual way — it measures how close the meaning of the chunk is to the meaning of the query, using standard embedding comparison.

The instruction-fit part is more interesting, and it’s worth being upfront that Databricks hasn’t published exactly how it works — what follows is a reasonable inference, not a confirmed detail. The model likely looks at the instructions alongside the chunk’s text and its metadata (things like publication date or document type) and checks whether the chunk actually follows the rules. Some checks are probably clean and mechanical — like confirming a date falls within the last 12 months. Others are fuzzier — like “prefer recent” with no exact cutoff — where the model likely relies on judgment rather than a strict yes/no. Either way, the result lands as a single score: close to 1 if the chunk follows the instructions well, close to 0 if it clearly breaks them.

Chunk 3 (analyst note, Q3 2026, on-topic): similarity 0.85 × instruction-fit 0.60 = 0.51
Chunk 5 (10-Q filing, Q3 2026, official): similarity 0.79 × instruction-fit 0.95 = 0.75 ← picked as pivot
Chunk 9 (weight-loss drug commentary): similarity 0.91 × instruction-fit 0.10 = 0.09 ← excluded despite high similarity

Even though chunk 9 reads as the most similar to the query out of the three, it still gets ruled out — because it breaks the exclusion rule. Chunk 5 wins instead, despite a lower similarity score, simply because it’s the official, recent document the instructions actually called for.

2. Cross-formulation agreement, filtered by instructions

A chunk that turns up across more than one parallel query formulation is a natural sign of relevance — but showing up repeatedly isn’t enough on its own. It still has to clear the instruction filter before it can qualify as a pivot.

[Choosing pivot for Group 1]

Chunk 3 appears in Query 1 ("total revenue") and Query 2 ("segment breakdown")
 → cross-formulation hit, AND satisfies "official + recent" → eligible pivot
Chunk 9 appears in Query 2 ("segment breakdown") and Query 3 ("weight-loss drug sales")
 → cross-formulation hit, BUT fails "excluding weight-loss drug" → disqualified from anchoring Group 1

Showing up in more than one search is a good sign for a chunk — but it doesn’t matter unless the chunk also follows the instructions.

3. Learned selection

Since one model handles both query generation and reranking, pivot selection is plausibly just another learned output of that same model — shaped by labeled examples showing which chunk should have anchored a group, given a query and its instructions.

Training example:
 Query: "Total revenue excluding [product], official sources only"
 Correct pivot (from labeled data): the official filing chunk, not the highest-similarity commentary chunk

The model learns: when instructions include “official sources only,” weight source_type heavily in pivot selection — even above raw topical similarity.

Over many such examples, the model internalizes how much each instruction type should influence pivot choice — not as a hardcoded rule, but as a learned pattern baked into the same weights doing reranking.

Across all three, the common thread holds: instruction compliance acts as a gate or a weighting factor before or during pivot selection — not something ignored until reranking happens. A pivot chosen without that check risks anchoring its whole group around a chunk that shouldn’t even be in the running, which would undermine every comparison built on top of it.

Caveat: none of these three methods are confirmed by Databricks’ public documentation — this is the most architecturally consistent reasoning given what is confirmed, that instructions and examples propagate through every other stage of the pipeline. It would be an odd design choice for pivot selection to be the one stage in an instruction-driven pipeline that ignores those instructions entirely — the rest of the system was built specifically to carry rules through every step.

Forming the Groups

Pivot 1: chunk 3 (analyst note, total revenue) → Group 1
Pivot 2: chunk 7 (filing, drug-specific revenue) → Group 2
Pivot 3: chunk 9 (call transcript excerpt) → Group 3

Every group gets ranked at the same time as the others, not one after another. Inside each group, a chunk’s score comes from two things combined: how closely it matches its pivot, and how well it follows the instructions. So a chunk that reads almost identically to the pivot can still lose its spot if it comes from an outdated document, while a chunk that matches the pivot less closely on the surface can still come out ahead if it’s the one that actually satisfies the rules.

Group 1 ranking: 3 > 5 > 1
Group 2 ranking: 7 > 6 > 2
Group 3 ranking: 9 > 4 > 8

Notice chunk 7 — the drug’s own revenue figure (which needs to be excluded) — gets its own group instead of getting mixed into Group 1. That’s important, whether or not schema exists to enforce it with a hard filter: Group 1 stays focused on the total revenue, Group 2 stays focused on the number that needs to be subtracted (in case there’s no hard filter), and the two never get mixed up together.

This is the system’s second lever: more pivots mean more groups, and more groups mean better precision — again, paid for in parallel compute rather than time spent waiting.

Step 5 — Merge Group Rankings Into a Final Order

Three groups produce three internally consistent rankings, but none of them have been measured directly against one another yet. Combining them is deliberately lightweight:

Final order: 3, 7, 9, 5, 6, 4, 1, 2, 8

The most sensible approach reuses the signal the model already computed for each pivot — sorting primarily by pivot relevance, then fine-tuning within that using each chunk’s own instruction-aware score — rather than re-running a full comparison, which would undo the entire point of splitting into groups in the first place.

Caveat: as with pivot selection, Databricks hasn’t published the exact merge mechanism — this is reasoned inference based on what would be consistent with the rest of the architecture, not a confirmed implementation detail.

One detail worth flagging: the LLM doesn’t see the whole merged pool — only the top-n slice of it. Chunks further down the list scored lower on both topical relevance and instruction compliance, so passing them along would just give the LLM more to sort through, with a real chance of surfacing a stray, less trustworthy figure.

What ends up in that slice depends on whether schema was available. If product_line isn't part of the schema, the slice needs both the total revenue figure (chunk 3) and the drug-specific segment figure (chunk 7), since the LLM has to perform the subtraction itself. If schema does support the filter, the exclusion may already be resolved by the time retrieval finishes — so the slice that matters could just be chunk 3, already correctly scoped, with nothing left to subtract.

Putting It Together

Multi-pivot groupwise reranking is what gets the most relevant chunks into the top of the final listing — comparing each chunk against a strong reference point, rather than scoring it alone, is what makes that ranking trustworthy. Parallelism is what makes this happen fast — three groups get judged at the same time instead of one after another, so thoroughness doesn’t cost extra latency. And the result of combining both is that the chunks which are genuinely most relevant to the query and the instructions are the ones that end up with the highest ranking, not just the ones that happened to be scored first or scored alone.

Step 6 — Answer Generation

The final top-n chunks — carrying the total-revenue figure and, where needed, the drug-specific segment figure — reach the LLM along with the original query and instructions. Because these numbers arrive already cleanly separated, the subtraction the analyst actually asked for becomes something the LLM can just perform, rather than something it has to guess at or catch on its own. Generation stays anchored to both the retrieved evidence and the original constraints, so the final answer reflects exactly what was asked — word limit included.

Why Two Parallel Ideas, Not One

It’s easy to think of this as one long stream of parallel work from start to finish, as if each search kept its own comparison lane all the way through. In reality, there are two separate parallel stages, one after the other, with a merge in between:

[Parallel query generation] → merge → [Parallel pivot-group reranking] → merge → LLM
 (boosts recall) (boosts precision)

Searching wider and comparing more carefully are two different problems, solved by two different rounds of parallel work. Neither one depends on a sequential retry loop to catch what it might have missed — which is exactly why both can be scaled up for a harder question without the wait growing any longer.

One Model, Two Jobs

Query generation and reranking coming from a single trained model, not two stitched-together components, means the same instruction-awareness shaping how a search gets phrased also scores pivots during reranking — no risk of the two drifting apart. It also means only one model needs optimizing.

The Measured Impact

Search runs more than 3x faster. Full answer generation runs 2x faster. Time-to-first-token lands around 2 seconds, and end-to-end latency stays consistently under 10 seconds — all while matching frontier-model retrieval quality. None of that speed comes from cutting corners; it comes from trading a slow, sequential “search, notice something’s missing, search again” habit for several searches and comparisons happening side by side.

Where it wins: the same correctness, now delivered without the wait.

The architecture evolution, side by side

With all three side by side, the pattern is clear: each generation fixed exactly one thing the previous one couldn’t.

Two Wins, In Order

Why does Databricks Knowledge Assistant work as well as it does? Not one clever trick — two separate wins, in exactly the right order.

First, it stopped ignoring instructions. Rules became filters enforced before any document was retrieved, instead of an afterthought at the end — and even where schema can’t support a hard filter, the same instruction-awareness still finds its way in, through a dedicated search and a subtraction handled downstream instead.

Second, it stopped waiting. A slow “search, check, search again” loop became several searches and comparisons running at once, judged by pivots shaped by those same instructions.

For a question stacking four demands — the right company, the right quarter, a hard exclusion, strict formatting — that’s the whole story: a system that first learned to get the answer right, then learned to get it right fast.

A note on sourcing: some architectural details here — especially around Instructed-Retriever-1 — are reasoned inference, not confirmed by Databricks. For more, see the original Databricks blog.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI