Bug: Allow Databricks Service Principal Authentication for VectorSearchRetrieverTool.

[tomvmeer]

Using databricks-ai-bridge==0.7.1
databricks-langchain==0.7.1

When initiating the VectorSearchRetrieverTool instance for Langchain, (either with a workspace client or with environment variables), the only supported flow seems to be with personal access token.

See entry point in:
integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py line 62:
dbvs = DatabricksVectorSearch(**kwargs)

integrations/langchain/src/databricks_langchain/vectorstores.py lines 255-266:

client_args = client_args or {}
client_args.setdefault("disable_notice", True)
            if (
                workspace_client is not None
                and workspace_client.config.auth_type == "model_serving_user_credentials"
            ):
                client_args.setdefault(
                    "credential_strategy", CredentialStrategy.MODEL_SERVING_USER_CREDENTIALS
                )
            self.index = VectorSearchClient(**client_args).get_index(
                endpoint_name=endpoint, index_name=index_name
            )

When using a service principal, workspace_client.config.auth_type is not model_serving_user_credentials and thus the client args remain basically empty. The method would suggest you can set client_args yourself, but actually these are being removed by the retrieval tool:
integrations/langchain/src/databricks_langchain/vector_search_retriever_tool.py lines 52-62:

 def _validate_tool_inputs(self):
        kwargs = {
            "index_name": self.index_name,
            "embedding": self.embedding,
            "text_column": self.text_column,
            "doc_uri": self.doc_uri,
            "primary_key": self.primary_key,
            "columns": self.columns,
            "workspace_client": self.workspace_client,
            "include_score": self.include_score,
        }
        dbvs = DatabricksVectorSearch(**kwargs)

To reproduce, just try to instantiate a VectorSearchRetrieverTool instance with service principal credentials.

[tomvmeer]

Current work around is patching the _validate_tool_inputs method;

@model_validator(mode="after")
    def _validate_tool_inputs(self):
        kwargs = {
            "index_name": self.index_name,
            "embedding": self.embedding,
            "text_column": self.text_column,
            "doc_uri": self.doc_uri,
            "primary_key": self.primary_key,
            "columns": self.columns,
            "workspace_client": self.workspace_client,
            "include_score": self.include_score,
            "client_args": self.client_args # <----- PATCH
        }
        dbvs = DatabricksVectorSearch(**kwargs)
        self._vector_store = dbvs

        self.name = self._get_tool_name()
        self.description = self.tool_description or self._get_default_tool_description(
            IndexDetails(dbvs.index)
        )
        self.resources = self._get_resources(
            self.index_name,
            (self.embedding.endpoint if isinstance(self.embedding, DatabricksEmbeddings) else None),
            IndexDetails(dbvs.index),
        )
        return self

[Markpajr]

Came across this when searching for the ability to disable notice while developing in a notebook using the vectorstore retriever. I think the workaround you mentioned here should be added directly to the code, as it makes sense that if you use DatabricksVectorSearch.as_retriever, The client_args should be passed all the way through.

The current code means it's impossible to disable notice when using as_retriever.

[tpmullee]

I can implement this by adding client_args to VectorSearchRetrieverTool and passing it through to DatabricksVectorSearch(**kwargs)..

That way VectorSearchClient(**client_args) is honored (SP/M2M auth, and flags like disable_notice). + I’ll include unit tests and docstring updates. Unless maintainers prefer a different param name (e.g., vs_client_args) >> I’ll open a PR.

lmk & thanks!