Skip to main content

Run vector search

This recipe performs a vector search query on a Mosaic AI Vector Search index using the Databricks SDK for Python.

Code snippet​

app.py
import reflex as rx
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
openai_client = w.serving_endpoints.get_open_ai_client()
EMBEDDING_MODEL_ENDPOINT_NAME = "databricks-gte-large-en"

def get_embeddings(text: str):
    try:
        response = openai_client.embeddings.create(
            model=EMBEDDING_MODEL_ENDPOINT_NAME, input=text
        )
        return response.data[0].embedding
    except Exception as e:
        return f"Error generating embeddings: {e}"

def run_vector_search(index_name: str, columns: str, prompt: str):
    columns_list = [col.strip() for col in columns.split(",") if col.strip()]
    prompt_vector = get_embeddings(prompt)

    if prompt_vector is None or isinstance(prompt_vector, str):
        return str(prompt_vector)

    try:
        results = w.vector_search_indexes.query_index(
            index_name=index_name,
            columns=columns_list,
            query_vector=prompt_vector,
            num_results=3
        )
        return results.result.data_array
    except Exception as e:
        return f"Error running vector search: {e}"

class VectorSearchState(rx.State):
    index_name: str = ""
    columns: str = ""
    search_query: str = ""
    search_results: str = ""
    is_searching: bool = False

    async def perform_search(self):
        self.is_searching = True
        yield
        result = run_vector_search(self.index_name, self.columns, self.search_query)
        self.search_results = str(result)
        self.is_searching = False

Resources​

Permissions​

Your app service principal needs the following permissions:

  • USE CATALOG on the catalog that contains the Vector Search index
  • USE SCHEMA on the schema that contains the Vector Search index
  • SELECT on the Vector Search index

See Query a vector search endpoint for more information.

Dependencies​

requirements.txt
databricks-sdk
reflex