Update variables

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/creators/llm_model_creator.py

@@ -30,21 +30,15 @@ def get_model(cls, model_name: str) -> AzureOpenAIChatCompletionClient:
 
     @classmethod
     def get_authentication_properties(cls) -> dict:
-        if get_identity_type() == IdentityType.SYSTEM_ASSIGNED:
+        if get_identity_type() in [
+            IdentityType.SYSTEM_ASSIGNED,
+            IdentityType.USER_ASSIGNED,
+        ]:
             # Create the token provider
             api_key = None
             token_provider = get_bearer_token_provider(
                 DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
             )
-        elif get_identity_type() == IdentityType.USER_ASSIGNED:
-            # Create the token provider
-            api_key = None
-            token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(
-                    managed_identity_client_id=os.environ["ClientId"]
-                ),
-                "https://cognitiveservices.azure.com/.default",
-            )
         else:
             token_provider = None
             api_key = os.environ["OpenAI__ApiKey"]

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/ai_search.py

@@ -289,12 +289,11 @@ async def add_entry_to_index(document: dict, vector_fields: dict, index_name: st
                 document["Question"].encode()
             ).decode("utf-8")
 
-            if identity_type == IdentityType.SYSTEM_ASSIGNED:
+            if identity_type in [
+                IdentityType.SYSTEM_ASSIGNED,
+                IdentityType.USER_ASSIGNED,
+            ]:
                 credential = DefaultAzureCredential()
-            elif identity_type == IdentityType.USER_ASSIGNED:
-                credential = DefaultAzureCredential(
-                    managed_identity_client_id=os.environ["ClientID"]
-                )
             else:
                 credential = AzureKeyCredential(
                     os.environ["AIService__AzureSearchOptions__Key"]

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/data_dictionary_creator.py

@@ -645,19 +645,14 @@ async def send_request_to_llm(self, system_prompt: str, input: str):
         api_version = os.environ["OpenAI__ApiVersion"]
         model = os.environ["OpenAI__CompletionDeployment"]
 
-        if get_identity_type() == IdentityType.SYSTEM_ASSIGNED:
+        if get_identity_type() in [
+            IdentityType.SYSTEM_ASSIGNED,
+            IdentityType.USER_ASSIGNED,
+        ]:
             token_provider = get_bearer_token_provider(
                 DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
             )
             api_key = None
-        elif get_identity_type() == IdentityType.USER_ASSIGNED:
-            token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(
-                    managed_identity_client_id=os.environ["FunctionApp__ClientId"]
-                ),
-                "https://cognitiveservices.azure.com/.default",
-            )
-            api_key = None
         else:
             token_provider = None
             api_key = os.environ["OpenAI__ApiKey"]