Trace Red Teaming Agent (Microsoft Foundry)

# Setup up red teaming agent
import os

# Azure imports
from azure.identity import DefaultAzureCredential
from azure.ai.evaluation.red_team import RedTeam, RiskCategory, AttackStrategy

#Set up environment variables
os.environ["AZURE_SUBSCRIPTION_ID"] = ""
os.environ["AZURE_RESOURCE_GROUP"] = ""
os.environ["AZURE_PROJECT_NAME"] = ""
os.environ["PROJECT_ENDPOINT"] = ""
os.environ["ARIZE_SPACE_ID"] = ""
os.environ["ARIZE_API_KEY"] = ""
os.environ["PROJECT_NAME"] = "red-team-violence-examples"


## Using Azure AI Foundry Hub project
azure_ai_project = {
    "subscription_id": os.environ["AZURE_SUBSCRIPTION_ID"],
    "resource_group_name": os.environ["AZURE_RESOURCE_GROUP"],
    "project_name": os.environ["AZURE_PROJECT_NAME"],
}
azure_ai_project = os.environ["PROJECT_ENDPOINT"]

# Instantiate your AI Red Teaming Agent
red_team_agent = RedTeam(
    azure_ai_project=azure_ai_project, # required
    credential=DefaultAzureCredential() # required
)

# Specifying risk categories and number of attack objectives per risk categories you want the AI Red Teaming Agent to cover
red_team_agent = RedTeam(
    azure_ai_project=azure_ai_project, # required
    credential=DefaultAzureCredential(), # required
    risk_categories=[ # optional, defaults to all four risk categories
        RiskCategory.Violence,
        RiskCategory.HateUnfairness,
        RiskCategory.Sexual,
        RiskCategory.SelfHarm
    ],
    num_objectives=10, # optional, defaults to 10
)

#enable tracing for openai
from arize.otel import register

tracer_provider = register(
    space_id = os.environ["ARIZE_SPACE_ID"],
    api_key = os.environ["ARIZE_API_KEY"],
    project_name = os.environ["PROJECT_NAME"],
)

from openinference.instrumentation.openai import OpenAIInstrumentor
OpenAIInstrumentor().instrument(tracer_provider=tracer_provider)

# Set up a callback function to pass to the red teaming agent scan
import openai

# Define a simple callback function that simulates a chatbot
def simple_callback(query: str) -> str:
    # Insert your LLM or agent here
    openai_client = openai.OpenAI()
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system",
            "content": "You are a helpful AI assistant. Always maintain a polite and professional tone. Provide concise answers."
            },
            {"role": "user", "content": query}
        ],
        max_tokens=100,
    )

red_team_result = await red_team_agent.scan(target=simple_callback)