✨ fix vector search tool

✨ fix vector search tool

Author: renardeinside

src/databricks/labs/mcp/servers/unity_catalog/cli.py

@@ -32,6 +32,12 @@ class CliSettings(BaseSettings):
         validation_alias=AliasChoices("g", "genie_space_ids"),
     )
 
+    vector_search_num_results: int = Field(
+        default=5,
+        description="Number of results to return from vector search queries",
+        validation_alias=AliasChoices("vn", "vector_search_num_results", "vector_num_results"),
+    )
+
     def get_catalog_name(self):
         return self.schema_full_name.split(".")[0] if self.schema_full_name else None
 

src/databricks/labs/mcp/servers/unity_catalog/tools/vector_search.py

@@ -1,62 +1,82 @@
-import io
 import json
-from contextlib import redirect_stdout
-
-from databricks_openai import VectorSearchRetrieverTool
+from pydantic import BaseModel
 from databricks.sdk import WorkspaceClient
+from databricks.vector_search.client import VectorSearchClient
+from databricks.labs.mcp.servers.unity_catalog.tools.base_tool import BaseTool
+from databricks.labs.mcp.servers.unity_catalog.cli import CliSettings
+from mcp.types import TextContent, Tool as ToolSpec
 
-from mcp.types import Tool as ToolSpec, TextContent
+# Constant storing vector index content vector column name
+CONTENT_VECTOR_COLUMN_NAME = "__db_content_vector"
 
-from databricks.labs.mcp.servers.unity_catalog.tools.base_tool import BaseTool
+
+class QueryInput(BaseModel):
+    query: str
 
 
 class VectorSearchTool(BaseTool):
-    def __init__(self, tool_obj: VectorSearchRetrieverTool):
-        self.tool_obj = tool_obj
-        tool_info = tool_obj.tool["function"]
-        llm_friendly_tool_name = tool_info["name"]
+    def __init__(self, endpoint_name: str, index_name: str, tool_name: str, columns: list[str], num_results: int = 5):
+        self.endpoint_name = endpoint_name
+        self.index_name = index_name
+        self.tool_name = tool_name
+        self.columns = columns
+        self.num_results = num_results
+
         tool_spec = ToolSpec(
-            name=llm_friendly_tool_name,
-            description=tool_info["description"],
-            inputSchema=tool_info["parameters"],
+            name=tool_name,
+            description=f"Searches the vector index `{index_name}`.",
+            inputSchema=QueryInput.model_json_schema(),
         )
-        super().__init__(tool_spec=tool_spec)
+        super().__init__(tool_spec)
 
     def execute(self, **kwargs):
-        """
-        Executes the vector search tool with the provided arguments.
-        """
-        # Create a buffer to capture stdout from vector search client
-        # print statements
-        f = io.StringIO()
-        with redirect_stdout(f):
-            res = self.tool_obj.execute(**kwargs)
-            return [
-                TextContent(
-                    type="text",
-                    text=json.dumps(vs_res),
-                )
-                for vs_res in res
-            ]
+        model = QueryInput.model_validate(kwargs)
+        vsc = VectorSearchClient(disable_notice=True)
+
+        index = vsc.get_index(index_name=self.index_name)
+
+        results = index.similarity_search(
+            query_text=model.query,
+            columns=self.columns,
+            num_results=self.num_results,
+        )
+
+        docs = results.get("result", {}).get("data_array", [])
+
+        return [TextContent(type="text", text=json.dumps(docs, indent=2))]
+
+
+def get_table_columns(workspace_client: WorkspaceClient, full_table_name: str) -> list[str]:
+    table_info = workspace_client.tables.get(full_table_name)
+    return [
+        col.name
+        for col in table_info.columns
+        if col.name != CONTENT_VECTOR_COLUMN_NAME
+    ]
 
 
 def _list_vector_search_tools(
-    workspace_client: WorkspaceClient, catalog_name: str, schema_name: str
+    workspace_client: WorkspaceClient, catalog_name: str, schema_name: str, vector_search_num_results: int
 ) -> list[VectorSearchTool]:
     tools = []
     for table in workspace_client.tables.list(
         catalog_name=catalog_name, schema_name=schema_name
     ):
-        # TODO: support filtering tables by securable kind (e.g. by making securable
-        # kind accessible here)
         if not table.properties or "model_endpoint_url" not in table.properties:
             continue
-        tool_obj = VectorSearchRetrieverTool(index_name=table.full_name)
-        tools.append(VectorSearchTool(tool_obj))
+
+        endpoint = table.properties["model_endpoint_url"]
+        index_name = table.full_name
+        tool_name = f"vector_search_{table.name}"
+
+        columns = get_table_columns(workspace_client, index_name)
+
+        tools.append(VectorSearchTool(endpoint, index_name, tool_name, columns, vector_search_num_results))
+
     return tools
 
 
-def list_vector_search_tools(settings) -> list[VectorSearchTool]:
+def list_vector_search_tools(settings: CliSettings) -> list[VectorSearchTool]:
     workspace_client = WorkspaceClient()
     catalog_name, schema_name = settings.schema_full_name.split(".")
-    return _list_vector_search_tools(workspace_client, catalog_name, schema_name)
+    return _list_vector_search_tools(workspace_client, catalog_name, schema_name, settings.vector_search_num_results)

tests/test_vector_search.py

@@ -9,6 +9,7 @@
 class DummyTable:
     def __init__(self, full_name, properties):
         self.full_name = full_name
+        self.name = full_name.split(".")[-1]
         self.properties = properties
 
 
@@ -21,6 +22,17 @@ def list(self, catalog_name=None, schema_name=None):
             DummyTable(full_name="cat.sch.tbl2", properties={}),
         ]
 
+    def get(self, full_table_name):
+        # Mock get_table_columns behavior
+        class DummyColumn:
+            def __init__(self, name):
+                self.name = name
+
+        class DummyTableInfo:
+            columns = [DummyColumn("col1"), DummyColumn("col2"), DummyColumn("__db_content_vector")]
+
+        return DummyTableInfo()
+
 
 class DummyWorkspaceClient:
     def __init__(self):
@@ -35,47 +47,44 @@ class DummySettings:
     "databricks.labs.mcp.servers.unity_catalog.tools.vector_search.WorkspaceClient",
     new=DummyWorkspaceClient,
 )
-@mock.patch(
-    "databricks.labs.mcp.servers.unity_catalog.tools.vector_search.VectorSearchRetrieverTool"
-)
-def test_list_vector_search_tools_filters_and_returns_expected(
-    MockVectorSearchRetrieverTool,
-):
-    MockVectorSearchRetrieverTool.side_effect = lambda index_name: mock.Mock(
-        tool={"function": {"name": index_name, "description": "", "parameters": {}}},
-        index_name=index_name,
-    )
+def test_list_vector_search_tools_filters_and_returns_expected():
     settings = DummySettings()
     tools = list_vector_search_tools(settings)
     assert len(tools) == 1
     tool = tools[0]
     assert isinstance(tool, VectorSearchTool)
-    assert tool.tool_obj.index_name == "cat.sch.tbl1"
+    assert tool.index_name == "cat.sch.tbl1"
+    assert tool.columns == ["col1", "col2"]  # filtered out "__db_content_vector"
 
 
 def test_internal_list_vector_search_tools_direct():
-    with mock.patch(
-        "databricks.labs.mcp.servers.unity_catalog.tools.vector_search.VectorSearchRetrieverTool"
-    ) as MockVectorSearchRetrieverTool:
-        MockVectorSearchRetrieverTool.side_effect = lambda index_name: mock.Mock(
-            tool={
-                "function": {"name": index_name, "description": "", "parameters": {}}
-            },
-            index_name=index_name,
-        )
-        client = DummyWorkspaceClient()
-        tools = _list_vector_search_tools(client, "cat", "sch")
-        assert len(tools) == 1
-        assert tools[0].tool_obj.index_name == "cat.sch.tbl1"
-
-
-def test_vector_search_tool_execute():
-    tool_obj = mock.Mock()
-    tool_obj.tool = {
-        "function": {"name": "vs_tool", "description": "", "parameters": {}}
+    client = DummyWorkspaceClient()
+    tools = _list_vector_search_tools(client, "cat", "sch")
+    assert len(tools) == 1
+    assert isinstance(tools[0], VectorSearchTool)
+    assert tools[0].index_name == "cat.sch.tbl1"
+    assert tools[0].columns == ["col1", "col2"]
+
+
+@mock.patch("databricks.labs.mcp.servers.unity_catalog.tools.vector_search.VectorSearchClient")
+def test_vector_search_tool_execute(MockVectorSearchClient):
+    mock_index = mock.Mock()
+    mock_index.similarity_search.return_value = {
+        "result": {"data_array": [{"id": 1, "score": 0.9}]}
     }
-    tool_obj.execute.return_value = [{"foo": "bar"}]
-    tool = VectorSearchTool(tool_obj)
-    result = tool.execute(query="test")
+
+    # Make get_index return our mock_index
+    MockVectorSearchClient.return_value.get_index.return_value = mock_index
+
+    tool = VectorSearchTool(
+        endpoint_name="endpoint1",
+        index_name="cat.sch.tbl1",
+        tool_name="vector_search_test",
+        columns=["col1", "col2"],
+    )
+
+    result = tool.execute(query="test query")
+
     assert isinstance(result, list)
-    assert result[0].text == '{"foo": "bar"}'
+    assert result[0].text.strip().startswith("[")  # It should be JSON string
+    assert "score" in result[0].text