Migrate from databricks-vectorsearch to databricks-ai-search

integrations/dspy/src/databricks_dspy/retrievers/databricks_rm.py

@@ -39,7 +39,7 @@
 
         .. code-block:: python
 
-            from databricks.vector_search.client import VectorSearchClient
+            from databricks.ai_search.client import VectorSearchClient
             from databricks.sdk import WorkspaceClient
 
             # Create a Databricks workspace client
@@ -151,7 +151,7 @@
             try:
                 import mlflow
 
                 mlflow.models.set_retriever_schema(  # ty:ignore[possibly-missing-attribute]: caught below
                     primary_key="doc_id",
                     text_column="page_content",
                     doc_uri="doc_uri",

integrations/langchain/pyproject.toml

@@ -10,7 +10,7 @@ license = { text="Apache-2.0" }
 requires-python = ">=3.10"
 dependencies = [
     "langchain>=1.0.0",
-    "databricks-vectorsearch>=0.50",
+    "databricks-ai-search>=0.73",
     "databricks-ai-bridge>=0.19.0",
     "mlflow>=3.0.0",
     "pydantic>2.10.0",

integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py

@@ -1,7 +1,7 @@
 import json
 from typing import Type
 
-from databricks.vector_search.reranker import DatabricksReranker
+from databricks.ai_search.reranker import DatabricksReranker
 from databricks_ai_bridge.utils.vector_search import IndexDetails
 from databricks_ai_bridge.vector_search_retriever_tool import (
     FilterItem,
@@ -24,7 +24,7 @@ class VectorSearchRetrieverTool(BaseTool, VectorSearchRetrieverToolMixin):
     for building a retriever tool for agents.
 
     **Note**: Any additional keyword arguments passed to the constructor will be passed along to
-    `databricks.vector_search.client.VectorSearchIndex.similarity_search` when executing the tool. `See
+    `databricks.ai_search.index.VectorSearchIndex.similarity_search` when executing the tool. `See
     documentation <https://api-docs.databricks.com/python/vector-search/databricks.vector_search.html#databricks.vector_search.index.VectorSearchIndex.similarity_search>`_
     to see the full set of supported keyword arguments,
     e.g. `score_threshold`. Also, see documentation for

integrations/langchain/src/databricks_langchain/vectorstores.py

@@ -17,8 +17,8 @@
 )
 
 import numpy as np
+from databricks.ai_search.reranker import DatabricksReranker, Reranker
 from databricks.sdk import WorkspaceClient
-from databricks.vector_search.reranker import DatabricksReranker, Reranker
 from databricks_ai_bridge.utils.vector_search import (
     IndexDetails,
     RetrieverSchema,
@@ -67,7 +67,7 @@ class DatabricksVectorSearch(VectorStore):
                     and ``service_principal_client_secret`` to allow for
                     service principal authentication instead of personal access token authentication.
         reranker: Optional reranker to apply on the top results. Pass an instance of
-                    ``databricks.vector_search.reranker.DatabricksReranker`` with
+                    ``databricks.ai_search.reranker.DatabricksReranker`` with
                     ``columns_to_rerank=[...]``. The reranker reorders the initial results using
                     the specified text columns.
 
@@ -107,7 +107,7 @@ class DatabricksVectorSearch(VectorStore):
 
         .. code-block:: python
 
-            from databricks.vector_search.reranker import DatabricksReranker
+            from databricks.ai_search.reranker import DatabricksReranker
 
             vector_store = DatabricksVectorSearch(
                 index_name="<your-index-name>",
@@ -249,12 +249,12 @@ def __init__(
             )
 
         try:
-            from databricks.vector_search.client import VectorSearchClient
-            from databricks.vector_search.utils import CredentialStrategy
+            from databricks.ai_search.client import VectorSearchClient
+            from databricks.ai_search.utils import CredentialStrategy
         except ImportError as e:
             raise ImportError(
-                "Could not import databricks-vectorsearch python package. "
-                "Please install it with `pip install databricks-vectorsearch`."
+                "Could not import databricks-ai-search python package. "
+                "Please install it with `pip install databricks-ai-search`."
             ) from e
 
         if endpoint is not None:
@@ -422,7 +422,7 @@ def similarity_search(
             filter: Filters to apply to the query. Defaults to None.
             query_type: The type of this query. Supported values are "ANN" and "HYBRID".
             reranker: Allows reranking the results. Defaults to None.
-            kwargs: Additional keyword arguments to pass to `databricks.vector_search.client.VectorSearchIndex.similarity_search`. `See
+            kwargs: Additional keyword arguments to pass to `databricks.ai_search.index.VectorSearchIndex.similarity_search`. `See
                     documentation <https://api-docs.databricks.com/python/vector-search/databricks.vector_search.html#databricks.vector_search.index.VectorSearchIndex.similarity_search>`_
                     to see the full set of supported keyword arguments
 
@@ -464,7 +464,7 @@ def similarity_search_with_score(
             filter: Filters to apply to the query. Defaults to None.
             query_type: The type of this query. Supported values are "ANN" and "HYBRID".
             reranker: Allows reranking the results. Defaults to None.
-            kwargs: Additional keyword arguments to pass to `databricks.vector_search.client.VectorSearchIndex.similarity_search`. `See
+            kwargs: Additional keyword arguments to pass to `databricks.ai_search.index.VectorSearchIndex.similarity_search`. `See
                     documentation <https://api-docs.databricks.com/python/vector-search/databricks.vector_search.html#databricks.vector_search.index.VectorSearchIndex.similarity_search>`_
                     to see the full set of supported keyword arguments
 
@@ -538,7 +538,7 @@ def similarity_search_by_vector(
             filter: Filters to apply to the query. Defaults to None.
             query_type: The type of this query. Supported values are "ANN" and "HYBRID".
             reranker: Allows reranking the results. Defaults to None.
-            kwargs: Additional keyword arguments to pass to `databricks.vector_search.client.VectorSearchIndex.similarity_search`. `See
+            kwargs: Additional keyword arguments to pass to `databricks.ai_search.index.VectorSearchIndex.similarity_search`. `See
                     documentation <https://api-docs.databricks.com/python/vector-search/databricks.vector_search.html#databricks.vector_search.index.VectorSearchIndex.similarity_search>`_
                     to see the full set of supported keyword arguments
 
@@ -591,7 +591,7 @@ def similarity_search_by_vector_with_score(
             filter: Filters to apply to the query. Defaults to None.
             query_type: The type of this query. Supported values are "ANN" and "HYBRID".
             reranker: Allows reranking the results. Defaults to None.
-            kwargs: Additional keyword arguments to pass to `databricks.vector_search.client.VectorSearchIndex.similarity_search`. `See
+            kwargs: Additional keyword arguments to pass to `databricks.ai_search.index.VectorSearchIndex.similarity_search`. `See
                     documentation <https://api-docs.databricks.com/python/vector-search/databricks.vector_search.html#databricks.vector_search.index.VectorSearchIndex.similarity_search>`_
                     to see the full set of supported keyword arguments
 

integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py

@@ -6,9 +6,9 @@
 
 import mlflow
 import pytest
+from databricks.ai_search.utils import CredentialStrategy
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.credentials_provider import ModelServingUserCredentials
-from databricks.vector_search.utils import CredentialStrategy
 from databricks_ai_bridge.test_utils.vector_search import (  # noqa: F401
     ALL_INDEX_NAMES,
     DELTA_SYNC_INDEX,
@@ -301,7 +301,7 @@ def test_vector_search_client_model_serving_environment():
             host="testDogfod.com", credentials_strategy=ModelServingUserCredentials()
         )
 
-        with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+        with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
             mock_instance = mockVSClient.return_value
             mock_instance.get_index.side_effect = _get_index
             with patch("databricks.sdk.service.serving.ServingEndpointsAPI.get", return_value=None):
@@ -317,7 +317,7 @@ def test_vector_search_client_model_serving_environment():
 
 
 def test_vector_search_client_non_model_serving_environment():
-    with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+    with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
         mock_instance = mockVSClient.return_value
         mock_instance.get_index.side_effect = _get_index
         vsTool = VectorSearchRetrieverTool(
@@ -332,7 +332,7 @@ def test_vector_search_client_with_pat_workspace_client():
     w.config.auth_type = "pat"
     w.config.host = "https://testDogfod.com"
     w.config.token = "fakeToken"
-    with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+    with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
         with patch("databricks.sdk.service.serving.ServingEndpointsAPI.get", return_value=None):
             mock_instance = mockVSClient.return_value
             mock_instance.get_index.side_effect = _get_index
@@ -356,7 +356,7 @@ def test_vector_search_client_with_sp_workspace_client():
     w.config.client_id = "fakeClientId"
     w.config.client_secret = "fakeClientSecret"
 
-    with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+    with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
         with patch("databricks.sdk.service.serving.ServingEndpointsAPI.get", return_value=None):
             mock_instance = mockVSClient.return_value
             mock_instance.get_index.side_effect = _get_index

integrations/langchain/tests/unit_tests/test_vectorstores.py

@@ -4,8 +4,8 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from databricks.vector_search.index import VectorSearchIndex
-from databricks.vector_search.reranker import DatabricksReranker, Reranker
+from databricks.ai_search.index import VectorSearchIndex
+from databricks.ai_search.reranker import DatabricksReranker, Reranker
 from databricks_ai_bridge.test_utils.vector_search import (
     ALL_INDEX_NAMES,
     DELTA_SYNC_INDEX,

integrations/llamaindex/pyproject.toml

@@ -9,7 +9,7 @@ readme = "README.md"
 license = { text="Apache-2.0" }
 requires-python = ">=3.10"
 dependencies = [
-    "databricks-vectorsearch>=0.40",
+    "databricks-ai-search>=0.73",
     "databricks-ai-bridge>=0.1.0",
     "llama-index>=0.12.0",
     "unitycatalog-llamaindex[databricks]>=0.2.0",

integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py

@@ -47,8 +47,8 @@ def __init__(self, **data):
         VectorSearchRetrieverToolMixin.__init__(self, **data)
 
         # Initialize private attributes
-        from databricks.vector_search.client import VectorSearchClient
-        from databricks.vector_search.utils import CredentialStrategy
+        from databricks.ai_search.client import VectorSearchClient
+        from databricks.ai_search.utils import CredentialStrategy
 
         credential_strategy = None
         if (

integrations/llamaindex/tests/unit_tests/test_vector_search_retriever_tool.py

@@ -4,11 +4,11 @@
 from unittest.mock import MagicMock, create_autospec, patch
 
 import pytest
+from databricks.ai_search.index import VectorSearchIndex
+from databricks.ai_search.reranker import DatabricksReranker, Reranker
+from databricks.ai_search.utils import CredentialStrategy
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.credentials_provider import ModelServingUserCredentials
-from databricks.vector_search.index import VectorSearchIndex
-from databricks.vector_search.reranker import DatabricksReranker, Reranker
-from databricks.vector_search.utils import CredentialStrategy
 from databricks_ai_bridge.test_utils.vector_search import (  # noqa: F401
     ALL_INDEX_NAMES,
     DEFAULT_VECTOR_DIMENSION,
@@ -156,7 +156,7 @@ def test_vector_search_client_model_serving_environment():
             host="testDogfod.com", credentials_strategy=ModelServingUserCredentials()
         )
 
-        with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+        with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
             with patch("databricks.sdk.service.serving.ServingEndpointsAPI.get", return_value=None):
                 vsTool = VectorSearchRetrieverTool(
                     index_name="catalog.schema.my_index_name",
@@ -172,7 +172,7 @@ def test_vector_search_client_model_serving_environment():
 
 
 def test_vector_search_client_non_model_serving_environment():
-    with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+    with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
         vsTool = VectorSearchRetrieverTool(
             index_name="catalog.schema.my_index_name",
             text_column="abc",
@@ -182,7 +182,7 @@ def test_vector_search_client_non_model_serving_environment():
         mockVSClient.assert_called_once_with(disable_notice=True, credential_strategy=None)
 
     w = WorkspaceClient(host="testDogfod.com", token="fakeToken")
-    with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+    with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
         with patch("databricks.sdk.service.serving.ServingEndpointsAPI.get", return_value=None):
             vsTool = VectorSearchRetrieverTool(
                 index_name="catalog.schema.my_index_name",

integrations/openai/pyproject.toml

@@ -9,7 +9,7 @@ readme = "README.md"
 license = { text="Apache-2.0" }
 requires-python = ">=3.10"
 dependencies = [
-    "databricks-vectorsearch>=0.50",
+    "databricks-ai-search>=0.73",
     "databricks-ai-bridge>=0.19.0",
     "openai>=1.99.9",
     "pydantic>2.10.0",

integrations/openai/src/databricks_openai/vector_search_retriever_tool.py

@@ -2,7 +2,7 @@
 import logging
 from typing import Any, Dict, List, Optional, Tuple
 
-from databricks.vector_search.index import VectorSearchIndex
+from databricks.ai_search.index import VectorSearchIndex
 from databricks_ai_bridge.utils.vector_search import (
     IndexDetails,
     RetrieverSchema,
@@ -68,7 +68,7 @@ class VectorSearchRetrieverTool(VectorSearchRetrieverToolMixin):
             )
 
     **Note**: Any additional keyword arguments passed to the constructor will be passed along to
-    `databricks.vector_search.client.VectorSearchIndex.similarity_search` when executing the tool. `See
+    `databricks.ai_search.index.VectorSearchIndex.similarity_search` when executing the tool. `See
     documentation <https://api-docs.databricks.com/python/vector-search/databricks.vector_search.html#databricks.vector_search.index.VectorSearchIndex.similarity_search>`_
     to see the full set of supported keyword arguments,
     e.g. `score_threshold`. Also, see documentation for
@@ -99,10 +99,10 @@ class VectorSearchRetrieverTool(VectorSearchRetrieverToolMixin):
 
     @model_validator(mode="after")
     def _validate_tool_inputs(self):
-        from databricks.vector_search.client import (
+        from databricks.ai_search.client import (
             VectorSearchClient,  # import here so we can mock in tests
         )
-        from databricks.vector_search.utils import CredentialStrategy
+        from databricks.ai_search.utils import CredentialStrategy
 
         splits = self.index_name.split(".")
         if len(splits) != 3:

integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py

@@ -6,11 +6,11 @@
 
 import mlflow
 import pytest
+from databricks.ai_search.index import VectorSearchIndex
+from databricks.ai_search.reranker import DatabricksReranker, Reranker
+from databricks.ai_search.utils import CredentialStrategy
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.credentials_provider import ModelServingUserCredentials
-from databricks.vector_search.index import VectorSearchIndex
-from databricks.vector_search.reranker import DatabricksReranker, Reranker
-from databricks.vector_search.utils import CredentialStrategy
 from databricks_ai_bridge.test_utils.vector_search import (  # noqa: F401
     ALL_INDEX_NAMES,
     DELTA_SYNC_INDEX,
@@ -279,7 +279,7 @@ def test_vector_search_client_model_serving_environment():
             host="testDogfod.com", credentials_strategy=ModelServingUserCredentials()
         )
 
-        with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+        with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
             with patch("databricks.sdk.service.serving.ServingEndpointsAPI.get", return_value=None):
                 vsTool = VectorSearchRetrieverTool(
                     index_name="catalog.schema.my_index_name",
@@ -295,7 +295,7 @@ def test_vector_search_client_model_serving_environment():
 
 
 def test_vector_search_client_non_model_serving_environment():
-    with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+    with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
         vsTool = VectorSearchRetrieverTool(
             index_name="catalog.schema.my_index_name",
             text_column="abc",
@@ -310,7 +310,7 @@ def test_vector_search_client_with_pat_workspace_client():
     w.config.auth_type = "pat"
     w.config.host = "https://testDogfod.com"
     w.config.token = "fakeToken"
-    with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+    with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
         with patch("databricks.sdk.service.serving.ServingEndpointsAPI.get", return_value=None):
             VectorSearchRetrieverTool(
                 index_name="catalog.schema.my_index_name",
@@ -334,7 +334,7 @@ def test_vector_search_client_with_sp_workspace_client():
     w.config.client_id = "fakeClientId"
     w.config.client_secret = "fakeClientSecret"
 
-    with patch("databricks.vector_search.client.VectorSearchClient") as mockVSClient:
+    with patch("databricks.ai_search.client.VectorSearchClient") as mockVSClient:
         with patch("databricks.sdk.service.serving.ServingEndpointsAPI.get", return_value=None):
             VectorSearchRetrieverTool(
                 index_name="catalog.schema.my_index_name",

pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
   "typing_extensions>=4.15.0",
   "pydantic>=2.10.0",
   "databricks-sdk>=0.49.0",
-  "databricks-vectorsearch>=0.57",
+  "databricks-ai-search>=0.73",
   "pandas>=2.2.0",
   "tiktoken>=0.8.0",
   "tabulate>=0.9.0",

src/databricks_ai_bridge/test_utils/vector_search.py

@@ -4,7 +4,7 @@
 from unittest.mock import MagicMock, create_autospec, patch
 
 import pytest
-from databricks.vector_search.index import VectorSearchIndex
+from databricks.ai_search.index import VectorSearchIndex
 
 INPUT_TEXTS = ["foo", "bar", "baz"]
 DEFAULT_VECTOR_DIMENSION = 4
@@ -142,7 +142,7 @@ def _get_index(
     mock_client = MagicMock()
     mock_client.get_index.side_effect = _get_index
     with mock.patch(
-        "databricks.vector_search.client.VectorSearchClient",
+        "databricks.ai_search.client.VectorSearchClient",
         return_value=mock_client,
     ):
         yield

src/databricks_ai_bridge/vector_search_retriever_tool.py

@@ -4,8 +4,8 @@
 from typing import Any, Dict, List, Optional
 
 import mlflow
+from databricks.ai_search.reranker import Reranker
 from databricks.sdk import WorkspaceClient
-from databricks.vector_search.reranker import Reranker
 from mlflow.entities import SpanType
 from mlflow.models.resources import (
     DatabricksServingEndpoint,

tests/integration_tests/vector_search/conftest.py

@@ -55,7 +55,7 @@ def vector_search_client(workspace_client):
     VectorSearchClient uses different env vars than the SDK,
     so we pass credentials explicitly.
     """
-    from databricks.vector_search.client import VectorSearchClient
+    from databricks.ai_search.client import VectorSearchClient
 
     workspace_url = workspace_client.config.host.rstrip("/")