From 131252d6f9e72b7a33861557b4eb1bbc3fc434ec Mon Sep 17 00:00:00 2001 From: Jenny Date: Mon, 8 Jun 2026 11:34:45 -0700 Subject: [PATCH 1/5] Import VectorSearchIndex from its canonical module `databricks.vector_search.client` re-exported `VectorSearchIndex` only as a side effect of doing `from databricks.vector_search.index import VectorSearchIndex` at the top of `client.py`. A recent release of `databricks-vectorsearch` (and/or the new `databricks-ai-search` package) removes that line, breaking every consumer that did `from databricks.vector_search.client import VectorSearchIndex`. The canonical home of `VectorSearchIndex` has been `databricks.vector_ search.index` since at least 0.55. Import from there, with a try/except preferring the new `databricks.ai_search.index` location so the bridge keeps working regardless of which underlying package the user installs. Reproduces user's traceback from the app-templates LangGraph template: ImportError: cannot import name 'VectorSearchIndex' from 'databricks.vector_search.client' --- integrations/langchain/tests/unit_tests/test_vectorstores.py | 5 ++++- .../tests/unit_tests/test_vector_search_retriever_tool.py | 5 ++++- .../src/databricks_openai/vector_search_retriever_tool.py | 5 ++++- .../tests/unit_tests/test_vector_search_retriever_tool.py | 5 ++++- src/databricks_ai_bridge/test_utils/vector_search.py | 5 ++++- 5 files changed, 20 insertions(+), 5 deletions(-) diff --git a/integrations/langchain/tests/unit_tests/test_vectorstores.py b/integrations/langchain/tests/unit_tests/test_vectorstores.py index ec7c70c6b..670721f01 100644 --- a/integrations/langchain/tests/unit_tests/test_vectorstores.py +++ b/integrations/langchain/tests/unit_tests/test_vectorstores.py @@ -4,7 +4,10 @@ from unittest.mock import MagicMock, patch import pytest -from databricks.vector_search.client import VectorSearchIndex +try: + from databricks.ai_search.index import VectorSearchIndex +except ImportError: + from databricks.vector_search.index import VectorSearchIndex from databricks.vector_search.reranker import DatabricksReranker, Reranker from databricks_ai_bridge.test_utils.vector_search import ( ALL_INDEX_NAMES, 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 cb2fa8fd6..a58dc5014 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 @@ -6,7 +6,10 @@ import pytest from databricks.sdk import WorkspaceClient from databricks.sdk.credentials_provider import ModelServingUserCredentials -from databricks.vector_search.client import VectorSearchIndex +try: + from databricks.ai_search.index import VectorSearchIndex +except ImportError: + 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 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 2596f6457..921518e59 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,10 @@ import logging from typing import Any, Dict, List, Optional, Tuple -from databricks.vector_search.client import VectorSearchIndex +try: + from databricks.ai_search.index import VectorSearchIndex +except ImportError: + from databricks.vector_search.index import VectorSearchIndex from databricks_ai_bridge.utils.vector_search import ( IndexDetails, RetrieverSchema, 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 459130f05..a1169558d 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 @@ -8,7 +8,10 @@ import pytest from databricks.sdk import WorkspaceClient from databricks.sdk.credentials_provider import ModelServingUserCredentials -from databricks.vector_search.client import VectorSearchIndex +try: + from databricks.ai_search.index import VectorSearchIndex +except ImportError: + 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 diff --git a/src/databricks_ai_bridge/test_utils/vector_search.py b/src/databricks_ai_bridge/test_utils/vector_search.py index ffdddabce..8b6fc4bac 100644 --- a/src/databricks_ai_bridge/test_utils/vector_search.py +++ b/src/databricks_ai_bridge/test_utils/vector_search.py @@ -4,7 +4,10 @@ from unittest.mock import MagicMock, create_autospec, patch import pytest -from databricks.vector_search.client import VectorSearchIndex +try: + from databricks.ai_search.index import VectorSearchIndex +except ImportError: + from databricks.vector_search.index import VectorSearchIndex INPUT_TEXTS = ["foo", "bar", "baz"] DEFAULT_VECTOR_DIMENSION = 4 From 9ba3d63b550500845bdfc076e43c7d756c56cfec Mon Sep 17 00:00:00 2001 From: Jenny Date: Mon, 8 Jun 2026 12:06:23 -0700 Subject: [PATCH 2/5] Fix lint and install databricks-openai from source in langchain tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ruff format/check: add the blank line ruff format wants above the try/except import block; let ruff --fix re-sort the import sections so I001 is happy. - integrations/langchain depends on databricks-openai (via databricks_langchain.utils), but its pyproject left databricks-openai to come from PyPI. With PyPI now shipping databricks-vectorsearch >=0.74 (which removed the VectorSearchIndex re-export from databricks.vector_search.client), the PyPI databricks-openai 0.15.0 blows up at import time and the langchain test job can't even collect. Add `databricks-openai = { path = "../openai", editable = true }` to [tool.uv.sources] so CI uses the (fixed) source. The openai_cross_version_test matrix (v0.3.0/v0.4.0/v0.5.0) checks out historical tags whose openai integration has the same broken import that this PR fixes for HEAD. Those tests can't be fixed by this PR because the historical sources are frozen — they need either a pin on databricks-vectorsearch<0.74 in the matrix, or to be dropped. Tracking separately. --- integrations/langchain/pyproject.toml | 1 + integrations/langchain/tests/unit_tests/test_vectorstores.py | 1 + .../tests/unit_tests/test_vector_search_retriever_tool.py | 1 + .../openai/tests/unit_tests/test_vector_search_retriever_tool.py | 1 + src/databricks_ai_bridge/test_utils/vector_search.py | 1 + 5 files changed, 5 insertions(+) diff --git a/integrations/langchain/pyproject.toml b/integrations/langchain/pyproject.toml index 257ba1643..b274feb36 100644 --- a/integrations/langchain/pyproject.toml +++ b/integrations/langchain/pyproject.toml @@ -52,6 +52,7 @@ build-backend = "hatchling.build" [tool.uv.sources] databricks-ai-bridge = { path = "../../", editable = true } +databricks-openai = { path = "../openai", editable = true } [tool.hatch.build] include = [ diff --git a/integrations/langchain/tests/unit_tests/test_vectorstores.py b/integrations/langchain/tests/unit_tests/test_vectorstores.py index 670721f01..df790d1e3 100644 --- a/integrations/langchain/tests/unit_tests/test_vectorstores.py +++ b/integrations/langchain/tests/unit_tests/test_vectorstores.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock, patch import pytest + try: from databricks.ai_search.index import VectorSearchIndex except ImportError: 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 a58dc5014..a2ac1f61c 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 @@ -6,6 +6,7 @@ import pytest from databricks.sdk import WorkspaceClient from databricks.sdk.credentials_provider import ModelServingUserCredentials + try: from databricks.ai_search.index import VectorSearchIndex except ImportError: 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 a1169558d..0fcb8ae0c 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 @@ -8,6 +8,7 @@ import pytest from databricks.sdk import WorkspaceClient from databricks.sdk.credentials_provider import ModelServingUserCredentials + try: from databricks.ai_search.index import VectorSearchIndex except ImportError: diff --git a/src/databricks_ai_bridge/test_utils/vector_search.py b/src/databricks_ai_bridge/test_utils/vector_search.py index 8b6fc4bac..465f5507a 100644 --- a/src/databricks_ai_bridge/test_utils/vector_search.py +++ b/src/databricks_ai_bridge/test_utils/vector_search.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock, create_autospec, patch import pytest + try: from databricks.ai_search.index import VectorSearchIndex except ImportError: From 6fab1a443901b886ed97f66628db8f89cd6e02d4 Mon Sep 17 00:00:00 2001 From: Jenny Date: Mon, 8 Jun 2026 12:23:53 -0700 Subject: [PATCH 3/5] CI: pin databricks-vectorsearch<0.74 in openai cross-version matrix The openai_test (v0.3.0/v0.4.0/v0.5.0) jobs check out historical tags of integrations/openai and run them against current PyPI. Those historical sources import `VectorSearchIndex` from `databricks.vector_search.client` (the bug this PR fixes for HEAD), and `databricks-vectorsearch>=0.74` removed that re-export. The frozen sources can't be fixed, so constrain the env to the version range they were validated against. --- .github/workflows/main.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 8326de7f4..0dc6e7986 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -235,6 +235,12 @@ jobs: pip install . pip install integrations/openai --group dev pip install "${{ matrix.version.mlflow }}" + # Pin databricks-vectorsearch to the version range these historical + # openai-integration tags were validated against. >=0.74 removed the + # `VectorSearchIndex` re-export from `databricks.vector_search.client` + # that those historical sources import; fixing the source isn't an + # option (frozen tags), so we constrain the env instead. + pip install "databricks-vectorsearch<0.74" - name: Run tests run: | # Only testing initialization since functionality can change From 16ca3d98a2bb8b65c0576127130441fdd4a988a5 Mon Sep 17 00:00:00 2001 From: Jenny Date: Mon, 8 Jun 2026 19:27:13 -0700 Subject: [PATCH 4/5] Revert "CI: pin databricks-vectorsearch<0.74 in openai cross-version matrix" This reverts commit 6fab1a443901b886ed97f66628db8f89cd6e02d4. --- .github/workflows/main.yml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 0dc6e7986..8326de7f4 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -235,12 +235,6 @@ jobs: pip install . pip install integrations/openai --group dev pip install "${{ matrix.version.mlflow }}" - # Pin databricks-vectorsearch to the version range these historical - # openai-integration tags were validated against. >=0.74 removed the - # `VectorSearchIndex` re-export from `databricks.vector_search.client` - # that those historical sources import; fixing the source isn't an - # option (frozen tags), so we constrain the env instead. - pip install "databricks-vectorsearch<0.74" - name: Run tests run: | # Only testing initialization since functionality can change From 418ee3b8afdecea2a5a6e18c55634c97cf713ffc Mon Sep 17 00:00:00 2001 From: Jenny Date: Mon, 8 Jun 2026 21:52:20 -0700 Subject: [PATCH 5/5] Simplify: drop ai_search try/except, import canonical path directly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The earlier try/except favoring `databricks.ai_search.index` was speculative future-proofing for a possible package rename, but `databricks-ai-search==0.73` itself exposes `VectorSearchIndex` only as a back-compat alias for `AISearchIndex` (`databricks/ai_search/index.py:682`). Today and on every released `databricks-vectorsearch` from 0.50 onward, `databricks.vector_search. index.VectorSearchIndex` is the canonical class — so a one-line direct import suffices and the cascade was just indirection. Validated under simulated `databricks-vectorsearch>=0.74` breakage (legacy `databricks.vector_search.client.VectorSearchIndex` attribute removed): - reproduces the user's `ImportError` for the legacy path - fixed `databricks_openai.vector_search_retriever_tool` imports cleanly - all 53 openai unit tests pass under the simulated breakage --- .../langchain/tests/unit_tests/test_vectorstores.py | 6 +----- .../tests/unit_tests/test_vector_search_retriever_tool.py | 6 +----- .../src/databricks_openai/vector_search_retriever_tool.py | 5 +---- .../tests/unit_tests/test_vector_search_retriever_tool.py | 6 +----- src/databricks_ai_bridge/test_utils/vector_search.py | 6 +----- 5 files changed, 5 insertions(+), 24 deletions(-) diff --git a/integrations/langchain/tests/unit_tests/test_vectorstores.py b/integrations/langchain/tests/unit_tests/test_vectorstores.py index df790d1e3..0b17b967b 100644 --- a/integrations/langchain/tests/unit_tests/test_vectorstores.py +++ b/integrations/langchain/tests/unit_tests/test_vectorstores.py @@ -4,11 +4,7 @@ from unittest.mock import MagicMock, patch import pytest - -try: - from databricks.ai_search.index import VectorSearchIndex -except ImportError: - from databricks.vector_search.index import VectorSearchIndex +from databricks.vector_search.index import VectorSearchIndex from databricks.vector_search.reranker import DatabricksReranker, Reranker from databricks_ai_bridge.test_utils.vector_search import ( ALL_INDEX_NAMES, 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 a2ac1f61c..a987e3657 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 @@ -6,11 +6,7 @@ import pytest from databricks.sdk import WorkspaceClient from databricks.sdk.credentials_provider import ModelServingUserCredentials - -try: - from databricks.ai_search.index import VectorSearchIndex -except ImportError: - from databricks.vector_search.index import VectorSearchIndex +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 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 921518e59..5d3a96865 100644 --- a/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py +++ b/integrations/openai/src/databricks_openai/vector_search_retriever_tool.py @@ -2,10 +2,7 @@ import logging from typing import Any, Dict, List, Optional, Tuple -try: - from databricks.ai_search.index import VectorSearchIndex -except ImportError: - from databricks.vector_search.index import VectorSearchIndex +from databricks.vector_search.index import VectorSearchIndex from databricks_ai_bridge.utils.vector_search import ( IndexDetails, RetrieverSchema, 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 0fcb8ae0c..3f1efef8f 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 @@ -8,11 +8,7 @@ import pytest from databricks.sdk import WorkspaceClient from databricks.sdk.credentials_provider import ModelServingUserCredentials - -try: - from databricks.ai_search.index import VectorSearchIndex -except ImportError: - from databricks.vector_search.index import VectorSearchIndex +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 diff --git a/src/databricks_ai_bridge/test_utils/vector_search.py b/src/databricks_ai_bridge/test_utils/vector_search.py index 465f5507a..404078d49 100644 --- a/src/databricks_ai_bridge/test_utils/vector_search.py +++ b/src/databricks_ai_bridge/test_utils/vector_search.py @@ -4,11 +4,7 @@ from unittest.mock import MagicMock, create_autospec, patch import pytest - -try: - from databricks.ai_search.index import VectorSearchIndex -except ImportError: - from databricks.vector_search.index import VectorSearchIndex +from databricks.vector_search.index import VectorSearchIndex INPUT_TEXTS = ["foo", "bar", "baz"] DEFAULT_VECTOR_DIMENSION = 4