Automated Red Teaming Scans of Databricks Mosaic AI Model Serving Endpoints Using Protect AI Recon

Introduction

We are thrilled to announce the integration of Protect AI’s Recon with Databricks Mosaic AI Model Serving endpoints, a groundbreaking step in securing enterprise LLM Application deployments. With this integration, enterprises can harness Recon’s advanced red-teaming capabilities to proactively identify vulnerabilities, enhance LLM application integrity and ensure compliance with the latest AI governance standards.

This guide will walk you through the process of configuring and scheduling automated red teaming scans for Mosaic AI model serving endpoints using Recon.

1. Setting Up a Target in Recon

Step 1.1: Access the Targets Section

• Navigate to the Targets section in Recon.
• Click on the "New Target" button located in the top-right corner of the page.

Step 1.2: Provide Target Details

• Name the Target: Use a descriptive name (e.g., "Customer Support Chatbot API" or "Finetuned Llama Model v3.2").
• Select the Connection Method: Choose “Databricks”

• If the target is a public endpoint, proceed to the next step.
• If it is a private endpoint, Recon will provide a static IP address. Enable Access for the Recon-provided IP Address in the Databricks Workspace’s IP access list.

Step 1.3: Configure Secure Access

OAuth and Access Token

1. OAuth Configuration:
   • If your Databricks model endpoint uses OAuth, provide the required OAuth client ID, client secret during the target setup.
   • Recon will securely store these credentials and use them to obtain access tokens for each request.
2. Access Token:
   • If the model endpoint requires an access token for authentication:
     • Navigate to your Databricks workspace and generate a Personal Access Token:
       • In Databricks: Go to User Settings > Access Tokens > Generate New Token.
       • Provide an appropriate lifetime for the token.
     • Copy the generated token and paste it into the Authentication Token field in Recon.

Databricks Workspace Details

Workspace URL:

1. Enter your Databricks workspace URL (e.g., https://<workspace>.cloud.databricks.com) into the connection configuration.
2. Ensure that this URL is accessible from the Recon platform (static IP may need whitelisting for private endpoints).

Model Name

• Navigate to the URL of the serving endpoint in Databricks Mosaic AI Model Serving.
• Example endpoint: https://<workspace>.cloud.databricks.com/model/<model-name>/<version>/invocations.
• Input the model name and version (e.g. EmailResponderLLM v1.2) in the target details.
• These details help Recon identify and label results accurately in the reports.

Allowlisting Recon IP (For Private Endpoints)

• If the endpoint is private, allow the static IP provided by Recon in your Databricks Workspace’s IP access list. 

2. Configure Target Settings

Step 2.1: Configuring the Input Payload

NOTE: The provided defaults are designed to work seamlessly out of the box. Make changes only if necessary. 

Recon will replace {INPUT} with the scan prompts during execution. Default hyperparameters are pre-configured but can be modified if needed. For example:

{

  "top_p": 0.95,

  "stream": false,

  "messages": [

    {

      "role": "user",

      "content": "{INPUT}"

    }

  ],

  "max_tokens": 512,

  "temperature": 0.7

}

3. Advanced Configurations

Step 3.1: Rate Limits

When configuring an endpoint in Recon, you can specify rate limits of the target model to ensure that scanning does not exceed the endpoint’s capacity. 

Step 3.2: Guardrails / Content Filters

If your endpoint is an AI Gateway and/or has built-in content filters or safeguards to block unsafe or unethical responses, you can configure these details in Recon to test their effectiveness during scans.

4. Initiating a Red Teaming Scan

Step 4.1: Setting up a New Scan

• Navigate to the Scan Log section in Recon.
• Click on the "New Scan" button located in the top-right corner of the page.

Step 4.1: Choose a Scan Mode

Recon offers two modes:

• Attack Library Scan:
  • Uses a pre-built library of attack prompts.
  • Select attack categories like Safety, Jailbreak, or System Prompt Leaks.
• Agent Scan:
  • A dynamic, LLM-powered agent generates customized attack prompts based on responses.
  • Choose between automated or human-augmented modes for deeper testing.

Step 4.2: Start the Scan

• Navigate to the Scan section in Recon.
• Select the target you just configured.
• Pick the scan mode and configure any additional settings like severity thresholds or attack objectives.
• Click Start Scan.

5. Viewing and Analyzing Scan Results

NOTE: A scan can take from minutes to hours to complete depending on the type of scan, complexity & latency of your application. 

Step 5.1: Access Scan Reports

• Go to the Scan Logs section in Recon.
• Open the report for the completed scan.

Step 5.2: Review Key Metrics

• Risk Score: Assesses the overall vulnerability (0–100 scale).
• Attack Categories: Displays success rates across categories like Safety or Jailbreak.
• Attack Details: View prompts and compromised responses. Use the "View Details" option to examine conversation trees or payloads.

Step 5.3: Alignment with DASF2.0

Recon supports mapping vulnerabilities to the Databricks AI Security Framework (DASF) 2.0 [Add Link to DASF site]  alongside OWASP, MITRE ATLAS & NIST AI-RMF for comprehensive risk analysis and meeting enterprise standards.

6. Export or Share Findings

• Direct Download: 

• Easily share the report with stakeholders by exporting it. Click on icon in the top right corner

• Recon API:

• Utilize the Recon API to export the scan results and build detailed dashboards via Databricks Notebooks and/or conduct further in-depth analysis.

Reach out to your Protect AI Sales team for more information and guidance about Recon & Databricks integration.