Introduction
Organizations are quickly adopting Gen AI applications because they help save money and offer a competitive edge. These AI tools can automate many tasks that weren’t possible before large language models (LLMs) became popular. However, while building and launching Gen AI apps is now easier, keeping them safe and secure is still a challenge
In this blog, we’ll show you how simple it is to turn on and use AI Guardrails in Mosaic AI Gateway using an easy-to-use interface—no coding required! Without AI Guardrails, your Gen AI applications can easily generate harmful content. Apart from easy-to-use interface, you also got governance capabilities as well as it is LLM model agnostic. That means that you can use it for any LLM models as shown in the below diagram.
Mosaic AI Gateway Architecture
Let us assume that you have recently launched an AI chatbot for your organization’s customer services unit. If you have devised a mechanism to track every input prompt and output prompt, you may realize that your AI Chatbot received many input prompts from end users which are unsafe or harmful. Similarly, your AI chatbot may also have provided responses for those unsafe and harmful input prompts. Knowing this, you may immediately decide to stop AI chatbot services as it may do more harm than good to your organization until you figure out how to provide AI Guardrails to prevent the same.
Based on the use case, AI Guardrails may also help in complying with regulatory requirements related to GDPR, HIPPA, etc.
Databricks can provide AI Guardrails capabilities using Mosaic AI Gateway. The following diagram shows how AI Guardrail works with Mosaic AI Gateway in a simplified manner.
Let us understand the key terminologies before getting into how to use the same for your specific use case.
Key Terminologies
Mosaic AI :
Mosaic AI is the umbrella term for Databricks' suite of AI/ML solutions, including both Generative AI and classical Data Science and Machine Learning (DSML). It is a collaborative and data-native solution designed to maximize productivity across the entire machine learning lifecycle, from experimentation to production. Mosaic AI provides a comprehensive array of features for preparing, processing, and managing data in a self-service manner, and for operationalizing machine learning at scale.
Mosaic AI Gateway :
Mosaic AI Gateway is a centralized service that brings governance, monitoring, and production readiness to model serving endpoints. It streamlines the usage and management of Gen AI models and agents within an organization. It also allows you to run, secure, and govern AI traffic to democratize and accelerate AI adoption for your organization.
AI Guardrails :
AI Guardrails prevent unwanted data and unsafe data in requests and responses for Gen AI applications. Using AI Guardrails, you can configure and enforce data compliance at the model serving endpoint level and to reduce harmful content on any requests sent to the underlying model. The system blocks bad requests and responses and returns a default message to the user. Basically, it prevents the LLM from interacting with certain types of content.
Let us prepare the environment
Let us prepare the Databricks environment to use AI Guardrails with Mosaic AI Gateway. By default, Mosaic AI Gateway comes with your workspace. As it is a serverless offering, it is easy to use it in your workspace. In order to use Mosaic AI Gateway, you will have to enable it while creating the model serving endpoint along with enabling AI Guardrails on the same. You need to have access to an active Databricks workspace to enable the same..
If you don’t have access to an active Databricks workspace, you should be able to create a new one using Express Setup or via an existing cloud account. Please visit https://www.databricks.com/try-databricks for more information on this.
Step by step instructions
In order to use Mosaic AI Gateway for AI Guardrails, you will have to create a Serving Endpoint for your choice of LLM Model which you want to use for your Gen AI application and on which you want to apply AI Guardrails. Mosaic AI Model Serving supports both Foundation Model Serving Endpoint as well as External Model Serving Endpoint. You can select whichever endpoint is suitable for your requirements. If you are confused about how to create Model Serving Endpoint, then please look at the below documentation link for information on the same.
https://docs.databricks.com/aws/en/machine-learning/model-serving/create-foundation-model-endpoints
As part of the Mosaic AI Model Serving Endpoint creation process, you will get an option to enable/disable AI Guardrails. Apart from the GUI option to create AI Guardrails using Mosaic AI Gateway, you also have an option to do the same using API as well. However, for the sake of simplicity and ease, we will only focus on GUI options to achieve the same in this blog.
Step 1: Enable AI Guardrails
On your Serving Endpoint creation screen, under the AI Gateway section, you will see various features supported by AI Gateway as shown in the picture below. It is necessary to select the checkbox ☑️shown against the feature to enable that respective feature for your AI Gateway. Please select the checkbox shown against AI Guardrails.
If you want to monitor your model behaviour, then you can enable usage tracking as well as inference tables which will help you to keep a watch on how your model is behaving by logging all those details.
Once you select the checkbox ☑ for AI Guardrails, you will see multiple options to enable various different AI Guardrails from the list of available options as mentioned below:
A. Safely filter
B. PII Detection
Step 2: Enable Input/Output Safety filter
Once you select the checkbox ☑shown against AI Guardrails, you will see multiple options under Input and Output tabs. If you want to enable detecting and blocking unsafe or harmful content being used as input prompts to your Gen AI applications, including violent crime, self-harm (e.g. content related to suicide), hate speech, etc. then select the checkbox ☑shown against Input Filter. This will ensure that none of your users will be able to process any input prompts which are unsafe. Similarly, it may happen that the LLM Model you are using for your Gen AI application may end up generating some harmful or unwanted responses for some of the valid input prompts as well. To avoid that, you have an option to enable the Safety Filter for output prompt as well. In this case, irrespective of the input prompts, your Gen AI application will not generate any harmful or unwanted responses.
Step 3: Test or Validate Input/Output Safety filter
Once you are done with selecting appropriate AI Guardrails options from GUI, you can create the Mosaic AI Model Serving endpoint by clicking on the Create button. This will create and deploy your LLM Model Serving Endpoint in sometime and it will be ready to use. The easiest way to test or validate is that you can use your Model Serving Endpoint with the AI Playground GUI available in your Databricks workspace. Once your endpoint is deployed and ready for usage, you can click on Use → Try in Playground option as shown in the below screenshot to launchthe AI Playground GUI to test or validate your serving endpoint.
In the following screenshot, you can see that input prompt contains self harm related question and as the safety filter is enabled, it does not provide a response. Instead, it generates an internal error message and identifies the category of the safety filter.
It is recommended that you enable both Input Safety Filter as well as Output Safety Filter to safeguard your Gen AI applications completely. Please note that it is not mandatory to enable Safety Filter for both or either of Input or Output filters for enabling AI Guardrails using Mosaic AI Gateway as it provides multiple other options which we will discuss in the following steps now.
Currently, Databricks uses Llama Guard 2-8b (may change in future) as the safety filter. To learn more about the Llama Guard safety filter and what topics apply to the safety filter, see the Meta Llama Guard 2 8B model card.
Step 4: Enable Other AI Guardrails Options
By now, you would have realized how easy it is to enable AI Guardrails using Mosaic AI Gateway in Databricks. Apart from Input/Output Safety Filters which we have created and validated above, there are multiple other AI Guardrails options which you can enable and validate very easily using GUI without writing any single line of code. You can refer to the below chart which summarizes all the AI Guardrails options available to you.
Summing up
In conclusion, Mosaic AI Gateway provides a straightforward and effective method for implementing AI Guardrails on LLM Model Serving Endpoints. By enabling features such as safety filters and PII detection through a user-friendly GUI, organizations can ensure their Gen AI applications are secure, compliant, and free from harmful content. This empowers users to leverage the benefits of Gen AI while maintaining control and mitigating potential risks. There will be new features which will get added to AI Guardrails so please keep a watch on new announcements from Databricks on the same.
2 Comments
