databricks-vectorsearch 0.37__tar.gz

databricks_vectorsearch-0.37/PKG-INFO

@@ -0,0 +1,45 @@
+Metadata-Version: 2.1
+Name: databricks_vectorsearch
+Version: 0.37
+Summary: Databricks Vector Search Client
+Home-page: UNKNOWN
+Author: Databricks
+Author-email: feedback@databricks.com
+License: UNKNOWN
+Platform: UNKNOWN
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+
+**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 Downlodable 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.
+
+
+

databricks_vectorsearch-0.37/README.md

@@ -0,0 +1,31 @@
+**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 Downlodable 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.
+

databricks_vectorsearch-0.37/databricks/vector_search/__init__.py

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

databricks_vectorsearch-0.37/databricks/vector_search/client.py

@@ -0,0 +1,474 @@
+import time
+import json
+import datetime
+import math
+from databricks.vector_search.exceptions import InvalidInputException
+from databricks.vector_search.utils import OAuthTokenUtils
+from databricks.vector_search.utils import RequestUtils
+from databricks.vector_search.index import VectorSearchIndex
+from mlflow.utils import databricks_utils
+
+class VectorSearchClient:
+    """
+    A client for interacting with the Vector Search service.
+
+    This client provides methods for managing endpoints and indexes in the Vector Search service.
+    """
+
+    def __init__(
+        self,
+        workspace_url=None,
+        personal_access_token=None,
+        service_principal_client_id=None,
+        service_principal_client_secret=None,
+        azure_tenant_id=None,
+        azure_login_id=None,
+        disable_notice=False,
+    ):
+        """
+        Initialize the VectorSearchClient.
+
+        :param str workspace_url: The URL of the workspace.
+        :param str personal_access_token: The personal access token for authentication.
+        :param str service_principal_client_id: The client ID of the service principal for authentication.
+        :param str service_principal_client_secret: The client secret of the service principal for authentication.
+        :param str azure_tenant_id: The tenant ID of Azure for authentication.
+        :param str azure_login_id: The login ID of Azure for authentication.
+        :param bool disable_notice: Whether to disable the notice message.
+        """
+        self.workspace_url = workspace_url
+        self.personal_access_token = personal_access_token
+        self.service_principal_client_id = service_principal_client_id
+        self.service_principal_client_secret = service_principal_client_secret
+        self.azure_tenant_id = azure_tenant_id
+        self.azure_login_id = azure_login_id
+        self._is_notebook_pat = False
+        # whether or not credentials are explicitly passed in by user in client or inferred by client
+        # via mlflow utilities. If passed in by user, continue to use user credentials in index object.
+        # If not, can attempt automatic auth refresh for model serving.
+        self._using_user_passed_credentials = bool(
+            (self.service_principal_client_id and self.service_principal_client_secret) or \
+            (self.workspace_url and self.personal_access_token))
+
+        if not (
+            self.service_principal_client_id and
+            self.service_principal_client_secret
+        ):
+            if self.workspace_url is None:
+                host_creds = databricks_utils.get_databricks_host_creds()
+                self.workspace_url = host_creds.host
+            if self.personal_access_token is None:
+                self._is_notebook_pat = True
+                host_creds = databricks_utils.get_databricks_host_creds()
+                self.personal_access_token = host_creds.token
+
+        self._control_plane_oauth_token = None
+        self._control_plane_oauth_token_expiry_ts = None
+        self.validate(disable_notice=disable_notice)
+
+    def validate(self, disable_notice=False):
+        if not (self.personal_access_token or
+                (self.service_principal_client_id and 
+                    self.service_principal_client_secret)):
+            raise InvalidInputException(
+                "Please specify either personal access token or service principal client ID and secret."
+            )
+        if (self.service_principal_client_id and
+                self.service_principal_client_secret and
+                not self.workspace_url):
+            raise InvalidInputException(
+                "Service Principal auth flow requires workspace url"
+            )
+        if self._is_notebook_pat and not disable_notice:
+            print(
+                """[NOTICE] Using a notebook authentication token. Recommended for development only. For improved performance, please use Service Principal based authentication. To disable this message, pass disable_notice=True to VectorSearchClient()."""
+            )
+        elif self.personal_access_token and not disable_notice:
+            print(
+                """[NOTICE] Using a Personal Authentication Token (PAT). Recommended for development only. For improved performance, please use Service Principal based authentication. To disable this message, pass disable_notice=True to VectorSearchClient()."""
+            )
+
+
+    def _get_token_for_request(self):
+        if self.personal_access_token:
+            return self.personal_access_token
+        if (
+            self._control_plane_oauth_token
+            and self._control_plane_oauth_token_expiry_ts
+            and self._control_plane_oauth_token_expiry_ts - 100 > time.time()
+        ):
+            return self._control_plane_oauth_token
+        if self.service_principal_client_id and \
+                self.service_principal_client_secret:
+            authorization_details = []
+            oauth_token_data = OAuthTokenUtils.get_oauth_token(
+                    workspace_url=self.workspace_url,
+                    service_principal_client_id=self.service_principal_client_id,
+                    service_principal_client_secret=self.service_principal_client_secret,
+                    authorization_details=authorization_details,
+            ) if not self.azure_tenant_id else OAuthTokenUtils.get_azure_oauth_token(
+                    workspace_url=self.workspace_url,
+                    service_principal_client_id=self.service_principal_client_id,
+                    service_principal_client_secret=self.service_principal_client_secret,
+                    authorization_details=authorization_details,
+                    azure_tenant_id=self.azure_tenant_id,
+                    azure_login_id=self.azure_login_id
+            )
+            self._control_plane_oauth_token = oauth_token_data["access_token"]
+            self._control_plane_oauth_token_expiry_ts = time.time() + oauth_token_data["expires_in"]
+            return self._control_plane_oauth_token
+        raise Exception("You must specify service principal or PAT token")
+
+    def create_endpoint(self, name, endpoint_type="STANDARD"):
+        """
+        Create an endpoint.
+
+        :param str name: The name of the endpoint.
+        :param str endpoint_type: The type of the endpoint. Must be STANDARD or ENTERPRISE.
+        """
+        json_data = {"name": name, "endpoint_type": endpoint_type}
+        return RequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/endpoints",
+            token=self._get_token_for_request(),
+            method="POST",
+            json=json_data,
+        )
+
+    def get_endpoint(self, name):
+        """
+        Get an endpoint.
+
+        :param str name: The name of the endpoint.
+        """
+        return RequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/endpoints/{name}",
+            token=self._get_token_for_request(),
+            method="GET",
+        )
+
+    def list_endpoints(self):
+        """
+        List all endpoints.
+        """
+        return RequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/endpoints",
+            token=self._get_token_for_request(),
+            method="GET",
+        )
+
+    def delete_endpoint(self, name):
+        """
+        Delete an endpoint.
+
+        :param str name: The name of the endpoint.
+        """
+        return RequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/endpoints/{name}",
+            token=self._get_token_for_request(),
+            method="DELETE",
+        )
+    
+    def create_endpoint_and_wait(self, name, endpoint_type="STANDARD", verbose=False, timeout=datetime.timedelta(minutes=60)):
+        """
+        Create an endpoint and wait for it to be online.
+
+        :param str name: The name of the endpoint.
+        :param str endpoint_type: The type of the endpoint. Must be STANDARD or ENTERPRISE.
+        :param bool verbose: Whether to print status messages.
+        :param datetime.timedelta timeout: The time allowed until we timeout with an Exception.
+        """
+        all_endpoints = [endpoint['name'] for endpoint in self.list_endpoints()['endpoints']]
+        if name in all_endpoints:
+            raise Exception(f"Endpoint {name} already exists.")
+        if verbose:
+            print(f"Creating endpoint {name}.")
+        self.create_endpoint(name, endpoint_type)
+        self.wait_for_endpoint(name, verbose, timeout)
+
+    def wait_for_endpoint(self, name, verbose=False, timeout=datetime.timedelta(minutes=60)):
+        """
+        Wait for an endpoint to be online.
+
+        :param str name: The name of the endpoint.
+        :param bool verbose: Whether to print status messages.
+        :param datetime.timedelta timeout: The time allowed until we timeout with an Exception.
+        """
+
+        def get_endpoint_state():
+            endpoint = self.get_endpoint(name)
+            endpoint_state = endpoint["endpoint_status"]["state"]
+            return endpoint_state
+
+        start_time = datetime.datetime.now()
+        sleep_time_seconds = min(30, timeout.total_seconds())
+        # Possible states are "ONLINE", "OFFLINE", and "PROVISIONING".
+        endpoint_state = get_endpoint_state()
+        while endpoint_state != "ONLINE" and datetime.datetime.now() - start_time < timeout:
+            if endpoint_state == "OFFLINE":
+                raise Exception(f"Endpoint {name} is OFFLINE.")
+            if verbose:
+                running_time = int(math.floor((datetime.datetime.now() - start_time).total_seconds()))
+                print(f"Endpoint {name} is in state {endpoint_state}. Time: {running_time}s.")
+            time.sleep(sleep_time_seconds)
+            endpoint_state = get_endpoint_state()
+        if endpoint_state == "ONLINE":
+          if verbose:
+            print(f"Endpoint {name} is ONLINE.")
+        else:
+            raise Exception(f"Endpoint {name} did not become ONLINE within timeout of {timeout.total_seconds()}s.")
+
+
+    def list_indexes(self, name):
+        """
+        List all indexes for an endpoint.
+
+        :param str name: The name of the endpoint.
+        """
+        return RequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/endpoints/{name}/indexes",
+            token=self._get_token_for_request(),
+            method="GET",
+        )
+
+    def create_delta_sync_index(
+            self,
+            endpoint_name,
+            index_name,
+            primary_key,
+            source_table_name,
+            pipeline_type,
+            embedding_dimension=None,
+            embedding_vector_column=None,
+            embedding_source_column=None,
+            embedding_model_endpoint_name=None,
+            sync_computed_embeddings=False
+    ):
+        """
+        Create a delta sync index.
+
+        :param str endpoint_name: The name of the endpoint.
+        :param str index_name: The name of the index.
+        :param str primary_key: The primary key of the index.
+        :param str source_table_name: The name of the source table.
+        :param str pipeline_type: The type of the pipeline. Must be CONTINUOUS or TRIGGERED.
+        :param int embedding_dimension: The dimension of the embedding vector.
+        :param str embedding_vector_column: The name of the embedding vector column.
+        :param str embedding_source_column: The name of the embedding source column.
+        :param str embedding_model_endpoint_name: The name of the embedding model endpoint.
+        :param bool sync_computed_embeddings: Whether to automatically sync the vector index contents and computed embeddings to a new UC table,
+                                             table name will be ${index_name}_writeback_table.
+        """
+        assert pipeline_type, "Pipeline type cannot be None. Please use CONTINUOUS/TRIGGERED as the pipeline type."
+        json_data = {
+            "name": index_name,
+            "index_type": "DELTA_SYNC",
+            "primary_key": primary_key,
+            "delta_sync_index_spec": {
+                "source_table": source_table_name,
+                "pipeline_type": pipeline_type.upper(),
+            }
+        }
+        if embedding_vector_column:
+            assert embedding_dimension, "Embedding dimension must be specified if source column is used"
+            assert not sync_computed_embeddings, "Sync computed embedding is not supported with embedding vector column"
+            json_data["delta_sync_index_spec"]["embedding_vector_columns"] = [
+                {
+                    "name": embedding_vector_column,
+                    "embedding_dimension": embedding_dimension
+                }
+            ]
+        elif embedding_source_column:
+            assert embedding_model_endpoint_name, \
+                "You must specify Embedding Model Endpoint"
+            json_data["delta_sync_index_spec"]["embedding_source_columns"] = [
+                {
+                    "name": embedding_source_column,
+                    "embedding_model_endpoint_name": embedding_model_endpoint_name
+                }
+            ]
+            if sync_computed_embeddings:
+                json_data["delta_sync_index_spec"]["embedding_writeback_table"] = f'{index_name}_writeback_table'
+
+        resp = RequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/endpoints/{endpoint_name}/indexes",
+            token=self._get_token_for_request(),
+            method="POST",
+            json=json_data,
+        )
+
+        index_url = resp.get('status', {}).get('index_url')
+        return VectorSearchIndex(
+            workspace_url=self.workspace_url,
+            index_url=index_url,
+            personal_access_token=self.personal_access_token,
+            service_principal_client_id=self.service_principal_client_id,
+            service_principal_client_secret=self.service_principal_client_secret,
+            name=resp["name"],
+            endpoint_name=endpoint_name,
+            azure_tenant_id=self.azure_tenant_id,
+            azure_login_id=self.azure_login_id,
+            use_user_passed_credentials=self._using_user_passed_credentials
+        )
+    
+    def create_delta_sync_index_and_wait(
+            self,
+            endpoint_name,
+            index_name,
+            primary_key,
+            source_table_name,
+            pipeline_type,
+            embedding_dimension=None,
+            embedding_vector_column=None,
+            embedding_source_column=None,
+            embedding_model_endpoint_name=None,
+            sync_computed_embeddings=False,
+            verbose=False,
+            timeout=datetime.timedelta(hours=24)):
+        """
+        Create a delta sync index and wait for it to be ready.
+
+        :param str endpoint_name: The name of the endpoint.
+        :param str index_name: The name of the index.
+        :param str primary_key: The primary key of the index.
+        :param str source_table_name: The name of the source table.
+        :param str pipeline_type: The type of the pipeline. Must be CONTINUOUS or TRIGGERED.
+        :param int embedding_dimension: The dimension of the embedding vector.
+        :param str embedding_vector_column: The name of the embedding vector column.
+        :param str embedding_source_column: The name of the embedding source column.
+        :param str embedding_model_endpoint_name: The name of the embedding model endpoint.
+        :param bool verbose: Whether to print status messages.
+        :param datetime.timedelta timeout: The time allowed until we timeout with an Exception.
+        :param bool sync_computed_embeddings: Whether to automatically sync the vector index contents and computed embeddings to a new UC table,
+                                             table name will be ${index_name}_writeback_table.
+        """
+        index = self.create_delta_sync_index(
+            endpoint_name,
+            index_name,
+            primary_key,
+            source_table_name,
+            pipeline_type,
+            embedding_dimension,
+            embedding_vector_column,
+            embedding_source_column,
+            embedding_model_endpoint_name,
+            sync_computed_embeddings)
+        index.wait_until_ready(verbose, timeout)
+        return index
+    
+
+    def create_direct_access_index(
+            self,
+            endpoint_name,
+            index_name,
+            primary_key,
+            embedding_dimension,
+            embedding_vector_column,
+            schema,
+            embedding_model_endpoint_name=None):
+        """
+        Create a direct access index.
+
+        :param str endpoint_name: The name of the endpoint.
+        :param str index_name: The name of the index.
+        :param str primary_key: The primary key of the index.
+        :param int embedding_dimension: The dimension of the embedding vector.
+        :param str embedding_vector_column: The name of the embedding vector column.
+        :param dict schema: The schema of the index.
+        :param str embedding_model_endpoint_name: The name of the optional embedding model endpoint to use when querying.
+        """
+        assert schema, """
+            Schema must be present when creating a direct access index.
+            Example schema: {"id": "integer", "text": "string", \
+                "text_vector": "array<float>", "bool_val": "boolean", \
+                    "float_val": "float", "date_val": "date"}"
+        """
+        json_data = {
+            "name": index_name,
+            "index_type": "DIRECT_ACCESS",
+            "primary_key": primary_key,
+            "direct_access_index_spec": {
+                "embedding_vector_columns": [
+                    {
+                        "name": embedding_vector_column,
+                        "embedding_dimension": embedding_dimension
+                    }
+                ],
+                "schema_json": json.dumps(schema)
+            },
+        }
+        if embedding_model_endpoint_name:
+            json_data["direct_access_index_spec"]["embedding_source_columns"] = [
+                {
+                    "embedding_model_endpoint_name": embedding_model_endpoint_name
+                }
+            ]
+        resp = RequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/endpoints/{endpoint_name}/indexes",
+            token=self._get_token_for_request(),
+            method="POST",
+            json=json_data,
+        )
+
+        index_url = resp.get('status', {}).get('index_url')
+        return VectorSearchIndex(
+            workspace_url=self.workspace_url,
+            index_url=index_url,
+            personal_access_token=self.personal_access_token,
+            service_principal_client_id=self.service_principal_client_id,
+            service_principal_client_secret=self.service_principal_client_secret,
+            name=resp["name"],
+            endpoint_name=endpoint_name,
+            azure_tenant_id=self.azure_tenant_id,
+            azure_login_id=self.azure_login_id,
+            use_user_passed_credentials=self._using_user_passed_credentials
+        )
+    
+    def _get_index_url(self, endpoint_name, index_name):
+        if endpoint_name:
+            url = f"{self.workspace_url}/api/2.0/vector-search/endpoints/{endpoint_name}/indexes/{index_name}"
+        else:
+            url = f"{self.workspace_url}/api/2.0/vector-search/indexes/{index_name}"
+        return url
+
+
+    def get_index(self, endpoint_name=None, index_name=None):
+        """
+        Get an index.
+
+        :param Option[str] endpoint_name: The optional name of the endpoint.
+        :param str index_name: The name of the index.
+        """
+        assert index_name, "Index name must be specified"
+        resp = RequestUtils.issue_request(
+            url=self._get_index_url(endpoint_name, index_name),
+            token=self._get_token_for_request(),
+            method="GET",
+        )
+        index_url = resp.get('status', {}).get('index_url')
+        response_endpoint_name = resp.get('endpoint_name')
+        return VectorSearchIndex(
+            workspace_url=self.workspace_url,
+            index_url=index_url,
+            personal_access_token=self.personal_access_token,
+            service_principal_client_id=self.service_principal_client_id,
+            service_principal_client_secret=self.service_principal_client_secret,
+            name=index_name,
+            endpoint_name=response_endpoint_name,
+            azure_tenant_id=self.azure_tenant_id,
+            azure_login_id=self.azure_login_id,
+            use_user_passed_credentials=self._using_user_passed_credentials
+        )
+
+    def delete_index(self, endpoint_name=None, index_name=None):
+        """
+        Delete an index.
+
+        :param Option[str] endpoint_name: The optional name of the endpoint.
+        :param str index_name: The name of the index.
+        """
+        assert index_name, "Index name must be specified"
+        return RequestUtils.issue_request(
+            url=self._get_index_url(endpoint_name, index_name),
+            token=self._get_token_for_request(),
+            method="DELETE",
+        )

databricks_vectorsearch-0.37/databricks/vector_search/exceptions.py

@@ -0,0 +1,3 @@
+
+class InvalidInputException(Exception):
+    pass

databricks_vectorsearch-0.37/databricks/vector_search/index.py

@@ -0,0 +1,319 @@
+import logging
+import json
+import time
+import datetime
+import math
+import deprecation
+from databricks.vector_search.utils import OAuthTokenUtils
+from databricks.vector_search.utils import RequestUtils
+from databricks.vector_search.utils import UrlUtils
+from mlflow.utils import databricks_utils
+
+
+class VectorSearchIndex:
+    """
+    VectorSearchIndex is a helper class that represents a Vector Search Index.
+
+    Those who wish to use this class should not instantiate it directly, but rather use the VectorSearchClient class.
+    """
+    def __init__(
+        self,
+        workspace_url,
+        index_url,
+        name,
+        endpoint_name,
+        personal_access_token=None,
+        service_principal_client_id=None,
+        service_principal_client_secret=None,
+        azure_tenant_id=None,
+        azure_login_id=None,
+        # whether or not credentials were explicitly passed in by user in client or inferred by client
+        # via mlflow utilities. If passed in by user, continue to use user credentials. If not, can
+        # attempt automatic auth refresh for model serving.
+        use_user_passed_credentials=False
+    ):
+        self.workspace_url = workspace_url
+        self.index_url = UrlUtils.add_https_if_missing(index_url) \
+            if index_url else None
+        self.name = name
+        self.endpoint_name = endpoint_name
+        self.personal_access_token = personal_access_token
+        self.service_principal_client_id = service_principal_client_id
+        self.service_principal_client_secret = service_principal_client_secret
+        if self.personal_access_token and \
+                not (self.service_principal_client_id and 
+                     self.service_principal_client_secret):
+            # In PAT flow, don't use index_url given for DP ingress
+            self.index_url = self.workspace_url + f"/api/2.0/vector-search/endpoints/{self.endpoint_name}/indexes/{self.name}"
+        self.index_url = self.index_url or (self.workspace_url + f"/api/2.0/vector-search/endpoints/{self.endpoint_name}/indexes/{self.name}") # Fallback to CP
+        self.azure_tenant_id = azure_tenant_id
+        self.azure_login_id = azure_login_id
+        self._control_plane_oauth_token = None
+        self._control_plane_oauth_token_expiry_ts = None
+        self._read_oauth_token = None
+        self._read_oauth_token_expiry_ts = None
+        self._write_oauth_token = None
+        self._write_oauth_token_expiry_ts = None
+        self._use_user_passed_credentials = use_user_passed_credentials
+
+    def _get_token_for_request(self, write=False, control_plane=False):
+        try:
+            # automatically refresh auth if not passed in by user and in model serving environment
+            if not self._use_user_passed_credentials and databricks_utils.is_in_databricks_model_serving_environment():
+                return databricks_utils.get_databricks_host_creds().token
+        except Exception as e:
+            logging.warning(f"Reading credentials from model serving environment failed with: {e} "
+                    f"Defaulting to cached vector search token")
+
+        if self.personal_access_token:  # PAT flow
+            return self.personal_access_token
+        if self.workspace_url in self.index_url:
+            control_plane = True
+        if (
+            control_plane and
+            self._control_plane_oauth_token and
+            self._control_plane_oauth_token_expiry_ts and
+            self._control_plane_oauth_token_expiry_ts - 100 > time.time()
+        ):
+            return self._control_plane_oauth_token
+        if (
+            write and
+            not control_plane and
+            self._write_oauth_token
+            and self._write_oauth_token_expiry_ts
+            and self._write_oauth_token_expiry_ts - 100 > time.time()
+        ):
+            return self._write_oauth_token
+        if (
+            not write and
+            not control_plane and
+            self._read_oauth_token
+            and self._read_oauth_token_expiry_ts
+            and self._read_oauth_token_expiry_ts - 100 > time.time()
+        ):
+            return self._read_oauth_token
+        if self.service_principal_client_id and \
+                self.service_principal_client_secret:
+            authorization_details = json.dumps([{
+                "type": "unity_catalog_permission",
+                "securable_type": "table",
+                "securable_object_name": self.name,
+                "operation": "WriteVectorIndex" if write else "ReadVectorIndex"
+            }]) if not control_plane else []
+            oauth_token_data = OAuthTokenUtils.get_oauth_token(
+                workspace_url=self.workspace_url,
+                service_principal_client_id=self.service_principal_client_id,
+                service_principal_client_secret=self.service_principal_client_secret,
+                authorization_details=authorization_details
+            ) if not self.azure_tenant_id else OAuthTokenUtils.get_azure_oauth_token(
+                workspace_url=self.workspace_url,
+                service_principal_client_id=self.service_principal_client_id,
+                service_principal_client_secret=self.service_principal_client_secret,
+                authorization_details=authorization_details,
+                azure_tenant_id=self.azure_tenant_id,
+                azure_login_id=self.azure_login_id
+            )
+            if control_plane:
+                self._control_plane_oauth_token = oauth_token_data["access_token"]
+                self._control_plane_oauth_token_expiry_ts = time.time() + oauth_token_data["expires_in"]
+                return self._control_plane_oauth_token
+            if write:
+                self._write_oauth_token = oauth_token_data["access_token"]
+                self._write_oauth_token_expiry_ts = time.time() + oauth_token_data["expires_in"]
+                return self._write_oauth_token
+            self._read_oauth_token = oauth_token_data["access_token"]
+            self._read_oauth_token_expiry_ts = time.time() + oauth_token_data["expires_in"]
+            return self._read_oauth_token
+        raise Exception("You must specify service principal or PAT token")
+
+    def upsert(self, inputs):
+        """
+        Upsert data into the index.
+
+        :param inputs: List of dictionaries to upsert into the index.
+        """
+        assert type(inputs) == list, "inputs must be of type: List of dictionaries"
+        assert all(
+            type(i) == dict for i in inputs
+        ), "inputs must be of type: List of dicts"
+        upsert_payload = {"inputs_json": json.dumps(inputs)}
+        return RequestUtils.issue_request(
+            url=f"{self.index_url}/upsert-data",
+            token=self._get_token_for_request(write=True),
+            method="POST",
+            json=upsert_payload
+        )
+
+    def delete(self, primary_keys):
+        """
+        Delete data from the index.
+
+        :param primary_keys: List of primary keys to delete from the index.
+        """
+        assert type(primary_keys) == list, "inputs must be of type: List"
+        delete_payload = {"primary_keys": primary_keys}
+        return RequestUtils.issue_request(
+            url=f"{self.index_url}/delete-data",
+            token=self._get_token_for_request(write=True),
+            method="DELETE",
+            json=delete_payload
+        )
+
+    def describe(self):
+        """
+        Describe the index. This returns metadata about the index.
+        """
+        return RequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/endpoints/{self.endpoint_name}/indexes/{self.name}",
+            token=self._get_token_for_request(control_plane=True),
+            method="GET",
+        )
+
+    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".
+        """
+        return RequestUtils.issue_request(
+            url=f"{self.workspace_url}/api/2.0/vector-search/endpoints/{self.endpoint_name}/indexes/{self.name}/sync",
+            token=self._get_token_for_request(control_plane=True),
+            method="POST",
+        )
+
+    def similarity_search(
+        self,
+        columns,
+        query_text=None,
+        query_vector=None,
+        filters=None,
+        num_results=5,
+        debug_level=1,
+        score_threshold=None,
+        query_method=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.
+        :param query_method: Query method to use for the query. Choices are "ANN" and "HYBRID".
+
+        """
+        json_data = {
+            "num_results": num_results,
+            "columns": columns,
+            "filters_json": json.dumps(filters) if filters else None,
+            "debug_level": debug_level
+        }
+        if query_text:
+            json_data["query"] = query_text
+            json_data["query_text"] = query_text
+        if query_vector:
+            json_data["query_vector"] = query_vector
+        if score_threshold:
+            json_data["score_threshold"] = score_threshold
+        if query_method:
+            json_data["query_method"] = query_method
+
+        response = RequestUtils.issue_request(
+            url=f"{self.index_url}/query",
+            token=self._get_token_for_request(),
+            method="GET",
+            json=json_data
+        )
+
+        out_put = response
+        while response["next_page_token"]:
+           response = self.__get_next_page(response["next_page_token"])
+           out_put["result"]["row_count"] += response["result"]["row_count"]
+           out_put["result"]["data_array"] += response["result"]["data_array"]
+
+        out_put.pop("next_page_token", None)
+        return out_put
+
+    def wait_until_ready(self, verbose=False, timeout=datetime.timedelta(hours=24)):
+        """
+        Wait for the index to be online.
+
+        :param bool verbose: Whether to print status messages.
+        :param datetime.timedelta timeout: The time allowed until we timeout with an Exception.
+        """
+                         
+        def get_index_state():
+            return self.describe()["status"]["detailed_state"]
+
+        start_time = datetime.datetime.now()
+        sleep_time_seconds = 30
+        # Provisioning states all contain `PROVISIONING`
+        # Online states all contain `ONLINE`.
+        # Offline states all contain `OFFLINE`.
+        index_state = get_index_state()
+        while "ONLINE" not in index_state and datetime.datetime.now() - start_time < timeout:
+            if "OFFLINE" in index_state:
+                raise Exception(f"Index {self.name} is offline")
+            if verbose:
+                running_time = int(math.floor((datetime.datetime.now() - start_time).total_seconds()))
+                print(f"Index {self.name} is in state {index_state}. Time: {running_time}s.")
+            time.sleep(sleep_time_seconds)
+            index_state = get_index_state()
+        if verbose:
+            print(f"Index {self.name} is in state {index_state}.")
+        if "ONLINE" not in index_state:
+            raise Exception(f"Index {self.name} did not become online within timeout of {timeout.total_seconds()}s.")
+
+    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, it will be used as the exclusive starting primary key.
+        """
+        json_data = {
+            "num_results": num_results,
+            "endpoint_name": self.endpoint_name,
+        }
+        if last_primary_key:
+            json_data["last_primary_key"] = last_primary_key
+
+        # TODO(ShengZhan): make this consistent with the rest.
+        url =  f"{self.workspace_url}/api/2.0/vector-search/indexes/{self.name}/scan"
+
+        return RequestUtils.issue_request(
+            url=url,
+            token=self._get_token_for_request(),
+            method="GET",
+            json=json_data
+       )
+
+    @deprecation.deprecated(deprecated_in="0.36", removed_in="0.37",
+                            current_version="0.36",
+                            details="Use the scan function instead")
+    def scan_index(self, num_results = 10, last_primary_key=None):
+        return self.scan(num_results, last_primary_key)
+
+    def __get_next_page(self, page_token):
+        """
+        Get the next page of results from a page token.
+        """
+        json_data = {
+            "page_token": page_token,
+            "endpoint_name": self.endpoint_name,
+        }
+        # TODO(ShengZhan): make this consistent with the rest.
+        url =  f"{self.workspace_url}/api/2.0/vector-search/indexes/{self.name}/query-next-page"
+
+        return RequestUtils.issue_request(
+            url=url,
+            token=self._get_token_for_request(),
+            method="GET",
+            json=json_data
+        )

databricks_vectorsearch-0.37/databricks/vector_search/utils.py

@@ -0,0 +1,139 @@
+import requests
+import logging
+import os
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
+from functools import lru_cache
+from databricks.vector_search.version import VERSION
+
+class OAuthTokenUtils:
+
+    @staticmethod
+    def get_azure_oauth_token(
+        workspace_url,
+        azure_tenant_id,
+        azure_login_id,
+        service_principal_client_id,
+        service_principal_client_secret,
+        authorization_details=None,
+    ):
+        url = f"https://login.microsoftonline.com/{azure_tenant_id}/oauth2/v2.0/token"
+        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
+        resource_identifier = azure_login_id
+        assert (azure_login_id and azure_tenant_id), "Both azure_login_id and azure_tenant_id must be specified"
+        data = {
+            "grant_type": "client_credentials",
+            "scope": f"{resource_identifier}/.default",
+            "client_id": service_principal_client_id,
+            "client_secret": service_principal_client_secret
+        }
+        azure_response = RequestUtils.issue_request(
+            url=url,
+            headers=headers,
+            method="POST",
+            data=data
+        )
+        aad_token = azure_response['access_token']
+        authorization_details = authorization_details or []
+        if not authorization_details:
+            return azure_response
+        url = workspace_url + "/oidc/v1/token"
+        headers = {
+            'Content-Type': 'application/x-www-form-urlencoded',
+            'Accept': '*/*',
+        }
+        data = {
+            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+            "assertion": aad_token,
+            "authorization_details": authorization_details
+        }
+        return RequestUtils.issue_request(
+            url=url,
+            headers=headers,
+            method="POST",
+            data=data
+        )
+
+    @staticmethod
+    def get_oauth_token(
+        workspace_url,
+        service_principal_client_id,
+        service_principal_client_secret,
+        authorization_details=None,
+    ):
+        authorization_details = authorization_details or []
+        url = workspace_url + "/oidc/v1/token"
+        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
+        data = {
+            "grant_type": "client_credentials",
+            "scope": "all-apis",
+            "authorization_details": authorization_details
+        }
+        logging.info(f"Issuing request to {url} with data {data} and headers {headers}")
+        response = RequestUtils.issue_request(
+            url=url,
+            auth=(service_principal_client_id, service_principal_client_secret),
+            headers=headers,
+            method="POST",
+            data=data
+        )
+        return response
+
+@lru_cache(maxsize=64)
+def _cached_get_request_session(
+        total_retries,
+        backoff_factor,
+        # To create a new Session object for each process, we use the process id as the cache key.
+        # This is to avoid sharing the same Session object across processes, which can lead to issues
+        # such as https://stackoverflow.com/q/3724900.
+        process_id):
+    session = requests.Session()
+    retry_strategy = Retry(
+        total=total_retries,  # Total number of retries
+        backoff_factor=backoff_factor,  # A backoff factor to apply between attempts
+        status_forcelist=[429],  # HTTP status codes to retry on
+    )
+    adapter = HTTPAdapter(max_retries=retry_strategy, pool_connections=50, pool_maxsize=50)
+    session.mount("https://", adapter)
+    session.mount("http://", adapter)
+    return session
+
+class RequestUtils:
+    session = _cached_get_request_session(
+        total_retries=3,
+        backoff_factor=1,
+        process_id=os.getpid())
+
+    @staticmethod
+    def issue_request(url, method, token=None, params=None, json=None, verify=True, auth=None, data=None, headers=None):
+        headers = headers or dict()
+        if token:
+            headers["Authorization"] = f"Bearer {token}"
+        headers["X-Databricks-Python-SDK-Version"] = VERSION
+        response = RequestUtils.session.request(
+            url=url,
+            headers=headers,
+            method=method,
+            params=params,
+            json=json,
+            verify=verify,
+            auth=auth,
+            data=data
+        )
+        try:
+            response.raise_for_status()
+        except Exception as e:
+            logging.warn(f"Error processing request {e}")
+            raise Exception(
+                f"Response content {response.content}, status_code {response.status_code}"
+            )
+        return response.json()
+
+
+class UrlUtils:
+
+    @staticmethod
+    def add_https_if_missing(url):
+        if not url.startswith("http://") and not url.startswith("https://"):
+            url = "https://" + url
+        return url

databricks_vectorsearch-0.37/databricks/vector_search/version.py

@@ -0,0 +1,2 @@
+# Update this when publishing a new version
+VERSION = "0.37"

databricks_vectorsearch-0.37/databricks_vectorsearch.egg-info/PKG-INFO

@@ -0,0 +1,45 @@
+Metadata-Version: 2.1
+Name: databricks-vectorsearch
+Version: 0.37
+Summary: Databricks Vector Search Client
+Home-page: UNKNOWN
+Author: Databricks
+Author-email: feedback@databricks.com
+License: UNKNOWN
+Platform: UNKNOWN
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+
+**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 Downlodable 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.
+
+
+

databricks_vectorsearch-0.37/databricks_vectorsearch.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+README.md
+setup.py
+databricks/__init__.py
+databricks/vector_search/__init__.py
+databricks/vector_search/client.py
+databricks/vector_search/exceptions.py
+databricks/vector_search/index.py
+databricks/vector_search/utils.py
+databricks/vector_search/version.py
+databricks_vectorsearch.egg-info/PKG-INFO
+databricks_vectorsearch.egg-info/SOURCES.txt
+databricks_vectorsearch.egg-info/dependency_links.txt
+databricks_vectorsearch.egg-info/requires.txt
+databricks_vectorsearch.egg-info/top_level.txt

databricks_vectorsearch-0.37/databricks_vectorsearch.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

databricks_vectorsearch-0.37/databricks_vectorsearch.egg-info/requires.txt

@@ -0,0 +1,4 @@
+mlflow-skinny<3,>=2.11.3
+protobuf<5,>=3.12.0
+requests>=2
+deprecation>=2

databricks_vectorsearch-0.37/databricks_vectorsearch.egg-info/top_level.txt

@@ -0,0 +1 @@
+databricks

databricks_vectorsearch-0.37/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

databricks_vectorsearch-0.37/setup.py

@@ -0,0 +1,28 @@
+import os
+from setuptools import find_packages, setup
+from importlib.machinery import SourceFileLoader
+
+with open("requirements.txt") as f:
+    required = f.read().splitlines()
+
+version = (
+    SourceFileLoader(
+        "version", os.path.join("databricks", "vector_search", "version.py")
+    )
+    .load_module()
+    .VERSION
+)
+
+setup(
+    name="databricks_vectorsearch",
+    version=version,
+    packages=find_packages(),
+    author="Databricks",
+    author_email="feedback@databricks.com",
+    license_files=("LICENSE.md"),
+    description="Databricks Vector Search Client",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    install_requires=required,
+    python_requires=">=3.7",
+)