Merge pull request Azure#8 from slister1001/restructure-red-team

Restructure RedTeamAgent

Author: slister1001

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

@@ -1,30 +1,3 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
-# ---------------------------------------------------------
-
-from enum import Enum
-from ._red_team_agent import RedTeamAgent, AttackStrategy
-
-class RiskCategory(str, Enum):
-    """Risk categories for attack objectives."""
-    HateUnfairness = "HateUnfairness"
-    Violence = "Violence"
-    Sexual = "Sexual"
-    SelfHarm = "SelfHarm" 
-    ProtectedMaterial = "ProtectedMaterial"
-    IndirectJailbreak = "IndirectJailbreak"
-    DirectJailbreak = "DirectJailbreak"
-    CodeVulnerability = "CodeVulnerability"
-    InferenceSensitiveAttributes = "InferenceSensitiveAttributes"
-
-class AttackObjectiveGenerator:
-    """Generator for creating attack objectives.
-
-    :param risk_categories: List of risk categories to generate attack objectives for
-    :type risk_categories: List[RiskCategory]
-    """
-    def __init__(self, risk_categories: list[RiskCategory], num_objectives: int = 10):
-        self.risk_categories = risk_categories
-        self.num_objectives = num_objectives
-
-__all__ = ["RedTeamAgent", "AttackStrategy", "RiskCategory", "AttackObjectiveGenerator"]
+# ---------------------------------------------------------

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_safety_evaluation/_safety_evaluation.py

@@ -18,7 +18,6 @@
 from azure.ai.evaluation.simulator import Simulator, AdversarialSimulator, AdversarialScenario, AdversarialScenarioJailbreak, IndirectAttackSimulator, DirectAttackSimulator
 from azure.ai.evaluation.simulator._utils import JsonLineList
 from azure.ai.evaluation._common.utils import validate_azure_ai_project
-from pyrit.prompt_target import PromptChatTarget #TODO: Remove this once eval logic for red team agent is moved to red team agent
 from azure.ai.evaluation._model_configurations import AzureOpenAIModelConfiguration, OpenAIModelConfiguration
 from azure.core.credentials import TokenCredential
 import json
@@ -460,7 +459,7 @@ def _check_target_is_callback(target:Callable) -> bool:
     def _validate_inputs(
             self,
             evaluators: List[_SafetyEvaluator],
-            target: Union[Callable, AzureOpenAIModelConfiguration, OpenAIModelConfiguration, PromptChatTarget], #TODO: Remove this once eval logic for red team agent is moved to red team agent
+            target: Union[Callable, AzureOpenAIModelConfiguration, OpenAIModelConfiguration],
             num_turns: int = 1,
             scenario: Optional[Union[AdversarialScenario, AdversarialScenarioJailbreak]] = None,
             source_text: Optional[str] = None,
@@ -478,10 +477,10 @@ def _validate_inputs(
         :param source_text: The source text to use as grounding document in the evaluation.
         :type source_text: Optional[str]
         ''' 
-        if not callable(target) and not isinstance(target, PromptChatTarget): #TODO: Remove this once eval logic for red team agent is moved to red team agent
+        if not callable(target):
             self._validate_model_config(target)
         #TODO: Remove self._check_target_is_callback(target)) once eval logic for red team agent is moved to red team agent
-        elif not isinstance(target, PromptChatTarget) and not (self._check_target_returns_str(target) or self._check_target_is_callback(target)): 
+        elif self._check_target_returns_str(target): 
             self.logger.error(f"Target function {target} does not return a string.")
             msg = f"Target function {target} does not return a string."
             raise EvaluationException(
@@ -574,7 +573,7 @@ def _calculate_defect_rate(self, evaluation_result_dict) -> EvaluationResult:
 
     async def __call__(
             self,
-            target: Union[Callable, AzureOpenAIModelConfiguration, OpenAIModelConfiguration, PromptChatTarget],
+            target: Union[Callable, AzureOpenAIModelConfiguration, OpenAIModelConfiguration],
             evaluators: List[_SafetyEvaluator] = [],
             evaluation_name: Optional[str] = None,
             num_turns : int = 1,

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

@@ -0,0 +1,11 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from enum import Enum
+from typing import List
+from .red_team_agent import RedTeamAgent
+from .attack_strategy import AttackStrategy
+from .attack_objective_generator import AttackObjectiveGenerator, RiskCategory
+
+__all__ = ["RedTeamAgent", "AttackStrategy", "RiskCategory", "AttackObjectiveGenerator"]

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team_agent/attack_objective_generator.py

@@ -0,0 +1,27 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from enum import Enum
+from typing import List
+
+class RiskCategory(str, Enum):
+    """Risk categories for attack objectives."""
+    HateUnfairness = "hate_unfairness"
+    Violence = "violence"
+    Sexual = "sexual"
+    SelfHarm = "self_harm" 
+    ProtectedMaterial = "protected_material"
+    IndirectJailbreak = "indirect_jailbreak"
+    DirectJailbreak = "direct_jailbreak"
+    CodeVulnerability = "code_vulnerability"
+    InferenceSensitiveAttributes = "inference_sensitive_attributes"
+
+class AttackObjectiveGenerator:
+    """Generator for creating attack objectives.
+
+    :param risk_categories: List of risk categories to generate attack objectives for
+    :type risk_categories: List[RiskCategory]
+    """
+    def __init__(self, risk_categories: List[RiskCategory], num_objectives: int = 10):
+        self.risk_categories = risk_categories
+        self.num_objectives = num_objectives

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team_agent/attack_strategy.py

@@ -0,0 +1,49 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+from enum import Enum
+from typing import List
+
+class AttackStrategy(Enum):
+    """Strategies for attacks."""
+    EASY = "easy"
+    MODERATE = "moderate"
+    DIFFICULT = "difficult"
+    AnsiAttack = "ansi_attack"
+    AsciiArt = "ascii_art"
+    AsciiSmuggler = "ascii_smuggler"
+    Atbash = "atbash"
+    Base64 = "base64"
+    Binary = "binary"
+    Caesar = "caesar"
+    CharacterSpace = "character_space"
+    CharSwap = "char_swap"
+    Diacritic = "diacritic"
+    Flip = "flip"
+    Leetspeak = "leetspeak"
+    MaliciousQuestion = "malicious_question"
+    Math = "math"
+    Morse = "morse"
+    Persuasion = "persuasion"
+    ROT13 = "rot13"
+    RepeatToken = "repeat_token"
+    SuffixAppend = "suffix_append"
+    StringJoin = "string_join"
+    Tense = "tense"
+    Tone = "tone"
+    Translation = "translation"
+    UnicodeConfusable = "unicode_confusable"
+    UnicodeSubstitution = "unicode_substitution"
+    Url = "url"
+    Variation = "variation"
+    Baseline = "baseline"
+    Jailbreak = "jailbreak"
+
+    @classmethod
+    def Compose(cls, items: List["AttackStrategy"]) -> List["AttackStrategy"]:
+        for item in items:
+            if not isinstance(item, cls):
+                raise ValueError("All items must be instances of AttackStrategy")
+        if len(items) > 2: 
+            raise ValueError("Composed strategies must have at most 2 items")
+        return items

sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_safety_evaluation/_callback_chat_target.py renamed to sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team_agent/callback_chat_target.py

@@ -1,7 +1,6 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
-
 import logging
 from typing import Any, Callable, Dict, List, Optional
 
@@ -24,7 +23,7 @@ def __init__(
         """
         Initializes an instance of the CallbackChatTarget class.
 
-        It is intended to be used with PYRIT where users define a callback function
+        It is intended to be used with PyRIT where users define a callback function
         that handles sending a prompt to a target and receiving a response.
         The CallbackChatTarget class is a wrapper around the callback function that allows it to be used
         as a target in the PyRIT framework.