Build it: RAG in SQL on Databricks with vector_search and ai_query, and a retrieval check

Part 3 called a model from SQL on a table. This part joins the module's pieces into the assistant the shop asked for: a customer's question in, an answer from the help articles out, with the articles it came from. It needs the Databricks side of part 2: the table workspace.shop.help_chunks and its index workspace.shop.help_chunks_index, synced. To rebuild from nothing, run part 1's 2 scripts in the workspace, with the help files in a volume and the table named workspace.shop.help_chunks, then part 2's notebook cells. The queries run in a SQL editor on a serverless SQL warehouse.

Build it: find the chunks

vector_search searches an index from SQL. Every argument is passed by name:

SELECT chunk_id, heading
FROM vector_search(
  index => 'workspace.shop.help_chunks_index',
  query_text => 'When do I get my money back?',
  num_results => 3,
  query_type => 'HYBRID'
)

query_text is embedded with the index's own model, so the question and the chunks are compared in the same space. query_type => 'HYBRID' combines that similarity search with a keyword search, which catches exact words a meaning search can miss, such as an order number or a product code; Databricks recommends it as the starting point for most uses. The function returns all the columns the index synced, by default every column of the table, so article, title, heading and text come back with each chunk_id.

For the returns question, the chunk to hope for is returns-2, "How long a refund takes". Part 1 put "Returns and refunds" into its text for exactly this kind of question, which never says "refund".

Answer from them

The answer is 1 more step: hand the chunks' text to a model with the question, and tell it to use nothing else. The test questions go in a table, so the same query serves 1 question or a thousand:

CREATE OR REPLACE TABLE workspace.shop.help_questions (question_id INT, question STRING, expected_article STRING);

INSERT INTO workspace.shop.help_questions VALUES
  (1, 'When do I get my money back?', 'returns'),
  (2, 'How much is delivery on an order of 25?', 'delivery'),
  (3, 'My parcel never came', 'delivery'),
  (4, 'Can I use 2 ways to pay for 1 order?', 'payments'),
  (5, 'Can I send back opened food?', 'returns'),
  (6, 'Do gift cards expire?', NULL);

Question 6 has no expected article on purpose: part 2 removed the section that answered it. Then the answer query:

WITH retrieved AS (
  SELECT q.question_id, q.question,
         array_join(collect_list(s.text), '\n\n') AS context,
         collect_list(s.chunk_id) AS sources
  FROM workspace.shop.help_questions AS q,
  LATERAL (
    SELECT * FROM vector_search(
      index => 'workspace.shop.help_chunks_index',
      query_text => q.question,
      num_results => 3,
      query_type => 'HYBRID')
  ) AS s
  GROUP BY q.question_id, q.question
)
SELECT question_id, question,
       ai_query(
         'system.ai.meta-llama-3-3-70b-instruct',
         'You answer customer questions for an online shop, using only the help articles below. '
           || 'If they do not contain the answer, say so and suggest contacting support. '
           || 'Answer in at most 3 sentences.\n\nHelp articles:\n' || context
           || '\n\nQuestion: ' || question,
         failOnError => false
       ).response AS answer,
       sources
FROM retrieved
ORDER BY question_id

• LATERAL runs the search once for each question, with that row's question as query_text, as Databricks' own example of many searches does.
• context joins the 3 chunks' text into 1 block; each chunk starts with its title and heading, from part 1, so the model sees what each passage is about. collect_list does not promise an order, which is fine here, because the prompt treats the 3 passages alike.
• The prompt limits the model to the passages and says what to do when they do not hold the answer. Without that line, a model asked about gift card expiry may answer from what shops usually do, or from an old version of the article it never saw.
• sources are the chunk ids the search returned, from the index, not from the model, so they cannot be made up. Show them under the answer, with the headings, so a customer or an agent can read the article itself.

Read the answers to questions 2 and 6. Question 2 needs the rule that standard delivery is free over 40.00 and costs 4.50 below it, after part 2's change; question 6 should say the articles do not answer it.

Build it: measure the retrieval

The model can only answer from what the search hands it. So check the search on its own, before judging any answer: for each question with an expected article, is that article in the top 3?

WITH hits AS (
  SELECT q.question_id, q.question, q.expected_article,
         array_contains(collect_list(s.article), q.expected_article) AS found,
         collect_list(s.chunk_id) AS returned
  FROM workspace.shop.help_questions AS q,
  LATERAL (
    SELECT * FROM vector_search(
      index => 'workspace.shop.help_chunks_index',
      query_text => q.question,
      num_results => 3,
      query_type => 'HYBRID')
  ) AS s
  WHERE q.expected_article IS NOT NULL
  GROUP BY q.question_id, q.question, q.expected_article
)
SELECT question_id, question, expected_article, found, returned,
       round(avg(CAST(found AS INT)) OVER (), 2) AS share_found
FROM hits
ORDER BY question_id

share_found is the share of questions whose article came back, the number to watch; the rows with found false are the ones to read. The Retrieval EvaluationRetrieval EvaluationTesting the search step of a RAG system on its own, with questions whose answering passages you know, and scoring it with recall at k, the share of those passages found in the top k results. concept covers the measures beyond this one, such as where in the list the right chunk ranks.

5 questions are a start, not a measure. A real set holds the questions customers actually ask, from the contact form of part 3, a few dozen or more, each with the article that answers it. Run this query after every change to the chunking, the embedding model or the search type, and keep the change only if share_found does not fall. A miss is usually fixed in the chunks: a heading without the words customers use, or 2 answers in 1 section.

Keep it current

The index is only as current as its last sync. In the bundle from module 11, a second job reloads the chunks and then syncs the index, every night:

    help_index:
      name: help_index
      schedule:
        quartz_cron_expression: "0 30 2 * * ?"
        timezone_id: UTC
      max_concurrent_runs: 1
      tasks:
        - task_key: load_help
          spark_python_task:
            python_file: src/chunk_help.py
          environment_key: default
        - task_key: sync_index
          depends_on: [{ task_key: load_help }]
          spark_python_task:
            python_file: src/sync_index.py
          environment_key: default
      environments:
        - environment_key: default
          spec:
            environment_version: "2"
            dependencies:
              - databricks-ai-search

It sits under resources.jobs beside shop_daily, and the prod target needs a matching entry for help_index: its schedule again, with pause_status: UNPAUSED, as module 11 does for shop_daily. src/chunk_help.py is part 2's script with module 10's changes for Databricks: the help files read from a volume, the table named workspace.shop.help_chunks, and SparkSession.builder.getOrCreate() for the session. src/sync_index.py is part 2's 1 line:

from databricks.ai_search.client import AISearchClient

AISearchClient().get_index(index_name="workspace.shop.help_chunks_index").sync()

The MERGE changes only the chunks whose text changed, so on most nights the sync has nothing to embed.

The endpoint is billed from the moment an index exists on it, whether or not anyone searches. On a paid workspace, delete the index when you finish the module; the endpoint stops charging 24 hours after its last index is deleted.

What you built

• A table of help chunks, 1 per section, each with a stable id and the title and heading in its text.
• A load that changes only the chunks that changed, and an AI Search index that re-embeds only those.
• Customer messages routed and summarised in SQL, each processed and paid for once.
• An answer to a customer's question from the help articles, with its sources, and a retrieval check to measure every change against.

The AI engineering course builds the same assistant from the application side, in Python, with its prompts and evaluations. Module 14 puts the whole course together: the shop's pipeline from raw files to a dashboard and an AI answer, end to end.

CourseAI engineering course, Module 5: RAG that answers correctlyRetrieval-augmented generation (RAG) answers a question in 2 steps: a search finds the passages in your documents most likely to hold the answer, and a model…