Metadata-Version: 2.1
Name: databricks-ai-search
Version: 0.82
Summary: Databricks AI Search Client (formerly Vector Search)
Home-page: UNKNOWN
Author: Databricks
Author-email: feedback@databricks.com
License: UNKNOWN
Platform: UNKNOWN
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE.md
Requires-Dist: mlflow-skinny (<4.0,>=3.10.1)
Requires-Dist: requests (<3.0,>=2.32.5)
Requires-Dist: deprecation (<3.0,>=2.1.0)
Requires-Dist: httpx (<1.0,>=0.28.0)

# Databricks AI Search Python SDK

The Python SDK for [Databricks AI Search](https://docs.databricks.com/aws/en/ai-search/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

```bash
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

```python
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.

```bash
# 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`.

```bash
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.

```bash
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.


