Fix workload identity, bump litellm version

Author: ryansingman

deploy/helm/tlm/Chart.yaml

@@ -15,7 +15,7 @@ type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
 # Versions are expected to follow Semantic Versioning (https://semver.org/)
-version: 0.1.40
+version: 0.1.41
 
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application. Versions are not expected to

infra/tlm/environments/staging/main.tf

@@ -18,6 +18,7 @@ locals {
 module "tlm" {
     source = "../../modules/tlm"
     environment = local.environment
+    entity = "internal"
     location = local.location
 
     create_openai_service = true
@@ -40,6 +41,9 @@ module "tlm" {
         }
     }
 
+    default_completion_model = "gpt-4o-mini"
+    default_embedding_model = "text-embedding-3-small"
+
     tags = local.tags
 }
 

infra/tlm/modules/tlm/acr_image_pull/app_registration.tf

@@ -1,12 +1,23 @@
-# make sure to export the app ID and tenant ID
+data "azuread_application_published_app_ids" "well_known" {}
 
-# password will be stored in image pull secret in cluster
+data "azuread_service_principal" "msgraph" {
+  client_id = data.azuread_application_published_app_ids.well_known.result["MicrosoftGraph"]
+}
 
 resource azuread_application_registration "this" {
-    display_name = "TLM Cross-organization ACR Image Pull"
+    display_name = "${var.entity} - TLM Cross-organization ACR Image Pull"
     sign_in_audience = "AzureADMultipleOrgs"
 }
 
+resource azuread_application_api_access "this" {
+    application_id = azuread_application_registration.this.id
+    api_client_id = data.azuread_application_published_app_ids.well_known.result["MicrosoftGraph"]
+
+    role_ids = [
+        data.azuread_service_principal.msgraph.app_role_ids["Application.Read.All"]
+    ]
+}
+
 resource azuread_application_password "this" {
     application_id = azuread_application_registration.this.id
 }

infra/tlm/modules/tlm/acr_image_pull/variables.tf

@@ -0,0 +1,4 @@
+variable "entity" {
+    type = string
+    description = "The name of the entity deploying TLM (used to set display name for the cross-organization ACR image pull app registration)"
+}

infra/tlm/modules/tlm/app/main.tf

@@ -31,4 +31,14 @@ resource helm_release "this" {
         name = "chat_backend.defaults.TLM_DEFAULT_MODEL_API_BASE"
         value = data.azurerm_cognitive_account.openai_service.endpoint
     }
+
+    set {
+        name = "chat_backend.defaults.TLM_DEFAULT_COMPLETION_MODEL"
+        value = var.default_completion_model
+    }
+
+    set {
+        name = "chat_backend.defaults.TLM_DEFAULT_EMBEDDING_MODEL"
+        value = var.default_embedding_model
+    }
 }

infra/tlm/modules/tlm/app/openai_identity_sa.tf

@@ -27,6 +27,10 @@ resource kubernetes_service_account "openai_identity_sa" {
         namespace = local.namespace
         annotations = {
             "azure.workload.identity/client-id" = azurerm_user_assigned_identity.openai_identity_sa.client_id
+            "azure.workload.identity/service-account-token-expiration" = "86400"
+        }
+        labels = {
+            "azure.workload.identity/use" = "true"
         }
     }
 

infra/tlm/modules/tlm/app/variables.tf

@@ -50,3 +50,13 @@ variable "image_pull_password" {
     description = "The password to pull images from the registry"
     sensitive = true
 }
+
+variable "default_completion_model" {
+    type = string
+    description = "The default completion model to use"
+}
+
+variable "default_embedding_model" {
+    type = string
+    description = "The default embedding model to use"
+}

infra/tlm/modules/tlm/main.tf

@@ -5,6 +5,12 @@ resource "azurerm_resource_group" "this" {
     tags = var.tags
 }
 
+module "acr_image_pull" {
+    source = "./acr_image_pull"
+
+    entity = var.entity
+}
+
 locals {
     openai_service_resource_group_name = var.create_openai_service ? azurerm_resource_group.this.name : var.openai_service_resource_group_name
 }
@@ -49,10 +55,6 @@ provider "helm" {
   }
 }
 
-module "acr_image_pull" {
-    source = "./acr_image_pull"
-}
-
 
 module "app" {
     source = "./app"
@@ -68,10 +70,19 @@ module "app" {
     image_pull_username = module.acr_image_pull.app_reg_client_id
     image_pull_password = module.acr_image_pull.app_reg_password
 
+    default_completion_model = var.default_completion_model
+    default_embedding_model = var.default_embedding_model
+
     tags = var.tags
 
     providers = {
         kubernetes = kubernetes
         helm = helm
     }
+
+    depends_on = [
+        module.acr_image_pull,
+        module.openai_service,
+        module.cluster
+    ]
 }

infra/tlm/modules/tlm/variables.tf

@@ -3,6 +3,11 @@ variable "environment" {
     description = "The environment to deploy the TLM to"
 }
 
+variable "entity" {
+    type = string
+    description = "The name of the entity deploying TLM (used to set display name for the cross-organization ACR image pull app registration)"
+}
+
 variable "location" {
     type = string
     description = "The location to deploy the TLM to"
@@ -43,3 +48,13 @@ variable "openai_deployments" {
     description = "The cognitive deployments to create (only set if create_openai_service is true)"
     default = {}
 }
+
+variable "default_completion_model" {
+    type = string
+    description = "The default completion model to use. Must be deployed in the OpenAI service."
+}
+
+variable "default_embedding_model" {
+    type = string
+    description = "The default embedding model to use. Must be deployed in the OpenAI service."
+}

services/chat-backend/pyproject.toml

@@ -8,6 +8,7 @@ dependencies = [
     "azure-identity>=1.19.0",
     "fastapi>=0.115.5",
     "gunicorn>=23.0.0",
+    "litellm==1.63.8",
     "openai>=1.55.1",
     "pydantic-settings>=2.6.1",
     "python-json-logger>=3.2.1",