adding in support for JWT tokens to openai backends

aisentry/facade/app.py

@@ -3,6 +3,7 @@
 import uuid
 from datetime import datetime
 from dapr.clients import DaprClient
+from azure.identity import DefaultAzureCredential
 import httpcore
 from enum import Enum
 from typing import Tuple
@@ -35,6 +36,11 @@
 load_dotenv(".env", override=True)
 
 
+# os.environ["AZURE_CLIENT_ID"] = "your_client_id"
+# os.environ["AZURE_TENANT_ID"] = "your_tenant_id"
+# os.environ["AZURE_FEDERATED_TOKEN_FILE"] = "/var/run/secrets/tokens/azure-identity-token"
+
+
 logger.info("Starting Ai-Sentry Facade app")
 app = Quart(__name__)
 
@@ -130,9 +136,18 @@ async def catch_all(path):
                 logger.info("detected use of api-key header - will use this for authentication")
                 logger.debug(f"Swapping out api-key inside header with {endpoint_info['api-key']} value")
                 openAI_request_headers['Api-Key'] = endpoint_info['api-key']
-            else:
+
+            if endpoint_info['api-key'] is not None:
                 logger.info("No api-key header detected - will use the default api-key for authentication")
                 openAI_request_headers['Api-Key'] = endpoint_info['api-key']
+
+            else:
+                logger.info("No api-key config detected - will use oAuth to talk to openAI backend services.")
+                #Get Access Token from workload identity
+                credential = DefaultAzureCredential()
+                token = credential.get_token("https://cognitiveservices.azure.com/.default")
+                openAI_request_headers['Authorization'] = f"Bearer {token.token}"
+
             json_body = json.loads(body)
 
             if 'messages' in json_body:

content/documentation/Workload-identity-config.md

@@ -0,0 +1,54 @@
+# AKS Workload Identity setup
+
+## MI creation
+```powershell
+az account set --subscription "subscriptionID"
+```
+
+```powershell
+az identity create --name "ai-sentry-be-mi" --resource-group "ai-sentry" --location "australiaeast" --subscription "879bb272-07db-4784-816a-a9fac90f49ae"
+```
+
+```bash
+export USER_ASSIGNED_CLIENT_ID="$(az identity show --resource-group "ai-sentry" --name "ai-sentry-be-mi" --query 'clientId' -otsv)"
+```
+## Grant MI access to openAI resources
+
+![alt text](..\images\openai_rbac.png)
+
+and assign your newly built managed identity to above role:
+
+![alt text](..\images\openai_rbac2.png)
+
+
+## Env variables for service account in AKS
+
+```bash
+export SERVICE_ACCOUNT_NAME="default"
+export SERVICE_ACCOUNT_NAMESPACE="ai-sentry"
+```
+
+## OIDC Issuer url
+
+```bash
+export AKS_OIDC_ISSUER="$(az aks show --name anevjes-aks --resource-group aks --query "oidcIssuerProfile.issuerUrl" -o tsv)"
+```
+
+## Create AKS Service Account
+
+```bash
+cat <<EOF | kubectl apply -f -
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  annotations:
+    azure.workload.identity/client-id: ${USER_ASSIGNED_CLIENT_ID}
+  name: ${SERVICE_ACCOUNT_NAME}
+  namespace: ${SERVICE_ACCOUNT_NAMESPACE}
+EOF
+```
+## Establish federated identity credential trust
+
+```powershell
+az identity federated-credential create --name ai-sentry-be-fed --identity-name ai-sentry-be-mi --resource-group ai-sentry --issuer ${AKS_OIDC_ISSUER} --subject system:serviceaccount:${SERVICE_ACCOUNT_NAMESPACE}:${SERVICE_ACCOUNT_NAME} --audience api://AzureADTokenExchange
+```

tests/http/adapter_test-ai-sentry.http

@@ -1,4 +1,4 @@
-POST http://4.200.49.22/openai/deployments/gpt-4o/chat/completions?api-version=2024-02-15-preview HTTP/1.1
+POST http://4.147.31.58/openai/deployments/gpt-4o/chat/completions?api-version=2024-02-15-preview HTTP/1.1
 ai-sentry-consumer: Product-car-review
 ai-sentry-log-level: PII_STRIPPING_ENABLED
 ai-sentry-backend-pool: pool1