Skip to content

Commit ffa5859

Browse files
Copilotslister1001
andcommitted
Add language support to RedTeam with SupportedLanguages enum
Co-authored-by: slister1001 <[email protected]>
1 parent 95a992d commit ffa5859

File tree

4 files changed

+204
-2
lines changed

4 files changed

+204
-2
lines changed

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from ._attack_strategy import AttackStrategy
88
from ._attack_objective_generator import RiskCategory
99
from ._red_team_result import RedTeamResult
10+
from ..simulator._constants import SupportedLanguages
1011
except ImportError:
1112
raise ImportError(
1213
"Could not import Pyrit. Please install the dependency with `pip install azure-ai-evaluation[redteam]`."
@@ -18,4 +19,5 @@
1819
"AttackStrategy",
1920
"RiskCategory",
2021
"RedTeamResult",
22+
"SupportedLanguages",
2123
]

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_red_team.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
EVALUATION_PASS_FAIL_MAPPING,
3131
TokenScope,
3232
)
33+
from azure.ai.evaluation.simulator._constants import SupportedLanguages
3334
from azure.ai.evaluation._evaluate._utils import _get_ai_studio_url
3435
from azure.ai.evaluation._evaluate._utils import extract_workspace_triad_from_trace_provider
3536
from azure.ai.evaluation._version import VERSION
@@ -138,6 +139,8 @@ class RedTeam:
138139
:type application_scenario: Optional[str]
139140
:param custom_attack_seed_prompts: Path to a JSON file containing custom attack seed prompts (can be absolute or relative path)
140141
:type custom_attack_seed_prompts: Optional[str]
142+
:param language: Language to use for attack objectives generation. Defaults to English.
143+
:type language: SupportedLanguages
141144
:param output_dir: Directory to save output files (optional)
142145
:type output_dir: Optional[str]
143146
"""
@@ -239,6 +242,7 @@ def __init__(
239242
num_objectives: int = 10,
240243
application_scenario: Optional[str] = None,
241244
custom_attack_seed_prompts: Optional[str] = None,
245+
language: SupportedLanguages = SupportedLanguages.English,
242246
output_dir=".",
243247
):
244248
"""Initialize a new Red Team agent for AI model evaluation.
@@ -260,13 +264,16 @@ def __init__(
260264
:type application_scenario: Optional[str]
261265
:param custom_attack_seed_prompts: Path to a JSON file with custom attack prompts
262266
:type custom_attack_seed_prompts: Optional[str]
267+
:param language: Language to use for attack objectives generation. Defaults to English.
268+
:type language: SupportedLanguages
263269
:param output_dir: Directory to save evaluation outputs and logs. Defaults to current working directory.
264270
:type output_dir: str
265271
"""
266272

267273
self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
268274
self.credential = credential
269275
self.output_dir = output_dir
276+
self.language = language
270277
self._one_dp_project = is_onedp_project(azure_ai_project)
271278

272279
# Initialize logger without output directory (will be updated during scan)
@@ -743,7 +750,7 @@ async def get_jailbreak_prefixes_with_retry():
743750
# Use the RAI service to get attack objectives
744751
try:
745752
self.logger.debug(
746-
f"API call: get_attack_objectives({risk_cat_value}, app: {application_scenario}, strategy: {strategy})"
753+
f"API call: get_attack_objectives({risk_cat_value}, app: {application_scenario}, strategy: {strategy}, language: {self.language.value})"
747754
)
748755
# strategy param specifies whether to get a strategy-specific dataset from the RAI service
749756
# right now, only tense requires strategy-specific dataset
@@ -753,6 +760,7 @@ async def get_jailbreak_prefixes_with_retry():
753760
risk_category=other_risk,
754761
application_scenario=application_scenario or "",
755762
strategy="tense",
763+
language=self.language.value,
756764
scan_session_id=self.scan_session_id,
757765
)
758766
else:
@@ -761,6 +769,7 @@ async def get_jailbreak_prefixes_with_retry():
761769
risk_category=other_risk,
762770
application_scenario=application_scenario or "",
763771
strategy=None,
772+
language=self.language.value,
764773
scan_session_id=self.scan_session_id,
765774
)
766775
if isinstance(objectives_response, list):

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/simulator/_model_tools/_generated_rai_client.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ async def get_attack_objectives(
100100
risk_category: Optional[str] = None,
101101
application_scenario: str = None,
102102
strategy: Optional[str] = None,
103+
language: str = "en",
103104
scan_session_id: Optional[str] = None,
104105
) -> Dict:
105106
"""Get attack objectives using the auto-generated operations.
@@ -112,6 +113,8 @@ async def get_attack_objectives(
112113
:type application_scenario: str
113114
:param strategy: Optional strategy to filter the attack objectives
114115
:type strategy: Optional[str]
116+
:param language: Language code for the attack objectives (e.g., "en", "es", "fr")
117+
:type language: str
115118
:param scan_session_id: Optional unique session ID for the scan
116119
:type scan_session_id: Optional[str]
117120
:return: The attack objectives
@@ -122,7 +125,7 @@ async def get_attack_objectives(
122125
response = self._client.get_attack_objectives(
123126
risk_types=[risk_type],
124127
risk_category=risk_category,
125-
lang="en",
128+
lang=language,
126129
strategy=strategy,
127130
headers={"x-ms-client-request-id": scan_session_id},
128131
)
Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
import pytest
2+
from unittest.mock import AsyncMock, MagicMock, patch
3+
from azure.ai.evaluation.red_team._red_team import RedTeam, RiskCategory
4+
from azure.ai.evaluation.simulator._constants import SupportedLanguages
5+
from azure.core.credentials import TokenCredential
6+
7+
8+
@pytest.fixture
9+
def mock_azure_ai_project():
10+
return {
11+
"subscription_id": "test-subscription",
12+
"resource_group_name": "test-resource-group",
13+
"project_name": "test-project",
14+
}
15+
16+
17+
@pytest.fixture
18+
def mock_credential():
19+
return MagicMock(spec=TokenCredential)
20+
21+
22+
class TestRedTeamLanguageSupport:
23+
"""Test language support functionality in RedTeam class."""
24+
25+
def test_red_team_init_default_language(self, mock_azure_ai_project, mock_credential):
26+
"""Test that RedTeam initializes with default English language."""
27+
with patch("azure.ai.evaluation.red_team._red_team.GeneratedRAIClient"), \
28+
patch("azure.ai.evaluation.red_team._red_team.setup_logger") as mock_setup_logger, \
29+
patch("azure.ai.evaluation.red_team._red_team.initialize_pyrit"), \
30+
patch("azure.ai.evaluation.red_team._red_team._AttackObjectiveGenerator"):
31+
32+
mock_logger = MagicMock()
33+
mock_setup_logger.return_value = mock_logger
34+
35+
agent = RedTeam(
36+
azure_ai_project=mock_azure_ai_project,
37+
credential=mock_credential,
38+
risk_categories=[RiskCategory.Violence],
39+
num_objectives=5
40+
)
41+
42+
# Verify default language is English
43+
assert agent.language == SupportedLanguages.English
44+
45+
def test_red_team_init_custom_language(self, mock_azure_ai_project, mock_credential):
46+
"""Test that RedTeam initializes with custom language."""
47+
with patch("azure.ai.evaluation.red_team._red_team.GeneratedRAIClient"), \
48+
patch("azure.ai.evaluation.red_team._red_team.setup_logger") as mock_setup_logger, \
49+
patch("azure.ai.evaluation.red_team._red_team.initialize_pyrit"), \
50+
patch("azure.ai.evaluation.red_team._red_team._AttackObjectiveGenerator"):
51+
52+
mock_logger = MagicMock()
53+
mock_setup_logger.return_value = mock_logger
54+
55+
# Test with Spanish language
56+
agent = RedTeam(
57+
azure_ai_project=mock_azure_ai_project,
58+
credential=mock_credential,
59+
risk_categories=[RiskCategory.Violence],
60+
num_objectives=5,
61+
language=SupportedLanguages.Spanish
62+
)
63+
64+
assert agent.language == SupportedLanguages.Spanish
65+
66+
@pytest.mark.parametrize("language", [
67+
SupportedLanguages.English,
68+
SupportedLanguages.Spanish,
69+
SupportedLanguages.French,
70+
SupportedLanguages.German,
71+
SupportedLanguages.Italian,
72+
SupportedLanguages.Portuguese,
73+
SupportedLanguages.Japanese,
74+
SupportedLanguages.Korean,
75+
SupportedLanguages.SimplifiedChinese,
76+
])
77+
def test_red_team_init_all_supported_languages(self, mock_azure_ai_project, mock_credential, language):
78+
"""Test that RedTeam initializes correctly with all supported languages."""
79+
with patch("azure.ai.evaluation.red_team._red_team.GeneratedRAIClient"), \
80+
patch("azure.ai.evaluation.red_team._red_team.setup_logger") as mock_setup_logger, \
81+
patch("azure.ai.evaluation.red_team._red_team.initialize_pyrit"), \
82+
patch("azure.ai.evaluation.red_team._red_team._AttackObjectiveGenerator"):
83+
84+
mock_logger = MagicMock()
85+
mock_setup_logger.return_value = mock_logger
86+
87+
agent = RedTeam(
88+
azure_ai_project=mock_azure_ai_project,
89+
credential=mock_credential,
90+
risk_categories=[RiskCategory.Violence],
91+
num_objectives=5,
92+
language=language
93+
)
94+
95+
assert agent.language == language
96+
97+
@pytest.mark.asyncio
98+
async def test_get_attack_objectives_passes_language(self, mock_azure_ai_project, mock_credential):
99+
"""Test that _get_attack_objectives passes language parameter to generated RAI client."""
100+
with patch("azure.ai.evaluation.red_team._red_team.GeneratedRAIClient") as mock_rai_client_class, \
101+
patch("azure.ai.evaluation.red_team._red_team.setup_logger") as mock_setup_logger, \
102+
patch("azure.ai.evaluation.red_team._red_team.initialize_pyrit"), \
103+
patch("azure.ai.evaluation.red_team._red_team._AttackObjectiveGenerator"):
104+
105+
mock_logger = MagicMock()
106+
mock_setup_logger.return_value = mock_logger
107+
108+
# Set up mock RAI client instance
109+
mock_rai_client = MagicMock()
110+
mock_rai_client.get_attack_objectives = AsyncMock(return_value=[
111+
{
112+
"id": "test-id",
113+
"messages": [{"role": "user", "content": "test prompt"}],
114+
"metadata": {"target_harms": [{"risk-type": "violence"}]}
115+
}
116+
])
117+
mock_rai_client_class.return_value = mock_rai_client
118+
119+
# Create RedTeam instance with Spanish language
120+
agent = RedTeam(
121+
azure_ai_project=mock_azure_ai_project,
122+
credential=mock_credential,
123+
risk_categories=[RiskCategory.Violence],
124+
num_objectives=5,
125+
language=SupportedLanguages.Spanish
126+
)
127+
128+
agent.generated_rai_client = mock_rai_client
129+
agent.scan_session_id = "test-session"
130+
131+
# Call _get_attack_objectives
132+
await agent._get_attack_objectives(
133+
risk_category=RiskCategory.Violence,
134+
application_scenario="test scenario",
135+
strategy="baseline"
136+
)
137+
138+
# Verify that get_attack_objectives was called with Spanish language
139+
mock_rai_client.get_attack_objectives.assert_called_once()
140+
call_args = mock_rai_client.get_attack_objectives.call_args
141+
assert call_args.kwargs["language"] == "es" # Spanish language code
142+
143+
@pytest.mark.asyncio
144+
async def test_get_attack_objectives_tense_strategy_passes_language(self, mock_azure_ai_project, mock_credential):
145+
"""Test that _get_attack_objectives passes language parameter for tense strategy."""
146+
with patch("azure.ai.evaluation.red_team._red_team.GeneratedRAIClient") as mock_rai_client_class, \
147+
patch("azure.ai.evaluation.red_team._red_team.setup_logger") as mock_setup_logger, \
148+
patch("azure.ai.evaluation.red_team._red_team.initialize_pyrit"), \
149+
patch("azure.ai.evaluation.red_team._red_team._AttackObjectiveGenerator"):
150+
151+
mock_logger = MagicMock()
152+
mock_setup_logger.return_value = mock_logger
153+
154+
# Set up mock RAI client instance
155+
mock_rai_client = MagicMock()
156+
mock_rai_client.get_attack_objectives = AsyncMock(return_value=[
157+
{
158+
"id": "test-id",
159+
"messages": [{"role": "user", "content": "test prompt"}],
160+
"metadata": {"target_harms": [{"risk-type": "violence"}]}
161+
}
162+
])
163+
mock_rai_client_class.return_value = mock_rai_client
164+
165+
# Create RedTeam instance with French language
166+
agent = RedTeam(
167+
azure_ai_project=mock_azure_ai_project,
168+
credential=mock_credential,
169+
risk_categories=[RiskCategory.Violence],
170+
num_objectives=5,
171+
language=SupportedLanguages.French
172+
)
173+
174+
agent.generated_rai_client = mock_rai_client
175+
agent.scan_session_id = "test-session"
176+
177+
# Call _get_attack_objectives with tense strategy
178+
await agent._get_attack_objectives(
179+
risk_category=RiskCategory.Violence,
180+
application_scenario="test scenario",
181+
strategy="tense"
182+
)
183+
184+
# Verify that get_attack_objectives was called with French language and tense strategy
185+
mock_rai_client.get_attack_objectives.assert_called_once()
186+
call_args = mock_rai_client.get_attack_objectives.call_args
187+
assert call_args.kwargs["language"] == "fr" # French language code
188+
assert call_args.kwargs["strategy"] == "tense"

0 commit comments

Comments
 (0)