Usage Examples

import os 
from openai import AsyncOpenAI, BadRequestError

API_KEY=os.getenv("AZURE_OPENAI_KEY")
ENDPOINT = "https://<YOUR AZURE OPEN AI INSTANCE>.openai.azure.com/openai/v1"
TARGET_MODEL = "<YOUR OPEN AI MODEL NAME>" # e.g. "gpt-4o-mini"

# Define handler function that calls the target model
async def handler(prompt, history, session_id, target_system_prompt):

    openai_client = AsyncOpenAI(
        base_url=ENDPOINT,
        api_key=API_KEY
    )
    
    # Build conversation messages
    messages = [{"role": "system", "content": YOUR_SYSTEM_PROMPT}]
    messages.extend(history)  # Include conversation history for multi-turn attacks
    messages.append({"role": "user", "content": prompt})

    try: 
        # Call the target model
        oai_response = await openai_client.chat.completions.create(
            messages=messages,
            model=TARGET_MODEL,
        )
        target_response = oai_response.choices[0].message.content
    
    except BadRequestError as e: 
        print(e)
        print(e.message)
        
        target_response = e.message

    # Optional: Print progress
    print(f"[{session_id[:8]}] Processed attack prompt")

    return target_response

# Browser configuration
CHAT_URL = "https://chat.example.com"
NEW_SESSION_BUTTON_SELECTOR = "button.new-chat"
INPUT_SELECTOR = "textarea#chat-input"
OUTPUT_SELECTOR = "div.message.assistant"
MESSAGE_CONTENT_SELECTOR = ".message-content"
SUBMIT_METHOD = "enter"
SUBMIT_BUTTON_SELECTOR = "#send-btn"
HEADLESS = False

# Target configuration
TARGET_MODEL_NAME = "browser-ui-model"
TARGET_SYSTEM_PROMPT = """You are a helpful AI assistant."""

# Initialize browser
browser_context = {"playwright": None, "browser": None}

browser_context["playwright"] = await async_playwright().start()
browser_context["browser"] = await browser_context["playwright"].chromium.launch(headless=HEADLESS)

# Handler with multi-tab session management
session_pages = {}
session_message_counts = {}


async def handler(prompt, history, session_id, target_system_prompt):
    # Create tab for new sessions
    if session_id not in session_pages:
        page = await browser_context["browser"].new_page()
        await page.goto(CHAT_URL)

        if NEW_SESSION_BUTTON_SELECTOR:
            await page.wait_for_selector(NEW_SESSION_BUTTON_SELECTOR, timeout=10000)
            await page.click(NEW_SESSION_BUTTON_SELECTOR)
            await asyncio.sleep(0.5)

        await page.wait_for_selector(INPUT_SELECTOR, timeout=10000)
        session_pages[session_id] = page
        session_message_counts[session_id] = 0

    # Use session-specific tab
    page = session_pages[session_id]
    initial_count = session_message_counts[session_id]

    # Interact with page
    await page.fill(INPUT_SELECTOR, prompt)

    if SUBMIT_METHOD == "enter":
        await page.press(INPUT_SELECTOR, "Enter")
    else:
        await page.click(SUBMIT_BUTTON_SELECTOR)

    # Poll indefinitely for new message
    expected_count = initial_count + 1
    while True:
        messages = await page.query_selector_all(OUTPUT_SELECTOR)
        if len(messages) >= expected_count:
            last_message = messages[-1]

            if MESSAGE_CONTENT_SELECTOR:
                content_el = await last_message.query_selector(MESSAGE_CONTENT_SELECTOR)
                if content_el:
                    response_text = await content_el.inner_text()
                else:
                    response_text = await last_message.inner_text()
            else:
                response_text = await last_message.inner_text()

            if response_text:
                session_message_counts[session_id] = len(messages)
                return response_text

        await asyncio.sleep(1)

EVAL_NAME = "<NAME OF THE EVALUATION>"
SESSIONS_PER_TECHNIQUE = 2 # number of attempts to make per technique, = to "Sessions per Technique" in the console UI
MAX_TURNS = 5 # max number of turns for the attacker to achieve its objective, = to "Attacker Max Turns..." in the console UI
EXECUTION_STRATEGY_TYPE = "random" # could also be "single" or "static_prompt_set", = to "Execution Strategy" in the console UI
N_RANDOM_TECHNIQUES = 2  # Mix up to n (here: 2) additional random techniques per session for diversity, = to "Random Technique Count" in the console UI

session = await client.evaluation_sessions.red_team.start_session(
    name=EVAL_NAME,
    target_model=TARGET_MODEL_NAME,
    target_system_prompt=TARGET_SYSTEM_PROMPT,
    sessions_per_technique=SESSIONS_PER_TECHNIQUE,
    max_turns=MAX_TURNS,
    execution_strategy_type=EXECUTION_STRATEGY_TYPE,
    n_random_techniques=N_RANDOM_TECHNIQUES
)

# Don't forget to start the session with `run_with_callback` or `run_with_callback_parallel`! 
_ = await session.run_with_callback(handler=handler)