Next Metricfy

Building a RAG Server with a Vector Database in Databricks-Part 1

—

by

richardchi3565
in

Retrieval-Augmented Generation (RAG) has quickly become one of the most practical platforms for producing an AI chatbot. Instead of relying solely on a pre-defined knowledge base, RAG server responses in your own data — making it accurate, up-to-date, and explainable.

Databricks is an excellent built in vector index platform for a RAG server. It combines the Mosaic AI Vector Search managed vector database, MLflow for model tracking and registering, Unity Catalog for governance, and seamless access across workspaces . In this post, we’ll walk through how to build a fully functional RAG server on Databricks from scratch .

Architecture Overview

Here’s the high-level architecture we’ll build:

Note: Prerequisite items are 1)Serverless compute 2)Unity Catalog in order to create out-of-box vector database in Databricks server.

▼ Data loading -Part 1

▼ Embedding column -Part 1

▼ Vector Index -Part 1

▼ Model Register-Part 2

▼ Final Response-Part 2

▼ Service Endpoint-Part 3

Data Loading

In  Databricks menu in left panel, go to the Data Engineering/Data Ingestion then select Create or modify table to upload a csv file. Also you can select a AWS S3 bucket or Azure Blob Storage connector

Sample data is located at https://www.kaggle.com/datasets/kyanyoga/sample-sales-data 

In python

# Create the target Delta table first (optional, COPY INTO can create it)
spark.sql("CREATE TABLE IF NOT EXISTS my_catalog.my_schema.my_table (col1 STRING, col2 INT)")
# Execute the COPY INTO command
spark.sql(f"""
  COPY INTO my_catalog.my_schema.my_table
  FROM '/Volumes/<catalog-name>/<schema-name>/<volume-name>/<file-name>.csv'
  FILEFORMAT = CSV
  OPTIONS (
    'header' = 'true',
    'inferSchema' = 'true'
  )
""")

Embedding Column

Once data is loaded into Datalake as a raw table, you can create a cleaned fact table with an embedding text column for vector index.

In Python

from pyspark.ml.feature import VectorAssembler
from pyspark.sql.types import ArrayType, FloatType
from pyspark.sql.functions import col, concat_ws, lit, concat
# Assuming your source Delta table is named "main.default.numerical_data"
table_name = "workspace.default.sales_data_raw"
df = spark.read.table(table_name)
 
df_combined = df.withColumn("combined_text", concat_ws(" ", col("customername"), lit("makes"), col("sales").cast("string"), lit("in"), col("year_id").cast("string"),lit("at"), col("city"), col("state"), lit("in"), col("country"), col("productline"), lit("is"), col("productcode"), lit("in"),col("territory"),  lit("deal size was"), col("dealsize")))

Vector Index

In  Databricks menu in left panel, go to the Playground then select an Endpoint such as Gemma 3 and Tool. Click on Add tools and select a Vector search tab and the select predefined your own vector endpoint. Start tying in command prompt.

In Python

from databricks.vector_search.client import VectorSearchClient
vsc = VectorSearchClient()
endpoint_name = "your_vector_search_endpoint_name"
index_name = "your_catalog.your_schema.your_direct_access_index_name"
index = vsc.create_direct_access_index(
    endpoint_name=endpoint_name,
    index_name=index_name,
    primary_key="id",
    embedding_dimension=1024, # Dimension of your precomputed vectors
    embedding_vector_column="combinded_text",
    # schema defines additional columns
    schema={"id": "string", "combinded_text": "array<float>", "text": "string"}

Vector Index Testing

In  Databricks menu in left panel, go to the Compute then select Vector Search tab. Create a vector search endpoint by selecting embedding column called “combinded_text”

In Python

import sys, os
# Append the directory containing your module to sys.path
sys.path.append(os.path.abspath('/databricks/python/lib/python3.12'))
sys.path.append(os.path.abspath('/databricks/python/lib/python3.12/site-packages'))
#print(sys.path) 
#%pip install databricks-langchain
#%restart_python
from databricks_langchain import ChatDatabricks
from databricks.vector_search.client import VectorSearchClient
from langchain_classic.chains import RetrievalQA
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.vectorstores import DatabricksVectorSearch
import mlflow
# --- Configuration ---
WORKSPACE_URL="[Your Workspace URL]"
VECTOR_ENDPOINT="[Your Endpoint]"
VECTOR_INDEX_NAME = "[Your Index Name]"
LLM_ENDPOINT = "databricks-meta-llama-3-70b-instruct" # Example foundation model
# Define the prompt template
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant. Use the following context to answer the user's question. If you don't know the answer, just say that you don't know, don't try to make up an answer.\nContext: {context}"),
    ("human", "{question}")
])
# Function to create the LangChain agent with Vector Search retriever
def create_agent():
    # Initialize the Vector Search client and retriever
    vsc =  VectorSearchClient(workspace_url=WORKSPACE_URL)
    # Note: the .as_retriever() method is available in the databricks_langchain package
    index  = vsc.get_index(
    endpoint_name=VECTOR_ENDPOINT,
    index_name=VECTOR_INDEX_NAME
)
    vectorstore = DatabricksVectorSearch(index)
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
 
    # Initialize the LLM
    llm =ChatDatabricks(
    endpoint="databricks-meta-llama-3-1-405b-instruct", # Example endpoint name
    temperature=0.1,
    max_tokens=256)
    # Create the RAG chain
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
        chain_type_kwargs={"prompt": prompt}
    )
 
    return qa_chain
 
print(create_agent()['result'])

Comments

Leave a Reply