From f5094cc90892e525c492e3454d39533400c74d55 Mon Sep 17 00:00:00 2001 From: Cameron Hutchison Date: Wed, 11 Dec 2024 10:30:02 -0600 Subject: [PATCH 1/3] Allow argument to be passed to VectorSearchClient to enable users to use Service Principals instead of PATs --- .../langchain/src/databricks_langchain/vectorstores.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/integrations/langchain/src/databricks_langchain/vectorstores.py b/integrations/langchain/src/databricks_langchain/vectorstores.py index e67f315..3fbf31f 100644 --- a/integrations/langchain/src/databricks_langchain/vectorstores.py +++ b/integrations/langchain/src/databricks_langchain/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: @@ -212,6 +216,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) and _INDEX_NAME_PATTERN.match(index_name)): raise ValueError( @@ -230,7 +235,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: From 81220c44efffdd583737c21813958133fea81315 Mon Sep 17 00:00:00 2001 From: Cameron Hutchison Date: Tue, 17 Dec 2024 18:46:15 -0600 Subject: [PATCH 2/3] Update integrations/langchain/src/databricks_langchain/vectorstores.py Co-authored-by: Prithvi Kannan <46332835+prithvikannan@users.noreply.github.com> --- .../langchain/src/databricks_langchain/vectorstores.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/integrations/langchain/src/databricks_langchain/vectorstores.py b/integrations/langchain/src/databricks_langchain/vectorstores.py index 3fbf31f..1af654b 100644 --- a/integrations/langchain/src/databricks_langchain/vectorstores.py +++ b/integrations/langchain/src/databricks_langchain/vectorstores.py @@ -235,7 +235,9 @@ def __init__( ) from e try: - self.index = VectorSearchClient(**client_args).get_index( + self.index = VectorSearchClient(**(client_args or {})).get_index( + endpoint_name=endpoint, index_name=index_name + ) endpoint_name=endpoint, index_name=index_name ) except Exception as e: From ad99744f55bb707afee6e3c14af97d926d9c3a42 Mon Sep 17 00:00:00 2001 From: Cameron Hutchison Date: Wed, 18 Dec 2024 10:38:50 -0600 Subject: [PATCH 3/3] Update vectorstores.py Resolve weird line duplication from accepting the proposed change --- integrations/langchain/src/databricks_langchain/vectorstores.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/vectorstores.py b/integrations/langchain/src/databricks_langchain/vectorstores.py index 1af654b..f513dfa 100644 --- a/integrations/langchain/src/databricks_langchain/vectorstores.py +++ b/integrations/langchain/src/databricks_langchain/vectorstores.py @@ -238,8 +238,6 @@ def __init__( self.index = VectorSearchClient(**(client_args or {})).get_index( endpoint_name=endpoint, index_name=index_name ) - endpoint_name=endpoint, index_name=index_name - ) except Exception as e: if endpoint is None and "Wrong vector search endpoint" in str(e): raise ValueError(