From 9087b4c4ca2dc7a965530dfbd6e04a75d893fc2f Mon Sep 17 00:00:00 2001 From: Cameron Hutchison Date: Wed, 11 Dec 2024 11:00:46 -0600 Subject: [PATCH] Allow argument to be passed to VectorSearchClient to enable users to use Service Principals instead of PATs --- libs/databricks/langchain_databricks/vectorstores.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/libs/databricks/langchain_databricks/vectorstores.py b/libs/databricks/langchain_databricks/vectorstores.py index 1e2aa0a..ee8ec74 100644 --- a/libs/databricks/langchain_databricks/vectorstores.py +++ b/libs/databricks/langchain_databricks/vectorstores.py @@ -86,6 +86,10 @@ class DatabricksVectorSearch(VectorStore): Make sure the text column specified is in the index. columns: The list of column names to get when doing the search. Defaults to ``[primary_key, text_column]``. + client_args: Additional arguments to pass to the VectorSearchClient. + Allows you to pass in values like `service_principal_client_id` + and `service_principal_client_secret` for to allow for + service principal authentication instead of personal access token authentication. Instantiate: @@ -214,6 +218,7 @@ def __init__( embedding: Optional[Embeddings] = None, text_column: Optional[str] = None, columns: Optional[List[str]] = None, + client_args: Optional[Dict[str, Any]] = None, ): if not isinstance(index_name, str): raise ValueError( @@ -232,7 +237,7 @@ def __init__( ) from e try: - self.index = VectorSearchClient().get_index( + self.index = VectorSearchClient(**client_args).get_index( endpoint_name=endpoint, index_name=index_name ) except Exception as e: