Databricks AI Search Python SDK

The Python SDK for Databricks AI Search, the fully managed vector and keyword search service in the Databricks platform. AI Search is the new name for the product previously known as Databricks Vector Search.

Installation

pip install databricks-ai-search

The legacy databricks-vectorsearch package continues to work: it is a thin re-export shim of databricks-ai-search that emits a deprecation warning on import. New projects should install databricks-ai-search.

Requires Python 3.10 or later.

Usage

from databricks.ai_search.client import AISearchClient

# Initialize the client (auto-detects credentials in Databricks notebooks)
client = AISearchClient()

# Create an endpoint and index
endpoint = client.create_endpoint(name="my-endpoint", endpoint_type="STANDARD")
index = client.create_delta_sync_index(
    endpoint_name="my-endpoint",
    index_name="my_catalog.my_schema.my_index",
    primary_key="id",
    embedding_dimension=768,
    embedding_vector_column="embedding",
    source_table_name="my_catalog.my_schema.my_source_table",
    pipeline_type="TRIGGERED",
)

# Query the index
results = index.similarity_search(
    query_vector=[0.1, 0.2, 0.3, ...],
    columns=["id", "text"],
    num_results=10,
)

AISearchClient is the canonical class name; VectorSearchClient is preserved as a backward-compat alias (VectorSearchClient is AISearchClient evaluates True). The same applies to AISearchIndex / VectorSearchIndex and AISearchException / VectorSearchException.

Authentication

The client supports the same authentication methods as the legacy SDK:

• Auto-detection (recommended for Databricks notebooks) — no arguments required
• Personal Access Token (PAT) — AISearchClient(workspace_url=..., personal_access_token=...)
• Service Principal (OAuth M2M, AWS / GCP) — AISearchClient(workspace_url=..., service_principal_client_id=..., service_principal_client_secret=...)
• Azure Service Principal — additionally pass azure_tenant_id=...
• Model Serving with user credentials — AISearchClient(credential_strategy=CredentialStrategy.MODEL_SERVING_USER_CREDENTIALS)

Service principal tokens request all-apis by default; if the secret is not allowed it, the SDK tries vector-search model-serving ai-gateway, then vector-search model-serving, then vector-search ai-gateway, then vector-search, and remembers the first one granted for the rest of the process. Refused token requests therefore appear in audit logs on the first token request for each service principal secret and workspace in a process, at most one per scope tried. Threads that make that first request at the same time each try the scopes in turn, so each can log its own refusals; every one of them still gets a token for the first granted scope. A secret allowed none of these scopes is refused for each of them on every token request. Broader scopes are not retried while the remembered one is still granted: if a secret's allowed scopes change, restart the process (or rotate the secret) to pick up a broader scope. Passing oauth_scopes turns the fallback off. To choose scopes, pass e.g. oauth_scopes=["vector-search", "model-serving"]; include ai-gateway if any index embeds through a Unity Catalog model service.

See the docstrings on AISearchClient.__init__ for details.

Use with AI coding agents

The package ships a guide for coding agents and a databricks-ai-search command.

# Print the guide, or one topic: quickstart, high-qps, debugging, load-testing
databricks-ai-search guide
databricks-ai-search guide --topic high-qps

# Install it where an agent reads it, in the current directory (or --dir PATH)
databricks-ai-search guide --install claude     # .claude/skills/databricks-ai-search/
databricks-ai-search guide --install agents-md  # a marked section in AGENTS.md (also: codex)
databricks-ai-search guide --install cursor     # .cursor/rules/databricks-ai-search.mdc

Installing never overwrites a file the command did not create unless you pass --force (claude and cursor only), never writes through a symlink, and replaces only its own marked section of AGENTS.md, keeping the file's line endings. Each install writes a new file and renames it over the old one: the file mode is kept, but the owner, ACLs, extended attributes, and hard links are not.

databricks-ai-search doctor checks one index without changing anything: auth mode, OAuth scope, endpoint state and target QPS, index type and embedding source, query route, embedding serving endpoints, and one test query. On a managed-embedding index the test query calls the embedding endpoint; --no-query skips it. It reads DATABRICKS_HOST plus either DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET or DATABRICKS_TOKEN (or --profile NAME from ~/.databrickscfg), redacts the configured token and secret and common token formats from its output, and exits 1 if a check fails. Text returned by the server is reported separately as server_message.

databricks-ai-search doctor --index my_catalog.my_schema.my_index
databricks-ai-search doctor --index my_catalog.my_schema.my_index --json

Flags: --endpoint NAME (default: the index's endpoint), --oauth-scopes "SCOPE ..." (as passed to AISearchClient(oauth_scopes=...)), --azure-tenant-id ID (Microsoft Entra ID service principal), --profile NAME (read credentials from ~/.databrickscfg), --no-query (skip the test query), and --json. Exit codes: 0 when no check fails, 1 when any check fails, 2 for invalid arguments. Notebook and Model Serving credentials that AISearchClient() detects on its own are not tested.

databricks-ai-search load-test runs the endpoint page's load test on a STANDARD endpoint that has a target QPS, as a one-time job, and prints its verdict (PASS, REVIEW or FAIL), findings, recommended next steps, and per-stage throughput and latency. databricks-ai-search set-target-qps changes or resets an endpoint's target QPS. Both change nothing without --yes: they print what they would do and exit 5 (or exit 1 if a check blocks the change, and set-target-qps exits 0 when the endpoint already has the value). load-test exits 0 PASS (or submitted with --no-wait), 3 FAIL, 4 REVIEW, 1 failed or blocked (including unrecognized_verdict, a verdict from a newer notebook: upgrade databricks-ai-search, then rerun with --run-id), 5 not confirmed. Read the outcome field of --json output for what happened. A load test sends real queries and runs serverless compute, and target QPS capacity is billed whether or not traffic arrives, so both cost money. See databricks-ai-search load-test --help for which credential is used for what. --query-type HYBRID tests hybrid (vector and keyword) queries; test with the query type your application uses, since an ANN PASS overstates hybrid capacity. Hybrid search uses about 4x the resources of ANN in our Target QPS benchmarks. For an index with self-managed embeddings, --text-column names the source table's string column, synced to the index, that query text comes from.

databricks-ai-search load-test --index my_catalog.my_schema.my_index --stages quick
databricks-ai-search load-test --index my_catalog.my_schema.my_index --stages quick --yes \
  --sp-client-id <client-id> --sp-secret-scope <scope> --sp-secret-key <key>
databricks-ai-search load-test --index my_catalog.my_schema.my_index --query-type HYBRID \
  --text-column body
databricks-ai-search set-target-qps --endpoint my-endpoint --target-qps 500 --wait --yes

DB license

Copyright (2022) Databricks, Inc.

This library (the "Software") may not be used except in connection with the Licensee's use of the Databricks Platform Services pursuant to an Agreement (defined below) between Licensee (defined below) and Databricks, Inc. ("Databricks"). This Software shall be deemed part of the Downloadable Services under the Agreement, or if the Agreement does not define Downloadable Services, Subscription Services, or if neither are defined then the term in such Agreement that refers to the applicable Databricks Platform Services (as defined below) shall be substituted herein for "Downloadable Services". Licensee's use of the Software must comply at all times with any restrictions applicable to the Downloadable Services and Subscription Services, generally, and must be used in accordance with any applicable documentation.

Additionally, and notwithstanding anything in the Agreement to the contrary:

• THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

If you have not agreed to an Agreement or otherwise do not agree to these terms, you may not use the Software.

This license terminates automatically upon the termination of the Agreement or Licensee's breach of these terms.

Agreement: the agreement between Databricks and Licensee governing the use of the Databricks Platform Services, which shall be, with respect to Databricks, the Databricks Terms of Service located at www.databricks.com/termsofservice, and with respect to Databricks Community Edition, the Community Edition Terms of Service located at www.databricks.com/ce-termsofuse, in each case unless Licensee has entered into a separate written agreement with Databricks governing the use of the applicable Databricks Platform Services.

Databricks Platform Services: the Databricks services or the Databricks Community Edition services, according to where the Software is used.

Licensee: the user of the Software, or, if the Software is being used on behalf of a company, the company.