User security context to model args instead of user json (#1318)

Author: darnissan

README.md

@@ -151,6 +151,7 @@ Note: RBAC assignments can take a few minutes before becoming effective.
     |AZURE_OPENAI_SYSTEM_MESSAGE|No|You are an AI assistant that helps people find information.|A brief description of the role and tone the model should use|
     |AZURE_OPENAI_STREAM|No|True|Whether or not to use streaming for the response. Note: Setting this to true prevents the use of prompt flow.|
     |AZURE_OPENAI_EMBEDDING_NAME|Only if using vector search using an Azure OpenAI embedding model||The name of your embedding model deployment if using vector search.
+    |MS_DEFENDER_ENABLED|Yes|True|Whether or not the Microsoft Defender for Cloud's threat protection for AI workloads plan is enabled on your subscription or not , for more details [Microsoft Defender for Cloud documentation](https://learn.microsoft.com/azure/defender-for-cloud/gain-end-user-context-ai).|
 
     See the [documentation](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/reference#example-response-2) for more information on these parameters.
 

README_azd.md

@@ -50,6 +50,15 @@ Run the following commands based on what you want to customize:
 * `azd env set AZURE_FORMRECOGNIZER_SERVICE_RESOURCE_GROUP {Name of existing resource group that Form Recognizer service is provisioned to}`.
 * `azd env set AZURE_FORMRECOGNIZER_SKU_NAME {Name of Form Recognizer SKU}`. Defaults to 'S0'.
 
+If you have enabled Microsoft Defender for Cloud's threat protection for AI workloads on your Azure OpenAI resource and want to add user context to alerts, you can do so by using the azd commands below. If you haven't set up threat protection yet, please follow this guide: [Microsoft Defender for Cloud documentation](https://learn.microsoft.com/azure/defender-for-cloud/gain-end-user-context-ai).
+
+To add security context to your alerts, run the following command:
+
+```bash
+azd env set MS_DEFENDER_ENABLED true
+```
+
+
 1. Run `azd auth login` to login to your Azure account.
 1. Run `azd up` to provision Azure resources and deploy this sample to those resources. This also runs a script to build the search index based on files in the `./data` folder.
 1. After the application has been successfully deployed you will see a URL printed to the console.  Click that URL to interact with the application in your browser.

app.py

@@ -274,12 +274,12 @@ def prepare_model_args(request_body, request_headers):
                     messages.append(messages_helper)
 
 
-    user_json = None
+    user_security_context = None
     if (MS_DEFENDER_ENABLED):
         authenticated_user_details = get_authenticated_user_details(request_headers)
-        conversation_id = request_body.get("conversation_id", None)
         application_name = app_settings.ui.title
-        user_json = get_msdefender_user_json(authenticated_user_details, request_headers, conversation_id, application_name)
+        user_security_context = get_msdefender_user_json(authenticated_user_details, request_headers, application_name )  # security component introduced here https://learn.microsoft.com/en-us/azure/defender-for-cloud/gain-end-user-context-ai
+    
 
     model_args = {
         "messages": messages,
@@ -288,8 +288,7 @@ def prepare_model_args(request_body, request_headers):
         "top_p": app_settings.azure_openai.top_p,
         "stop": app_settings.azure_openai.stop_sequence,
         "stream": app_settings.azure_openai.stream,
-        "model": app_settings.azure_openai.model,
-        "user": user_json
+        "model": app_settings.azure_openai.model
     }
 
     if len(messages) > 0:
@@ -340,6 +339,10 @@ def prepare_model_args(request_body, request_headers):
                         "embedding_dependency"
                     ]["authentication"][field] = "*****"
 
+    if model_args.get("extra_body") is None:
+        model_args["extra_body"] = {}
+    if user_security_context:  # security component introduced here https://learn.microsoft.com/en-us/azure/defender-for-cloud/gain-end-user-context-ai     
+                model_args["extra_body"]["user_security_context"]= user_security_context.to_dict()
     logging.debug(f"REQUEST BODY: {json.dumps(model_args_clean, indent=4)}")
 
     return model_args
@@ -1057,4 +1060,4 @@ async def generate_title(conversation_messages) -> str:
         return messages[-2]["content"]
 
 
-app = create_app()
+app = create_app()

backend/security/ms_defender_utils.py

@@ -1,15 +1,21 @@
-import json
+from typing import Dict, Any
+from dataclasses import dataclass, asdict, field
+import os
 
-def get_msdefender_user_json(authenticated_user_details, request_headers, conversation_id, application_name):
-    auth_provider = authenticated_user_details.get('auth_provider')
+
+@dataclass
+class UserSecurityContext:
+    application_name: str = field(default=None)    
+    end_user_id: str = field(default=None)
+    end_user_tenant_id: str = field(default=None)
+    source_ip: str = field(default=None)
+    def to_dict(self) -> Dict[str, Any]:
+        return {k: v for k, v in asdict(self).items() if v is not None}
+            
+
+def get_msdefender_user_json(authenticated_user_details, request_headers, application_name) -> UserSecurityContext:
     source_ip = request_headers.get('Remote-Addr', '')
-    header_names = ['User-Agent', 'X-Forwarded-For', 'Forwarded', 'X-Real-IP', 'True-Client-IP', 'CF-Connecting-IP']
-    user_args = {
-        "EndUserId": authenticated_user_details.get('user_principal_id'),
-        "EndUserIdType": "EntraId" if auth_provider == "aad" else auth_provider,
-        "SourceIp": source_ip.split(':')[0], #remove port
-        "SourceRequestHeaders": {header: request_headers[header] for header in header_names if header in request_headers},
-        "ConversationId": conversation_id,
-        "ApplicationName": application_name,
-    }
-    return json.dumps(user_args)
+    end_user_id = authenticated_user_details.get('user_principal_id')
+    source_ip= source_ip.split(':')[0]
+    return UserSecurityContext(end_user_id=end_user_id, source_ip=source_ip, application_name=application_name, end_user_tenant_id=None)
+