From 92e7f9847e7b00ae277e1191c2d1acb45d3c1dbb Mon Sep 17 00:00:00 2001 From: Leonard Avevor Date: Tue, 14 Jan 2025 13:16:20 +0100 Subject: [PATCH] feat: Support rbac auth type in conversation_logger --- .../utilities/helpers/azure_search_helper.py | 36 ++++++++++++------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/code/backend/batch/utilities/helpers/azure_search_helper.py b/code/backend/batch/utilities/helpers/azure_search_helper.py index e6a7e94dd..c0314645d 100644 --- a/code/backend/batch/utilities/helpers/azure_search_helper.py +++ b/code/backend/batch/utilities/helpers/azure_search_helper.py @@ -265,15 +265,27 @@ def get_conversation_logger(self): ), ] - return AzureSearch( - azure_search_endpoint=self.env_helper.AZURE_SEARCH_SERVICE, - azure_search_key=( - self.env_helper.AZURE_SEARCH_KEY - if self.env_helper.is_auth_type_keys() - else None - ), - index_name=self.env_helper.AZURE_SEARCH_CONVERSATIONS_LOG_INDEX, - embedding_function=self.llm_helper.get_embedding_model().embed_query, - fields=fields, - user_agent="langchain chatwithyourdata-sa", - ) + if self.env_helper.AZURE_AUTH_TYPE == "rbac": + credential = DefaultAzureCredential() + return AzureSearch( + azure_search_endpoint=self.env_helper.AZURE_SEARCH_SERVICE, + azure_search_key=None, # Remove API key + index_name=self.env_helper.AZURE_SEARCH_CONVERSATIONS_LOG_INDEX, + embedding_function=self.llm_helper.get_embedding_model().embed_query, + fields=fields, + user_agent="langchain chatwithyourdata-sa", + credential=credential # Add token credential or send none so it is auto handled by AzureSearch library + ) + else: + return AzureSearch( + azure_search_endpoint=self.env_helper.AZURE_SEARCH_SERVICE, + azure_search_key=( + self.env_helper.AZURE_SEARCH_KEY + if self.env_helper.is_auth_type_keys() + else None + ), + index_name=self.env_helper.AZURE_SEARCH_CONVERSATIONS_LOG_INDEX, + embedding_function=self.llm_helper.get_embedding_model().embed_query, + fields=fields, + user_agent="langchain chatwithyourdata-sa", + )