databricks-ai-search 0.73__py3-none-any.whl

databricks/__init__.py

@@ -0,0 +1,7 @@
+# See: https://packaging.python.org/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages
+#
+# This file must only contain the following line, or other packages in the databricks.* namespace
+# may not be importable. The contents of this file must be byte-for-byte equivalent across all packages.
+# If they are not, parallel package installation may lead to clobbered and invalid files.
+# Also see https://github.com/databricks/databricks-sdk-py/issues/343.
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

databricks/ai_search/__init__.py

@@ -0,0 +1,3 @@
+from databricks.ai_search.version import VERSION as __version__
+
+__all__ = ["__version__"]

databricks/ai_search/async_index.py

@@ -0,0 +1,348 @@
+import asyncio
+from typing import Any, Dict, List, Optional, Union
+
+from databricks.ai_search.async_utils import AsyncRequestUtils
+from databricks.ai_search.index import AISearchIndex
+from databricks.ai_search.reranker import Reranker
+from databricks.ai_search.utils import (
+    CredentialStrategy,
+    normalize_response_keys_to_snake_case,
+)
+
+
+class AsyncAISearchIndex:
+    """
+    AsyncAISearchIndex is a helper class that represents a AI Search Index for ``async/await`` code.
+
+    Those who wish to use this class should not instantiate it directly, but rather use
+    :meth:`AISearchClient.get_async_index` to obtain an instance.
+
+    Index lifecycle operations (create, delete, ``wait_until_ready``) are available
+    only on :class:`AISearchClient`.
+
+    Example::
+
+        async with client.get_async_index(
+            endpoint_name="my-endpoint", index_name="my-index"
+        ) as idx:
+            results = await idx.similarity_search(
+                query_vector=[0.1, 0.2, ...], columns=["id"], num_results=10
+            )
+    """
+
+    def __init__(
+        self,
+        workspace_url: str,
+        index_url: str,
+        name: str,
+        endpoint_name: str,
+        mlserving_endpoint_name: Optional[str] = None,
+        personal_access_token: Optional[str] = None,
+        service_principal_client_id: Optional[str] = None,
+        service_principal_client_secret: Optional[str] = None,
+        azure_tenant_id: Optional[str] = None,
+        azure_login_id: Optional[str] = None,
+        use_user_passed_credentials: bool = False,
+        credential_strategy: Optional[CredentialStrategy] = None,
+        get_reranker_url_callable: Optional[callable] = None,
+        mlserving_endpoint_name_for_query: Optional[str] = None,
+        total_retries: int = 3,
+        backoff_factor: float = 1,
+        backoff_jitter: float = 0.2,
+    ):
+        """
+        Initialize an AsyncAISearchIndex instance.
+
+        :param str workspace_url: The URL of the Databricks workspace.
+        :param str index_url: The direct URL to the vector search index endpoint.
+        :param str name: The name of the vector search index.
+        :param str endpoint_name: The name of the vector search endpoint.
+        :param str mlserving_endpoint_name: The name of the model serving endpoint used for embedding generation during ingestion.
+        :param str personal_access_token: Personal access token for authentication.
+        :param str service_principal_client_id: Service principal client ID for authentication.
+        :param str service_principal_client_secret: Service principal client secret for authentication.
+        :param str azure_tenant_id: Azure tenant ID for Azure-based authentication.
+        :param str azure_login_id: Azure login ID (Databricks Azure Application ID) for authentication.
+        :param bool use_user_passed_credentials: Whether credentials were explicitly provided by the user (True) or inferred automatically (False).
+        :param CredentialStrategy credential_strategy: The credential strategy to use for authentication.
+        :param callable get_reranker_url_callable: A callable function to retrieve the reranker-compatible index URL when needed.
+        :param str mlserving_endpoint_name_for_query: The name of the model serving endpoint to use for queries (if different from ingestion endpoint).
+        :param int total_retries: Total number of retries for requests. Defaults to 3.
+        :param float backoff_factor: Backoff factor for retry delays. Defaults to 1.
+        :param float backoff_jitter: Random jitter proportion (0-1) to add to backoff delays. Defaults to 0.2.
+        """
+        self._sync = AISearchIndex(
+            workspace_url=workspace_url,
+            index_url=index_url,
+            name=name,
+            endpoint_name=endpoint_name,
+            mlserving_endpoint_name=mlserving_endpoint_name,
+            personal_access_token=personal_access_token,
+            service_principal_client_id=service_principal_client_id,
+            service_principal_client_secret=service_principal_client_secret,
+            azure_tenant_id=azure_tenant_id,
+            azure_login_id=azure_login_id,
+            use_user_passed_credentials=use_user_passed_credentials,
+            credential_strategy=credential_strategy,
+            get_reranker_url_callable=get_reranker_url_callable,
+            mlserving_endpoint_name_for_query=mlserving_endpoint_name_for_query,
+            total_retries=total_retries,
+            backoff_factor=backoff_factor,
+            backoff_jitter=backoff_jitter,
+        )
+        self.workspace_url = workspace_url
+        self.name = name
+        self.endpoint_name = endpoint_name
+        self.index_url = self._sync.index_url
+        self.total_retries = total_retries
+        self.backoff_factor = backoff_factor
+        self.backoff_jitter = backoff_jitter
+        self._closed = False
+        self._token_lock = asyncio.Lock()
+
+    async def upsert(self, inputs):
+        """
+        Upsert data into the index.
+
+        :param inputs: List of dictionaries to upsert into the index.
+        """
+        payload = self._sync._build_upsert_payload(inputs)
+        token = await self._aget_token(write=True)
+        return await AsyncRequestUtils.issue_request(
+            url=f"{self.index_url}/upsert-data",
+            token=token,
+            method="POST",
+            json=payload,
+            total_retries=self.total_retries,
+            backoff_factor=self.backoff_factor,
+            backoff_jitter=self.backoff_jitter,
+        )
+
+    async def delete(self, primary_keys):
+        """
+        Delete data from the index.
+
+        :param primary_keys: List of primary keys to delete from the index.
+        """
+        payload = self._sync._build_delete_payload(primary_keys)
+        token = await self._aget_token(write=True)
+        return await AsyncRequestUtils.issue_request(
+            url=f"{self.index_url}/delete-data",
+            token=token,
+            method="DELETE",
+            json=payload,
+            total_retries=self.total_retries,
+            backoff_factor=self.backoff_factor,
+            backoff_jitter=self.backoff_jitter,
+        )
+
+    async def describe(self):
+        """
+        Describe the index. This returns metadata about the index.
+        """
+        token = await self._aget_token(control_plane=True)
+        return await AsyncRequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/indexes/{self.name}",
+            token=token,
+            method="GET",
+            total_retries=self.total_retries,
+            backoff_factor=self.backoff_factor,
+            backoff_jitter=self.backoff_jitter,
+        )
+
+    async def sync(self):
+        """
+        Sync the index. This is used to sync the index with the source delta table.
+        This only works with managed delta sync index with pipeline type="TRIGGERED".
+        """
+        token = await self._aget_token(control_plane=True)
+        return await AsyncRequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/indexes/{self.name}/sync",
+            token=token,
+            method="POST",
+            total_retries=self.total_retries,
+            backoff_factor=self.backoff_factor,
+            backoff_jitter=self.backoff_jitter,
+        )
+
+    async def scan(self, num_results=10, last_primary_key=None):
+        """
+        Given all the data in the index sorted by primary key, this returns the next
+        ``num_results`` data after the primary key specified by ``last_primary_key``.
+        If ``last_primary_key`` is None, it returns the first ``num_results``.
+
+        Please note if there's ongoing updates to the index, the scan results may not be consistent.
+
+        :param num_results: Number of results to return.
+        :param last_primary_key: Last primary key from previous pagination, used as the exclusive starting primary key.
+        """
+        payload = self._sync._build_scan_payload(num_results, last_primary_key)
+        token = await self._aget_token()
+        response = await AsyncRequestUtils.issue_request(
+            url=f"{self.index_url}/scan",
+            token=token,
+            method="GET",
+            json=payload,
+            total_retries=self.total_retries,
+            backoff_factor=self.backoff_factor,
+            backoff_jitter=self.backoff_jitter,
+        )
+        return normalize_response_keys_to_snake_case(response)
+
+    async def similarity_search(
+        self,
+        columns: List[str],
+        query_text: Optional[str] = None,
+        query_vector: Optional[List[float]] = None,
+        filters: Optional[Union[str, Dict[str, Any]]] = None,
+        num_results: int = 5,
+        debug_level: int = 0,
+        score_threshold: Optional[float] = None,
+        query_type: Optional[str] = None,
+        columns_to_rerank: Optional[List[str]] = None,
+        disable_notice: bool = False,
+        reranker: Optional[Reranker] = None,
+        total_retries: int = 3,
+        backoff_factor: float = 1,
+        backoff_jitter: float = 0.2,
+        *,
+        query_columns: Optional[List[str]] = None,
+        sort_columns: Optional[List[str]] = None,
+        facets: Optional[List[str]] = None,
+        _query_experiment_config: Optional[Dict[str, Any]] = None,
+    ):
+        """similarity_search(columns: List[str], query_text: Optional[str] = None, query_vector: Optional[List[float]] = None, filters: Optional[Union[str, Dict[str, Any]]] = None, num_results: int = 5, debug_level: int = 0, score_threshold: Optional[float] = None, query_type: Optional[str] = None, columns_to_rerank: Optional[List[str]] = None, disable_notice: bool = False, reranker: Optional[Reranker] = None, total_retries: int = 3, backoff_factor: float = 1, backoff_jitter: float = 0.2, query_columns: Optional[List[str]] = None, sort_columns: Optional[List[str]] = None, facets: Optional[List[str]] = None)
+
+        Perform a similarity search on the index. This returns the top K results that are most similar to the query.
+
+        :param columns: List of column names to return in the results.
+        :param query_text: Query text to search for.
+        :param query_vector: Query vector to search for.
+        :param filters: Filters to apply to the query.
+        :param num_results: Number of results to return.
+        :param debug_level: Debug level to use for the query.
+        :param score_threshold: Score threshold to use for the query. If reranker is used, the score threshold is applied before reranking.
+        :param query_type: Query type of this query. Choices are "ANN", "HYBRID", and "FULL_TEXT".
+        :param query_columns: Text columns to search for query_text. When empty, all text columns are searched.
+        :param sort_columns: Sort results by column values instead of the default relevance ordering. Each clause has the form "<column> ASC" or "<column> DESC", for example ["rating DESC", "price ASC"].
+        :param facets: Facets to compute over the matched results. Each entry is one of: "<column>" (top 10 distinct values by count), "<column> TOP <n>" (top n distinct values, n > 0), or "<column> BUCKETS [[from,to],...]" (inclusive numeric ranges). TOP and BUCKETS are case-insensitive; a column may appear at most once.
+        :param columns_to_rerank: (Deprecated) List of column names to use for reranking the results.
+            Use the ``reranker`` parameter instead.
+        :param disable_notice: Whether to disable the notice message.
+        :param reranker: Optional reranker to apply on the top results. Pass an instance of
+            :class:`databricks.ai_search.reranker.DatabricksReranker` with
+            ``columns_to_rerank=[...]``. The reranker reorders the initial results using
+            the specified text columns.
+        :type reranker: Optional[:class:`databricks.ai_search.reranker.Reranker`]
+        :param total_retries: Total number of retries for the request. Set to 0 to disable retries.
+        :param backoff_factor: Backoff factor to apply between retry attempts. The delay between retries
+            is calculated as {backoff_factor} * (2 ** (retry_count - 1)) seconds. For example, with
+            backoff_factor=1, delays are 0.5s, 1s, 2s, 4s, etc.
+        :param backoff_jitter: Random jitter to add to backoff delays to avoid thundering herd problem.
+            Value between 0 and 1 representing the proportion of jitter to apply.
+
+        Example:
+            Use the Databricks reranker to improve the ordering of hybrid search results:
+
+            .. code-block:: python
+
+                from databricks.ai_search.reranker import DatabricksReranker
+
+                results = index.similarity_search(
+                    query_text=\"How to create a AI Search index\",
+                    columns=[\"id\", \"text\", \"parent_doc_summary\", \"date\"],
+                    # The final number of results to return. The reranker will automatically overfetch 50 documents and rerank them.
+                    num_results=10,
+                    query_type=\"hybrid\",
+                    # Needed for debug info to get any warnings and time to rerank the results.
+                    debug_level=1,
+                    # The text reranked will be concatenated and if it is longer than 2000 characters, it will be truncated.
+                    # Include shorter, important columns first.
+                    reranker=DatabricksReranker(columns_to_rerank=[\"parent_doc_summary\", \"text\", \"other_column\"]),
+                )
+                # Check if reranking was successful and how much additional time it took to rerank the results.
+                if "warnings" in results['debug_info']:
+                    print(results['debug_info']['warnings'])
+                else:
+                    print(f"Reranking was successful and took {results['debug_info']['reranker_time']}ms")
+        """
+        json_data, is_reranker = self._sync._build_query_payload(
+            columns=columns,
+            query_text=query_text,
+            query_vector=query_vector,
+            filters=filters,
+            num_results=num_results,
+            debug_level=debug_level,
+            score_threshold=score_threshold,
+            query_type=query_type,
+            columns_to_rerank=columns_to_rerank,
+            disable_notice=disable_notice,
+            reranker=reranker,
+            query_columns=query_columns,
+            sort_columns=sort_columns,
+            facets=facets,
+            _query_experiment_config=_query_experiment_config,
+        )
+        query_url = await asyncio.to_thread(self._sync._resolve_query_url, is_reranker)
+
+        token = await self._aget_token(index_url=query_url, control_plane=is_reranker)
+        response = await AsyncRequestUtils.issue_request(
+            url=f"{query_url}/query",
+            token=token,
+            method="GET",
+            json=json_data,
+            total_retries=total_retries,
+            backoff_factor=backoff_factor,
+            backoff_jitter=backoff_jitter,
+        )
+        response = normalize_response_keys_to_snake_case(response)
+
+        accumulated = response
+        while response.get("next_page_token"):
+            response = await self._aget_next_page(query_url, response["next_page_token"], control_plane=is_reranker)
+            accumulated["result"]["row_count"] += response["result"]["row_count"]
+            accumulated["result"]["data_array"] += response["result"]["data_array"]
+
+        accumulated.pop("next_page_token", None)
+        return accumulated
+
+    async def aclose(self):
+        """
+        Close this index instance.
+        """
+        self._closed = True
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.aclose()
+
+    async def _aget_token(self, write=False, control_plane=False, index_url=None):
+        async with self._token_lock:
+            return await asyncio.to_thread(
+                self._sync._get_token_for_request,
+                write=write,
+                control_plane=control_plane,
+                index_url=index_url,
+            )
+
+    async def _aget_next_page(self, index_url, page_token, control_plane=False):
+        payload = self._sync._build_next_page_payload(page_token)
+        token = await self._aget_token(control_plane=control_plane)
+        response = await AsyncRequestUtils.issue_request(
+            url=f"{index_url}/query-next-page",
+            token=token,
+            method="GET",
+            json=payload,
+            total_retries=self.total_retries,
+            backoff_factor=self.backoff_factor,
+            backoff_jitter=self.backoff_jitter,
+        )
+        return normalize_response_keys_to_snake_case(response)
+
+
+# Backward-compat alias for the AI Search rebrand. The canonical class is
+# AsyncAISearchIndex; AsyncVectorSearchIndex is the same class accessible under
+# the legacy name.
+AsyncVectorSearchIndex = AsyncAISearchIndex

databricks/ai_search/async_utils.py

@@ -0,0 +1,204 @@
+"""Async HTTP transport for the Databricks AI Search SDK.
+
+Mirrors :class:`databricks.ai_search.utils.RequestUtils` over
+``httpx.AsyncClient``. Status codes map to the same exception hierarchy as the
+sync path so callers can use a single ``except`` block regardless of which
+client they use.
+"""
+
+import asyncio
+import datetime
+import random
+from email.utils import parsedate_to_datetime
+from typing import Dict, Optional
+
+import httpx
+
+from databricks.ai_search.exceptions import (
+    BadRequest,
+    NotFound,
+    PermissionDenied,
+    ResourceConflict,
+    TooManyRequests,
+    AISearchException,
+)
+from databricks.ai_search.version import VERSION
+
+
+_RETRYABLE_STATUS_CODES = {429, 503, 504}
+# Match the sync path's urllib3.Retry network-error coverage: TimeoutException
+# spans the connect/read/write/pool timeouts, NetworkError spans connect/read/
+# write/close failures, and RemoteProtocolError catches mid-stream resets.
+_RETRYABLE_EXCEPTIONS = (
+    httpx.TimeoutException,
+    httpx.NetworkError,
+    httpx.RemoteProtocolError,
+)
+
+
+# Per-event-loop client cache so one httpx.AsyncClient (and its connection
+# pool) is reused across requests on the same loop. Sync mirrors this with
+# the functools.cache-decorated _cached_get_request_session in utils.py.
+_clients: Dict[int, httpx.AsyncClient] = {}
+
+
+def _parse_retry_after(value):
+    """Parse a Retry-After header value as either seconds or an HTTP-date.
+
+    Returns the delay in seconds, or None if the header cannot be parsed.
+    """
+    if value is None:
+        return None
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        pass
+    try:
+        retry_date = parsedate_to_datetime(value)
+    except (TypeError, ValueError):
+        return None
+    if retry_date is None:
+        return None
+    now = datetime.datetime.now(datetime.timezone.utc)
+    if retry_date.tzinfo is None:
+        retry_date = retry_date.replace(tzinfo=datetime.timezone.utc)
+    delta = (retry_date - now).total_seconds()
+    return max(delta, 0.0)
+
+
+async def _get_client_for_loop():
+    """Return a cached ``httpx.AsyncClient`` bound to the running loop.
+
+    ``httpx.AsyncClient`` holds connection pools that are tied to the event
+    loop that created them; sharing across loops produces obscure errors.
+    Keying the cache by ``id(loop)`` keeps construction lazy and per-loop.
+    The body has no ``await`` between the cache check and insert, so within
+    a single loop the check-and-set is atomic.
+    """
+    loop = asyncio.get_running_loop()
+    key = id(loop)
+    client = _clients.get(key)
+    if client is None or client.is_closed:
+        client = httpx.AsyncClient(timeout=None)
+        _clients[key] = client
+    return client
+
+
+async def aclose_all_clients():
+    """Close every cached ``httpx.AsyncClient``. Intended for tests and orderly shutdown.
+
+    Do not call while requests are in flight — this closes the underlying
+    connection pools, which will cause any awaiting coroutines to raise
+    ``httpx.ClientError``.
+    """
+    clients = list(_clients.values())
+    _clients.clear()
+    for client in clients:
+        try:
+            await client.aclose()
+        except (httpx.HTTPError, OSError):
+            pass
+
+
+class AsyncRequestUtils:
+    """Async equivalent of :class:`RequestUtils`.
+
+    Issues a single HTTP request with retry/backoff and translates non-2xx
+    responses into the SDK's typed exceptions. ``httpx`` has no built-in
+    retry support, so the loop is implemented here.
+    """
+
+    @staticmethod
+    async def issue_request(
+        url: str,
+        method: str,
+        token: Optional[str] = None,
+        params: Optional[dict] = None,
+        json: Optional[dict] = None,
+        auth: Optional[tuple] = None,
+        data: Optional[dict] = None,
+        headers: Optional[dict] = None,
+        total_retries: int = 3,
+        backoff_factor: float = 1.0,
+        backoff_jitter: float = 0.2,
+        client: Optional[httpx.AsyncClient] = None,
+    ):
+        """Issue an async HTTP request and return its JSON body.
+
+        Retries up to ``total_retries`` times on transient failures (HTTP
+        429/503/504 plus httpx ``TimeoutException`` / ``NetworkError`` /
+        ``RemoteProtocolError``). Honors ``Retry-After`` when present;
+        otherwise sleeps for
+        ``backoff_factor * (2 ** (attempt - 1)) * (1 + backoff_jitter * random())``
+        seconds between attempts. Non-retryable 4xx responses raise immediately.
+        """
+        headers = dict(headers) if headers else {}
+        if token:
+            headers["Authorization"] = f"Bearer {token}"
+        headers["X-Databricks-Python-SDK-Version"] = VERSION
+
+        owns_client = client is None
+        if owns_client:
+            client = await _get_client_for_loop()
+
+        last_exc: Optional[BaseException] = None
+        response: Optional[httpx.Response] = None
+        for attempt in range(total_retries + 1):
+            try:
+                response = await client.request(
+                    method=method,
+                    url=url,
+                    params=params,
+                    json=json,
+                    data=data,
+                    headers=headers,
+                    auth=auth,
+                )
+                if response.status_code in _RETRYABLE_STATUS_CODES and attempt < total_retries:
+                    delay = _parse_retry_after(response.headers.get("Retry-After"))
+                    if delay is None:
+                        delay = backoff_factor * (2**attempt) * (1 + backoff_jitter * random.random())
+                    await asyncio.sleep(delay)
+                    continue
+                break
+            except _RETRYABLE_EXCEPTIONS as exc:
+                last_exc = exc
+                if attempt >= total_retries:
+                    raise
+                delay = backoff_factor * (2**attempt) * (1 + backoff_jitter * random.random())
+                await asyncio.sleep(delay)
+
+        if response is None:
+            raise AISearchException(
+                f"Request to {url} failed: {last_exc!r}",
+                status_code=None,
+                response_content=None,
+            )
+
+        if response.is_success:
+            return response.json()
+
+        status_code = response.status_code
+        error_message = f"HTTP {status_code}"
+        try:
+            error_data = response.json()
+            if isinstance(error_data, dict):
+                error_message = error_data.get("message") or error_data.get("error") or str(error_data)
+        except ValueError:
+            # json.JSONDecodeError is a ValueError subclass; we catch the parent
+            # because the `json` parameter on this method shadows the stdlib name
+            # inside this scope.
+            error_message = response.text or f"HTTP {status_code} error"
+
+        content = response.content
+        if status_code == 400:
+            raise BadRequest(error_message, status_code=status_code, response_content=content)
+        if status_code == 403:
+            raise PermissionDenied(error_message, status_code=status_code, response_content=content)
+        if status_code == 404:
+            raise NotFound(error_message, status_code=status_code, response_content=content)
+        if status_code == 409:
+            raise ResourceConflict(error_message, status_code=status_code, response_content=content)
+        if status_code == 429:
+            raise TooManyRequests(error_message, status_code=status_code, response_content=content)
+        raise AISearchException(error_message, status_code=status_code, response_content=content)