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 ec7c70c6b..0b17b967b 100644
--- a/integrations/langchain/tests/unit_tests/test_vectorstores.py
+++ b/integrations/langchain/tests/unit_tests/test_vectorstores.py
@@ -4,7 +4,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from databricks.vector_search.client 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 cb2fa8fd6..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,7 +6,7 @@
 import pytest
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.credentials_provider import ModelServingUserCredentials
-from databricks.vector_search.client 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 2596f6457..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,7 +2,7 @@
 import logging
 from typing import Any, Dict, List, Optional, Tuple
 
-from databricks.vector_search.client 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 459130f05..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,7 +8,7 @@
 import pytest
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.credentials_provider import ModelServingUserCredentials
-from databricks.vector_search.client 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 ffdddabce..404078d49 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.client import VectorSearchIndex
+from databricks.vector_search.index import VectorSearchIndex
 
 INPUT_TEXTS = ["foo", "bar", "baz"]
 DEFAULT_VECTOR_DIMENSION = 4