From fc1a681b71a041aa78c8a39449bbfcac99d2926e Mon Sep 17 00:00:00 2001
From: Pamela Fox <pamela.fox@gmail.com>
Date: Mon, 3 Jun 2024 07:13:20 -0700
Subject: [PATCH] Use bicep and env vars for backends

---
 app/backend/app.py         | 22 +++++++++++++---------
 infra/main.bicep           | 20 ++++++++++++++++++++
 infra/main.parameters.json |  3 +++
 3 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/app/backend/app.py b/app/backend/app.py
index 6a66deb092..dc509eafa6 100644
--- a/app/backend/app.py
+++ b/app/backend/app.py
@@ -4,10 +4,10 @@
 import logging
 import mimetypes
 import os
-import httpx
 from pathlib import Path
-from typing import Any, AsyncGenerator, Dict, Union, cast, List
+from typing import Any, AsyncGenerator, Dict, List, Union, cast
 
+import httpx
 from azure.core.exceptions import ResourceNotFoundError
 from azure.identity.aio import DefaultAzureCredential, get_bearer_token_provider
 from azure.monitor.opentelemetry import configure_azure_monitor
@@ -76,11 +76,6 @@
 mimetypes.add_type("application/javascript", ".js")
 mimetypes.add_type("text/css", ".css")
 
-# Data for the backends could be supplied through config. This data is simply here to illustrate usage.
-backends: List[Backend] = [
-    Backend("cog-w2og7ojyhvoq6.openai.azure.com", 1),
-    Backend("cog-kfdf7d5q443bu.openai.azure.com", 1),
-]
 
 @bp.route("/")
 async def index():
@@ -438,13 +433,22 @@ async def setup_clients():
 
         api_version = os.getenv("AZURE_OPENAI_API_VERSION") or "2024-03-01-preview"
 
-        lb = AsyncLoadBalancer(backends)
+        client_args = {}
+        if AZURE_OPENAI_SERVICE_BACKEND2 := os.environ.get("AZURE_OPENAI_SERVICE_BACKEND2"):
+            backends: List[Backend] = [
+                Backend(f"{AZURE_OPENAI_SERVICE}.openai.azure.com", 1),
+                Backend(f"{AZURE_OPENAI_SERVICE_BACKEND2}.openai.azure.com", 1),
+            ]
+
+            lb = AsyncLoadBalancer(backends)
+            # Inject the load balancer as the transport in a new default httpx client
+            client_args["http_client"] = httpx.AsyncClient(transport=lb)
 
         openai_client = AsyncAzureOpenAI(
             api_version=api_version,
             azure_endpoint=endpoint,
             azure_ad_token_provider=token_provider,
-            http_client = httpx.AsyncClient(transport = lb)        # Inject the load balancer as the transport in a new default httpx client
+            **client_args,
         )
     elif OPENAI_HOST == "local":
         openai_client = AsyncOpenAI(
diff --git a/infra/main.bicep b/infra/main.bicep
index 7ac59b7272..ae26235030 100644
--- a/infra/main.bicep
+++ b/infra/main.bicep
@@ -43,6 +43,7 @@ param appServiceSkuName string // Set in main.parameters.json
 @allowed([ 'azure', 'openai', 'azure_custom' ])
 param openAiHost string // Set in main.parameters.json
 param isAzureOpenAiHost bool = startsWith(openAiHost, 'azure')
+param deployAzureOpenAiBackendTwo bool = false
 param azureOpenAiCustomUrl string = ''
 param azureOpenAiApiVersion string = ''
 
@@ -294,6 +295,7 @@ module backend 'core/host/appservice.bicep' = {
       AZURE_OPENAI_GPT4V_MODEL: gpt4vModelName
       // Specific to Azure OpenAI
       AZURE_OPENAI_SERVICE: isAzureOpenAiHost ? openAi.outputs.name : ''
+      AZURE_OPENAI_SERVICE_BACKEND2: isAzureOpenAiHost && deployAzureOpenAiBackendTwo ? openAiBackendTwo.outputs.name : ''
       AZURE_OPENAI_CHATGPT_DEPLOYMENT: chatGpt.deploymentName
       AZURE_OPENAI_EMB_DEPLOYMENT: embedding.deploymentName
       AZURE_OPENAI_GPT4V_DEPLOYMENT: useGPT4V ? gpt4vDeploymentName : ''
@@ -385,6 +387,23 @@ module openAi 'core/ai/cognitiveservices.bicep' = if (isAzureOpenAiHost) {
   }
 }
 
+module openAiBackendTwo 'core/ai/cognitiveservices.bicep' = if (isAzureOpenAiHost && deployAzureOpenAiBackendTwo) {
+  name: 'openai-backend-two'
+  scope: openAiResourceGroup
+  params: {
+    name: '${abbrs.cognitiveServicesAccounts}${resourceToken}-b2'
+    location: openAiResourceGroupLocation
+    tags: tags
+    publicNetworkAccess: publicNetworkAccess
+    bypass: bypass
+    sku: {
+      name: openAiSkuName
+    }
+    deployments: openAiDeployments
+    disableLocalAuth: true
+  }
+}
+
 // Formerly known as Form Recognizer
 // Does not support bypass
 module documentIntelligence 'core/ai/cognitiveservices.bicep' = {
@@ -766,6 +785,7 @@ output AZURE_OPENAI_GPT4V_MODEL string = gpt4vModelName
 
 // Specific to Azure OpenAI
 output AZURE_OPENAI_SERVICE string = isAzureOpenAiHost ? openAi.outputs.name : ''
+output AZURE_OPENAI_SERVICE_BACKEND2 string = isAzureOpenAiHost && deployAzureOpenAiBackendTwo ? openAiBackendTwo.outputs.name : ''
 output AZURE_OPENAI_API_VERSION string = isAzureOpenAiHost ? azureOpenAiApiVersion : ''
 output AZURE_OPENAI_RESOURCE_GROUP string = isAzureOpenAiHost ? openAiResourceGroup.name : ''
 output AZURE_OPENAI_CHATGPT_DEPLOYMENT string = isAzureOpenAiHost ? chatGpt.deploymentName : ''
diff --git a/infra/main.parameters.json b/infra/main.parameters.json
index bfaa0deff9..9b87e68f29 100644
--- a/infra/main.parameters.json
+++ b/infra/main.parameters.json
@@ -119,6 +119,9 @@
     "openAiHost": {
       "value": "${OPENAI_HOST=azure}"
     },
+    "deployAzureOpenAiBackendTwo": {
+      "value": "${DEPLOY_AZURE_OPENAI_BACKEND_TWO=false}"
+    },
     "azureOpenAiCustomUrl":{
       "value": "${AZURE_OPENAI_CUSTOM_URL}"
     },