Thursday, September 28, 2023

Best practices for OpenAI Chat apps: Go Keyless

As part of my role the Python advocacy team for Azure, I am now one of the maintainers on several ChatGPT samples, like my simple chat app and this popular chat + search app. In this series of blog posts, I'll share my learnings for writing chat-like applications. My experience is from apps with Python backends, but many of these practices apply cross-language.

Today's tip for OpenAI apps isn't really specific to OpenAI, but is a good practice for production-grade apps of any type: don't use API keys! If your app is using openai.com's OpenAI service, then you'll have to use keys, but if you're using Azure's OpenAI service, then you can authenticate with Azure Active Directory tokens instead.


The risks of keys

It's tempting to use keys, since the setup looks so straightforward - you only need your endpoint URL and key.

openai.api_type = "azure"
openai.api_version = "2023-05-15"
openai.api_base = os.getenv("AZURE_OPENAI_ENDPOINT") 
openai.api_key = os.getenv("AZURE_OPENAI_KEY")

But using API keys in a codebase can lead to all kinds of issues. To name a few:

  • The key could be accidentally checked into a source control, by a developer who replaces the getenv() call with a hardcoded string, or a developer who adds a .env file to a commit.
  • Once checked into source control, keys are exposed internally and are also at a greater risk of external exposure by malicious actors who gain access to the codebase.
  • In a large company, multiple developers might unknowingly use the same key, use up each other's resources, and discover their services are failing due to quota errors.

I've seen all of these situations play out, and I don't want them to happen to other developers. A more secure approach is to use authentication tokens, and that's what I use in my samples.

Authenticating to Azure OpenAI with Active Directory

This code authenticates to Azure OpenAI with the openai Python package and Azure Python SDK:

openai.api_type = "azure_ad"
openai.api_version = "2023-05-15"
openai.api_base = os.getenv("AZURE_OPENAI_ENDPOINT")
azure_credential = DefaultAzureCredential()
openai.api_key = azure_credential.get_token(
    "https://cognitiveservices.azure.com/.default")

The differences:

  • The api_type is set to "azure_ad" so that the openai package knows to send the headers with the Bearer Token set properly.
  • The code authenticates to Azure using DefaultAzureCredential which will iterate through many possible credential types until it finds a valid Azure login.
  • The code then gets a token from that credential and sets that as the api_key.

Accessing OpenAI locally

The next step is to make sure that whoever is running the code has permission to access the OpenAI service. By default, you will not, even if you created the OpenAI service yourself. That's a security measure to make sure you don't accidentally access production resources from a local machine (particularly helpful when your code deals with write operations on databases.).

To access an OpenAI resource, you need the "Cognitive Services OpenAI User" role (role ID '5e0bd9bd-7b93-4f28-af87-19fc36ad61bd'). That can be assigned using the Azure Portal, Azure CLI, or ARM/Bicep.

Assigning roles with the Azure CLI

First, set the following environment variables:

  • PRINCIPAL_ID: The principal ID of your logged in account.
  • SUBSCRIPTION_ID: The subscription ID of your logged in account.
  • RESOURCE_GROUP: The resource group of the OpenAI resource.

Then you can run this command using the Azure CLI:

az role assignment create \
        --role "5e0bd9bd-7b93-4f28-af87-19fc36ad61bd" \
        --assignee-object-id "$PRINCIPAL_ID" \
        --scope /subscriptions/"$SUBSCRIPTION_ID"/resourceGroups/"$RESOURCE_GROUP" \
        --assignee-principal-type User

Assigning roles with ARM/Bicep

We use the Azure Developer CLI to deploy all of our samples, which relies on Bicep files to declare the infrastructure-as-code. That results in more repeatable deploys, so it's a great approach for deploying production applications.

This Bicep resource creates the role, assuming a principalId parameter is set:

resource role 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
  name: guid(subscription().id, resourceGroup().id,
             principalId, roleDefinitionId)
  properties: {
    principalId: principalId
    principalType: 'User'
    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions',
                                 '5e0bd9bd-7b93-4f28-af87-19fc36ad61bd')
  }
}

You can also see how our sample's main.bicep uses a module to set up the role.

Assigning roles with the Azure Portal

If you are unable to use those automated approaches (preferred), it's also possible to use the Azure Portal to create the role:

  • Open the OpenAI resource
  • Select "Access Control (IAM)" from the left navigation
  • Select "+ Add" in the top menu
  • Search for "Cognitive Services User" and select it in the results
  • Select "Assign access to: User, group, or service principal"
  • Search for your email address
  • Select "Review and assign"

Accessing OpenAI from production hosts

The next step is to ensure your deployed application can also use a DefaultAzureCredential token to access the OpenAI resource. That requires setting up a Managed Identity and assigning that same role to the Managed identity. There are two kinds of managed identities: system-assigned and user-assigned. All Azure hosting platforms support managed identity. We'll start with App Service and system-assigned identities as an example.

Managed identity for App Service

This is how we create an App Service with a system-assigned identity in Bicep code:

resource appService 'Microsoft.Web/sites@2022-03-01' = {
  name: name
  location: location
  identity: { type: 'SystemAssigned'}
  ...
}

For more details, see this article on Managed Identity for App Service.

Assigning roles to the managed identity

The role assignment process is largely the same for the host as it was for a user, but the principal ID must be set to the managed identity's principal ID instead and the principal type is "ServicePrincipal".

For example, this Bicep assigns the role for an App Service system-assigned identity:

resource role 'Microsoft.Authorization/roleAssignments@2022-04-01' = {
  name: guid(subscription().id, resourceGroup().id,
             principalId, roleDefinitionId)
  properties: {
    principalId: appService.identity.principalId
    principalType: 'ServicePrincipal'
    roleDefinitionId: resourceId('Microsoft.Authorization/roleDefinitions',
                                 '5e0bd9bd-7b93-4f28-af87-19fc36ad61bd')
  }
}

User-assigned identity for Azure Container Apps

It's also possible to use a system-assigned identity for Azure Container Apps, using a similar approach as above. However, for our samples, we needed to use user-assigned identities so that we could give the same identity access to Azure Container Registry before the ACA app was provisioned. That's the advantage of a user-assigned identities, reuse across multiple resources.

First, we create a new identity outside of the ACA Bicep:

resource userIdentity 'Microsoft.ManagedIdentity/userAssignedIdentities@2023-01-31' = {
  name: '${prefix}-id-aca'
  location: location
}

Then we assign that identity to the ACA resource:

resource app 'Microsoft.App/containerApps@2022-03-01' = {
  name: name
  location: location
  identity: {
    type: 'UserAssigned'
    userAssignedIdentities: { '${userIdentity.id}': {} }
  }
  ...

When using a user-assigned identity, we need to modify our call to AzureDefaultCredential to tell it which identity to use, since you could potentially have multiple user-assigned identities (not just the single system-assigned identity for the hosting environment).

The following code retrieves the identity's ID from the environment variables and specifies it as the client_id for the Managed Identity credential:

default_credential = azure.identity.aio.ManagedIdentityCredential(
    client_id=os.getenv("AZURE_OPENAI_CLIENT_ID"))

Refreshing expired authentication tokens

The credentials returned from Azure AD do not last forever, so for any long running script or hosted application, you will need to refresh the tokens. Typically, the Azure Python SDK takes care of that for you, but since we use the openai package for Python apps, we need to implement token refresh ourselves.

For our application that uses the Quart web framework, we define a function that runs before every request to check if the globally stored token is close to expiring. If so, we fetch a new token and store it.

@bp.before_request
async def ensure_openai_token():
    if openai.api_type != "azure_ad":
        return
    openai_token = current_app.config[CONFIG_OPENAI_TOKEN]
    if openai_token.expires_on < time.time() + 60:
        openai_token = await current_app.config[CONFIG_CREDENTIAL].get_token(
            "https://cognitiveservices.azure.com/.default"
        )
        current_app.config[CONFIG_OPENAI_TOKEN] = openai_token
        openai.api_key = openai_token.token

For our script that ingests data from PDFs, we define a similar function that we call before every attempt to use the vector embedding function:

def refresh_openai_token():
    if (
        CACHE_KEY_TOKEN_TYPE in open_ai_token_cache
        and open_ai_token_cache[CACHE_KEY_TOKEN_TYPE] == "azure_ad"
        and open_ai_token_cache[CACHE_KEY_CREATED_TIME] + 300 < time.time()
    ):
        token_cred = open_ai_token_cache[CACHE_KEY_TOKEN_CRED]
        openai.api_key = token_cred.get_token(
            "https://cognitiveservices.azure.com/.default").token
        open_ai_token_cache[CACHE_KEY_CREATED_TIME] = time.time()

Accessing OpenAI in a local Docker container

At this point, you should be able to access OpenAI both for local development and in production. Unless, that is, you're developing with a local Docker container. By default, a Docker container does not have a way to access any of your local credentials, so you'll see authentication errors in the logs. It used to be possible to use a workaround with volumes to access the credential, but after Azure started encrypting the local credential, it's now an open question as to how to easily authenticate inside a local container.

Unfortunately, in this case, my current approach is to fallback to using a key for local development in a Docker container. Another interesting approach would be to use a mock ChatGPT service in a local environment, to avoid unnecessarily using up quota.


All together now

As you can see, it's not entirely straightforward to authenticate to OpenAI without keys, depending on how you're developing locally and where you're deploying.

The following code uses a key when it's set in the environment, uses a user-assigned Managed Identity when the identity ID is set in the environment, and otherwises uses DefaultAzureCredential:

openai.api_base = os.getenv("AZURE_OPENAI_ENDPOINT")
openai.api_version = "2023-03-15-preview"
if os.getenv("AZURE_OPENAI_KEY"):
    openai.api_type = "azure"
    openai.api_key = os.getenv("AZURE_OPENAI_KEY")
else:
    openai.api_type = "azure_ad"
    if client_id := os.getenv("AZURE_OPENAI_CLIENT_ID"):
        default_cred = azure.identity.aio.ManagedIdentityCredential(
        client_id=client_id)
    else:
        default_cred = azure.identity.aio.DefaultAzureCredential(
            exclude_shared_token_cache_credential=True)
    token = await default_cred.get_token(
        "https://cognitiveservices.azure.com/.default")
    openai.api_key = token.token

The technologies in this space are changing rapidly, so some of the more tricky aspects of keyless authentication will hopefully be easier in the future. In the meantime, try to avoid keys whenever possible.

Posted by
Labels: , ,

No comments:

Post a Comment

Subscribe to: Post Comments (Atom)

.