Support system assigned managed identity for embedding dependency (#1041

)
microsoft · Aug 6, 2024 · 6881d7f · 6881d7f
1 parent 9f17cfa
commit 6881d7f
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 15 deletions.
diff --git a/backend/settings.py b/backend/settings.py
@@ -175,16 +175,24 @@ def extract_embedding_dependency(self) -> Optional[dict]:
                 "type": "deployment_name",
                 "deployment_name": self.embedding_name
             }
-
-        elif self.embedding_endpoint and self.embedding_key:
-            return {
-                "type": "endpoint",
-                "endpoint": self.embedding_endpoint,
-                "authentication": {
-                    "type": "api_key",
-                    "api_key": self.embedding_key
+        elif self.embedding_endpoint:
+            if self.embedding_key:
+                return {
+                    "type": "endpoint",
+                    "endpoint": self.embedding_endpoint,
+                    "authentication": {
+                        "type": "api_key",
+                        "api_key": self.embedding_key
+                    }
+                }
+            else:
+                return {
+                    "type": "endpoint",
+                    "endpoint": self.embedding_endpoint,
+                    "authentication": {
+                        "type": "system_assigned_managed_identity"
+                    }
                 }
-            }
         else:   
             return None
 

diff --git a/tests/integration_tests/dotenv_templates/dotenv.jinja2 b/tests/integration_tests/dotenv_templates/dotenv.jinja2
@@ -1,7 +1,9 @@
 DATASOURCE_TYPE={{ DATASOURCE_TYPE }}
 AZURE_OPENAI_ENDPOINT={{ AZURE_OPENAI_ENDPOINT }}
 AZURE_OPENAI_MODEL={{ AZURE_OPENAI_MODEL }}
+{% if not USE_MI %}
 AZURE_OPENAI_KEY={{ AZURE_OPENAI_KEY }}
+{% endif %}
 AZURE_OPENAI_TEMPERATURE=0
 AZURE_OPENAI_TOP_P=1.0
 AZURE_OPENAI_MAX_TOKENS=1000
@@ -14,16 +16,18 @@ AZURE_OPENAI_EMBEDDING_NAME={{ AZURE_OPENAI_EMBEDDING_NAME }}
 {% endif %}
 {% if USE_AOAI_EMBEDDINGS %}
 AZURE_OPENAI_EMBEDDING_ENDPOINT={{ AZURE_OPENAI_ENDPOINT }}/openai/deployments/ada/embeddings?api-version=2023-03-15-preview
+{% if not USE_MI %}
 AZURE_OPENAI_EMBEDDING_KEY={{ AZURE_OPENAI_KEY }}
 {% endif %}
+{% endif %}
 {% if ENABLE_CHAT_HISTORY %}
 AZURE_COSMOSDB_ACCOUNT={{ AZURE_COSMOSDB_ACCOUNT }}
 AZURE_COSMOSDB_DATABASE={{ AZURE_COSMOSDB_DATABASE }}
 AZURE_COSMOSDB_CONVERSATIONS_CONTAINER={{ AZURE_COSMOSDB_CONVERSATIONS_CONTAINER }}
 AZURE_COSMOSDB_ACCOUNT_KEY={{ AZURE_COSMOSDB_ACCOUNT_KEY }}
 AZURE_COSMOSDB_ENABLE_FEEDBACK={{ AZURE_COSMOSDB_ENABLE_FEEDBACK }}
 {% endif %}
-{% if DATASOURCE_TYPE == "AZURE_COGNITIVE_SEARCH" %}
+{% if DATASOURCE_TYPE == "AzureCognitiveSearch" %}
 AZURE_SEARCH_SERVICE={{ AZURE_SEARCH_SERVICE }}
 AZURE_SEARCH_INDEX={{ AZURE_SEARCH_INDEX }}
 AZURE_SEARCH_KEY={{ AZURE_SEARCH_KEY }}
@@ -35,7 +39,7 @@ AZURE_SEARCH_URL_COLUMN=
 AZURE_SEARCH_VECTOR_COLUMNS=contentVector
 AZURE_SEARCH_QUERY_TYPE={{ AZURE_SEARCH_QUERY_TYPE }}
 AZURE_SEARCH_PERMITTED_GROUPS_COLUMN={{ AZURE_SEARCH_PERMITTED_GROUPS_COLUMN }}
-{% elif DATASOURCE_TYPE == "ELASTICSEARCH" %}
+{% elif DATASOURCE_TYPE == "Elasticsearch" %}
 ELASTICSEARCH_ENDPOINT={{ ELASTICSEARCH_ENDPOINT }}
 ELASTICSEARCH_ENCODED_API_KEY={{ ELASTICSEARCH_ENCODED_API_KEY }}
 ELASTICSEARCH_INDEX={{ ELASTICSEARCH_INDEX }}
@@ -48,7 +52,7 @@ ELASTICSEARCH_VECTOR_COLUMNS=text_embedding.predicted_value
 {% if USE_ELASTICSEARCH_EMBEDDINGS and ELASTICSEARCH_EMBEDDING_MODEL_ID %}
 ELASTICSEARCH_EMBEDDING_MODEL_ID={{ ELASTICSEARCH_EMBEDDING_MODEL_ID }}
 {% endif %}
-{% elif DATASOURCE_TYPE == "AZURE_COSMOS_DB" %}
+{% elif DATASOURCE_TYPE == "AzureCosmosDB" %}
 AZURE_COSMOSDB_MONGO_VCORE_CONNECTION_STRING={{ AZURE_COSMOSDB_MONGO_VCORE_CONNECTION_STRING }}
 AZURE_COSMOSDB_MONGO_VCORE_DATABASE={{ AZURE_COSMOSDB_MONGO_VCORE_DATABASE }}
 AZURE_COSMOSDB_MONGO_VCORE_CONTAINER={{ AZURE_COSMOSDB_MONGO_VCORE_CONTAINER }}
@@ -58,7 +62,7 @@ AZURE_COSMOSDB_MONGO_VCORE_FILENAME_COLUMN=
 AZURE_COSMOSDB_MONGO_VCORE_TITLE_COLUMN=
 AZURE_COSMOSDB_MONGO_VCORE_URL_COLUMN=
 AZURE_COSMOSDB_MONGO_VCORE_VECTOR_COLUMNS={{ AZURE_COSMOSDB_MONGO_VCORE_VECTOR_COLUMNS }}
-{% elif DATASOURCE_TYPE == "PINECONE" %}
+{% elif DATASOURCE_TYPE == "Pinecone" %}
 PINECONE_ENVIRONMENT={{ PINECONE_ENVIRONMENT }}
 PINECONE_API_KEY={{ PINECONE_API_KEY }}
 PINECONE_INDEX_NAME={{ PINECONE_INDEX_NAME }}
@@ -67,7 +71,7 @@ PINECONE_FILENAME_COLUMN=
 PINECONE_TITLE_COLUMN=
 PINECONE_URL_COLUMN=
 PINECONE_VECTOR_COLUMNS={{ PINECONE_VECTOR_COLUMNS }}
-{% elif DATASOURCE_TYPE == "AZURE_ML_INDEX" %}
+{% elif DATASOURCE_TYPE == "AzureMLIndex" %}
 AZURE_MLINDEX_NAME={{ AZURE_ML_INDEX_NAME }}
 AZURE_MLINDEX_VERSION={{ AZURE_ML_INDEX_VERSION }}
 AZURE_ML_PROJECT_RESOURCE_ID={{ AZURE_ML_PROJECT_RESOURCE_ID }}

diff --git a/tests/integration_tests/test_datasources.py b/tests/integration_tests/test_datasources.py
@@ -64,6 +64,10 @@ def use_aoai_embeddings(request):
 def use_elasticsearch_embeddings(request):
     return request.param
 
+@pytest.fixture(scope="function", params=[True, False], ids=["use_mi", "use_keys"])
+def use_mi(request):
+    return request.param
+
 
 @pytest.fixture(scope="function")
 def dotenv_rendered_template_path(
@@ -73,7 +77,8 @@ def dotenv_rendered_template_path(
     enable_chat_history,
     stream, 
     use_aoai_embeddings,
-    use_elasticsearch_embeddings
+    use_elasticsearch_embeddings, 
+    use_mi
 ):
     rendered_template_name = request.node.name.replace("[", "_").replace("]", "_")
     template_path = os.path.join(
@@ -92,6 +97,7 @@ def dotenv_rendered_template_path(
         dotenv_template_params["USE_ELASTICSEARCH_EMBEDDINGS"] = use_elasticsearch_embeddings
 
     dotenv_template_params["USE_AOAI_EMBEDDINGS"] = use_aoai_embeddings
+    dotenv_template_params["USE_MI"] = use_mi
 
     if use_aoai_embeddings or use_elasticsearch_embeddings:
         dotenv_template_params["AZURE_SEARCH_QUERY_TYPE"] = "vector"
@@ -140,6 +146,7 @@ async def test_dotenv(test_app: Quart, dotenv_template_params: dict[str, str]):
             }
         ]
     }
+
     test_client = test_app.test_client()
     response = await test_client.post(request_path, json=request_data)
     assert response.status_code == 200