From e9d80880a15f1e285a4bcb89e04def059faad423 Mon Sep 17 00:00:00 2001 From: Jenny Date: Mon, 8 Jun 2026 23:51:33 -0700 Subject: [PATCH 1/3] Migrate from databricks-vectorsearch to databricks-ai-search `databricks-vectorsearch` is deprecated and has been renamed to `databricks-ai-search`. Per upstream's deprecation warning: databricks-vectorsearch is deprecated and has been renamed to databricks-ai-search. Imports under 'databricks.vector_search.*' will continue to work as a thin re-export of 'databricks.ai_search.*', but new code should switch to 'pip install databricks-ai-search' and 'from databricks.ai_search.* import ...'. Follow that recommendation literally: * pyproject.toml: replace `databricks-vectorsearch>=*` with `databricks-ai-search>=0.73` (root, openai, langchain, llamaindex). * All `from databricks.vector_search.* import ...` -> ai_search. * Mock-patch targets and docstring references updated to match. Symbol names are unchanged because `databricks-ai-search==0.73` exposes back-compat aliases: - databricks/ai_search/client.py:1090 VectorSearchClient = AISearchClient - databricks/ai_search/index.py:682 VectorSearchIndex = AISearchIndex - databricks/ai_search/exceptions.py:47 VectorSearchException = AISearchException The reranker / utils symbols (`Reranker`, `DatabricksReranker`, `CredentialStrategy`) were already used as-is and exist at the same relative paths in ai_search. Validated: * Fresh venv: only databricks-ai-search==0.73 installed (no vectorsearch). * 53/53 openai unit tests pass. * 258/258 langchain unit tests pass (4 skipped as before). * Smoke: `from databricks_openai.vector_search_retriever_tool import VectorSearchRetrieverTool` resolves `VectorSearchIndex` to `databricks.ai_search.index` even with the `databricks.vector_search` package completely absent. --- .../retrievers/databricks_rm.py | 2 +- integrations/langchain/pyproject.toml | 2 +- .../vector_search_retriever_tool.py | 6 ++-- .../src/databricks_langchain/vectorstores.py | 30 +++++++++---------- .../test_vector_search_retriever_tool.py | 10 +++---- .../tests/unit_tests/test_vectorstores.py | 4 +-- integrations/llamaindex/pyproject.toml | 2 +- .../vector_search_retriever_tool.py | 4 +-- .../test_vector_search_retriever_tool.py | 12 ++++---- integrations/openai/pyproject.toml | 2 +- .../vector_search_retriever_tool.py | 10 +++---- .../test_vector_search_retriever_tool.py | 14 ++++----- pyproject.toml | 2 +- .../test_utils/vector_search.py | 4 +-- .../vector_search_retriever_tool.py | 2 +- .../vector_search/conftest.py | 2 +- 16 files changed, 54 insertions(+), 54 deletions(-) diff --git a/integrations/dspy/src/databricks_dspy/retrievers/databricks_rm.py b/integrations/dspy/src/databricks_dspy/retrievers/databricks_rm.py index a8dde4d79..e78288052 100644 --- a/integrations/dspy/src/databricks_dspy/retrievers/databricks_rm.py +++ b/integrations/dspy/src/databricks_dspy/retrievers/databricks_rm.py @@ -39,7 +39,7 @@ class DatabricksRM(dspy.Retrieve): .. 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 diff --git a/integrations/langchain/pyproject.toml b/integrations/langchain/pyproject.toml index b274feb36..1e24651c7 100644 --- a/integrations/langchain/pyproject.toml +++ b/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", diff --git a/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py b/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py index de05f9a68..8ee5baa63 100644 --- a/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py +++ b/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,8 +24,8 @@ 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 - documentation `_ + `databricks.ai_search.client.VectorSearchIndex.similarity_search` when executing the tool. `See + documentation `_ to see the full set of supported keyword arguments, e.g. `score_threshold`. Also, see documentation for :class:`~databricks_ai_bridge.vector_search_retriever_tool.VectorSearchRetrieverToolMixin` for additional supported constructor diff --git a/integrations/langchain/src/databricks_langchain/vectorstores.py b/integrations/langchain/src/databricks_langchain/vectorstores.py index 032556223..16e7ad676 100644 --- a/integrations/langchain/src/databricks_langchain/vectorstores.py +++ b/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="", @@ -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,8 +422,8 @@ 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 - documentation `_ + kwargs: Additional keyword arguments to pass to `databricks.ai_search.client.VectorSearchIndex.similarity_search`. `See + documentation `_ to see the full set of supported keyword arguments Returns: @@ -464,8 +464,8 @@ 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 - documentation `_ + kwargs: Additional keyword arguments to pass to `databricks.ai_search.client.VectorSearchIndex.similarity_search`. `See + documentation `_ to see the full set of supported keyword arguments Returns: @@ -538,8 +538,8 @@ 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 - documentation `_ + kwargs: Additional keyword arguments to pass to `databricks.ai_search.client.VectorSearchIndex.similarity_search`. `See + documentation `_ to see the full set of supported keyword arguments Returns: @@ -591,8 +591,8 @@ 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 - documentation `_ + kwargs: Additional keyword arguments to pass to `databricks.ai_search.client.VectorSearchIndex.similarity_search`. `See + documentation `_ to see the full set of supported keyword arguments Returns: diff --git a/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py b/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py index d2373dc0d..1e07a0111 100644 --- a/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py +++ b/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 diff --git a/integrations/langchain/tests/unit_tests/test_vectorstores.py b/integrations/langchain/tests/unit_tests/test_vectorstores.py index 0b17b967b..5664db3aa 100644 --- a/integrations/langchain/tests/unit_tests/test_vectorstores.py +++ b/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, diff --git a/integrations/llamaindex/pyproject.toml b/integrations/llamaindex/pyproject.toml index b25ec623a..490e968c6 100644 --- a/integrations/llamaindex/pyproject.toml +++ b/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", diff --git a/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py b/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py index 2c92a2383..b6b023d6d 100644 --- a/integrations/llamaindex/src/databricks_llamaindex/vector_search_retriever_tool.py +++ b/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 ( diff --git a/integrations/llamaindex/tests/unit_tests/test_vector_search_retriever_tool.py b/integrations/llamaindex/tests/unit_tests/test_vector_search_retriever_tool.py index a987e3657..42eaaf9d6 100644 --- a/integrations/llamaindex/tests/unit_tests/test_vector_search_retriever_tool.py +++ b/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", diff --git a/integrations/openai/pyproject.toml b/integrations/openai/pyproject.toml index 006356b4a..93da54a13 100644 --- a/integrations/openai/pyproject.toml +++ b/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", diff --git a/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py b/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py index 5d3a96865..8fa05f7bf 100644 --- a/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py +++ b/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,8 +68,8 @@ 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 - documentation `_ + `databricks.ai_search.client.VectorSearchIndex.similarity_search` when executing the tool. `See + documentation `_ to see the full set of supported keyword arguments, e.g. `score_threshold`. Also, see documentation for :class:`~databricks_ai_bridge.vector_search_retriever_tool.VectorSearchRetrieverToolMixin` for additional supported constructor @@ -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: diff --git a/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py b/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py index 3f1efef8f..cee2f3da3 100644 --- a/integrations/openai/tests/unit_tests/test_vector_search_retriever_tool.py +++ b/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", diff --git a/pyproject.toml b/pyproject.toml index 3044efabd..e5c097e70 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/databricks_ai_bridge/test_utils/vector_search.py b/src/databricks_ai_bridge/test_utils/vector_search.py index 404078d49..d576b8c1a 100644 --- a/src/databricks_ai_bridge/test_utils/vector_search.py +++ b/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 diff --git a/src/databricks_ai_bridge/vector_search_retriever_tool.py b/src/databricks_ai_bridge/vector_search_retriever_tool.py index edb430a37..9de5956bd 100644 --- a/src/databricks_ai_bridge/vector_search_retriever_tool.py +++ b/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, diff --git a/tests/integration_tests/vector_search/conftest.py b/tests/integration_tests/vector_search/conftest.py index c343eb6f4..5f85fe45c 100644 --- a/tests/integration_tests/vector_search/conftest.py +++ b/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("/") From 75951f0c113b65bc7e61a164e588b97a02d6604a Mon Sep 17 00:00:00 2001 From: Jenny Date: Tue, 9 Jun 2026 14:40:40 -0700 Subject: [PATCH 2/3] Revert docstring URLs to the working `vector_search.html` form MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit's bulk sed also rewrote the Sphinx doc links in six docstrings to point at `databricks.ai_search.html`. That page currently 403s on api-docs.databricks.com — the docs site has not yet republished under the `databricks-ai-search` name and there is no redirect from the old URL. Until the docs team ships docs under the new package name (and adds redirects from the old), point users at the URL that actually resolves today. The page that opens at `databricks.vector_search.html` documents the same `VectorSearchIndex.similarity_search` API the bridge calls, so the link target is still correct content-wise. This reverts only the URL strings; the actual import paths (`from databricks.ai_search.* import ...`) are unchanged. --- .../databricks_langchain/vector_search_retriever_tool.py | 2 +- .../langchain/src/databricks_langchain/vectorstores.py | 8 ++++---- .../src/databricks_openai/vector_search_retriever_tool.py | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py b/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py index 8ee5baa63..8d1c980dc 100644 --- a/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py +++ b/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py @@ -25,7 +25,7 @@ class VectorSearchRetrieverTool(BaseTool, VectorSearchRetrieverToolMixin): **Note**: Any additional keyword arguments passed to the constructor will be passed along to `databricks.ai_search.client.VectorSearchIndex.similarity_search` when executing the tool. `See - documentation `_ + documentation `_ to see the full set of supported keyword arguments, e.g. `score_threshold`. Also, see documentation for :class:`~databricks_ai_bridge.vector_search_retriever_tool.VectorSearchRetrieverToolMixin` for additional supported constructor diff --git a/integrations/langchain/src/databricks_langchain/vectorstores.py b/integrations/langchain/src/databricks_langchain/vectorstores.py index 16e7ad676..efb49c23a 100644 --- a/integrations/langchain/src/databricks_langchain/vectorstores.py +++ b/integrations/langchain/src/databricks_langchain/vectorstores.py @@ -423,7 +423,7 @@ def similarity_search( 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.ai_search.client.VectorSearchIndex.similarity_search`. `See - documentation `_ + documentation `_ to see the full set of supported keyword arguments Returns: @@ -465,7 +465,7 @@ def similarity_search_with_score( 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.ai_search.client.VectorSearchIndex.similarity_search`. `See - documentation `_ + documentation `_ to see the full set of supported keyword arguments Returns: @@ -539,7 +539,7 @@ def similarity_search_by_vector( 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.ai_search.client.VectorSearchIndex.similarity_search`. `See - documentation `_ + documentation `_ to see the full set of supported keyword arguments Returns: @@ -592,7 +592,7 @@ def similarity_search_by_vector_with_score( 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.ai_search.client.VectorSearchIndex.similarity_search`. `See - documentation `_ + documentation `_ to see the full set of supported keyword arguments Returns: diff --git a/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py b/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py index 8fa05f7bf..364a520dd 100644 --- a/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py +++ b/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py @@ -69,7 +69,7 @@ class VectorSearchRetrieverTool(VectorSearchRetrieverToolMixin): **Note**: Any additional keyword arguments passed to the constructor will be passed along to `databricks.ai_search.client.VectorSearchIndex.similarity_search` when executing the tool. `See - documentation `_ + documentation `_ to see the full set of supported keyword arguments, e.g. `score_threshold`. Also, see documentation for :class:`~databricks_ai_bridge.vector_search_retriever_tool.VectorSearchRetrieverToolMixin` for additional supported constructor From 3b9c200f4bdce126f9c16a902e5ec0f42488968a Mon Sep 17 00:00:00 2001 From: Jenny Date: Tue, 9 Jun 2026 15:12:12 -0700 Subject: [PATCH 3/3] Docstring: refer to VectorSearchIndex by its real module (`.index`, not `.client`) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The original docstrings said `databricks.vector_search.client. VectorSearchIndex.similarity_search` — the same wrong-path incidental-re-export bug that PR #435 fixed in the actual code. The earlier bulk sed in this PR carried the bug forward into the new package name, producing `databricks.ai_search.client.VectorSearchIndex` which does not even resolve (the new `databricks.ai_search.client` imports `AISearchIndex`, not `VectorSearchIndex`; only the alias on `.index` exists). Fix the docstrings to point at the real home of the class: databricks.ai_search.client.VectorSearchIndex → databricks.ai_search.index.VectorSearchIndex --- .../databricks_langchain/vector_search_retriever_tool.py | 2 +- .../langchain/src/databricks_langchain/vectorstores.py | 8 ++++---- .../src/databricks_openai/vector_search_retriever_tool.py | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py b/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py index 8d1c980dc..bf6e561b1 100644 --- a/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py +++ b/integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py @@ -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.ai_search.client.VectorSearchIndex.similarity_search` when executing the tool. `See + `databricks.ai_search.index.VectorSearchIndex.similarity_search` when executing the tool. `See documentation `_ to see the full set of supported keyword arguments, e.g. `score_threshold`. Also, see documentation for diff --git a/integrations/langchain/src/databricks_langchain/vectorstores.py b/integrations/langchain/src/databricks_langchain/vectorstores.py index efb49c23a..4aa98654d 100644 --- a/integrations/langchain/src/databricks_langchain/vectorstores.py +++ b/integrations/langchain/src/databricks_langchain/vectorstores.py @@ -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.ai_search.client.VectorSearchIndex.similarity_search`. `See + kwargs: Additional keyword arguments to pass to `databricks.ai_search.index.VectorSearchIndex.similarity_search`. `See documentation `_ 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.ai_search.client.VectorSearchIndex.similarity_search`. `See + kwargs: Additional keyword arguments to pass to `databricks.ai_search.index.VectorSearchIndex.similarity_search`. `See documentation `_ 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.ai_search.client.VectorSearchIndex.similarity_search`. `See + kwargs: Additional keyword arguments to pass to `databricks.ai_search.index.VectorSearchIndex.similarity_search`. `See documentation `_ 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.ai_search.client.VectorSearchIndex.similarity_search`. `See + kwargs: Additional keyword arguments to pass to `databricks.ai_search.index.VectorSearchIndex.similarity_search`. `See documentation `_ to see the full set of supported keyword arguments diff --git a/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py b/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py index 364a520dd..cc602a136 100644 --- a/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py +++ b/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py @@ -68,7 +68,7 @@ class VectorSearchRetrieverTool(VectorSearchRetrieverToolMixin): ) **Note**: Any additional keyword arguments passed to the constructor will be passed along to - `databricks.ai_search.client.VectorSearchIndex.similarity_search` when executing the tool. `See + `databricks.ai_search.index.VectorSearchIndex.similarity_search` when executing the tool. `See documentation `_ to see the full set of supported keyword arguments, e.g. `score_threshold`. Also, see documentation for