azure-ai-evaluation 1.0.0b3__py3-none-any.whl → 1.0.0b5__py3-none-any.whl

azure/ai/evaluation/simulator/_adversarial_simulator.py

@@ -6,20 +6,21 @@
 import asyncio
 import logging
 import random
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Literal, Optional, Union, cast
 
 from tqdm import tqdm
 
+from azure.ai.evaluation._common._experimental import experimental
+from azure.ai.evaluation._common.utils import validate_azure_ai_project
 from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
 from azure.ai.evaluation._http_utils import get_async_http_client
-from azure.ai.evaluation._model_configurations import AzureAIProject
 from azure.ai.evaluation.simulator import AdversarialScenario
 from azure.ai.evaluation.simulator._adversarial_scenario import _UnstableAdversarialScenario
+from azure.core.credentials import TokenCredential
 from azure.core.pipeline.policies import AsyncRetryPolicy, RetryMode
-from azure.identity import DefaultAzureCredential
 
 from ._constants import SupportedLanguages
-from ._conversation import CallbackConversationBot, ConversationBot, ConversationRole
+from ._conversation import CallbackConversationBot, ConversationBot, ConversationRole, ConversationTurn
 from ._conversation._conversation import simulate_conversation
 from ._model_tools import (
     AdversarialTemplateHandler,
@@ -28,11 +29,13 @@ from ._model_tools import (
     RAIClient,
     TokenScope,
 )
+from ._model_tools._template_handler import AdversarialTemplate, TemplateParameters
 from ._utils import JsonLineList
 
 logger = logging.getLogger(__name__)
 
 
+@experimental
 class AdversarialSimulator:
     """
     Initializes the adversarial simulator with a project scope.
@@ -44,41 +47,28 @@ class AdversarialSimulator:
     :type credential: ~azure.core.credentials.TokenCredential
     """
 
-    def __init__(self, *, azure_ai_project: AzureAIProject, credential=None):
+    def __init__(self, *, azure_ai_project: dict, credential):
         """Constructor."""
-        # check if azure_ai_project has the keys: subscription_id, resource_group_name and project_name
-        if not all(key in azure_ai_project for key in ["subscription_id", "resource_group_name", "project_name"]):
-            msg = "azure_ai_project must contain keys: subscription_id, resource_group_name, project_name"
+
+        try:
+            self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
+        except EvaluationException as e:
             raise EvaluationException(
-                message=msg,
-                internal_message=msg,
+                message=e.message,
+                internal_message=e.internal_message,
                 target=ErrorTarget.ADVERSARIAL_SIMULATOR,
-                category=ErrorCategory.MISSING_FIELD,
-                blame=ErrorBlame.USER_ERROR,
-            )
-        # check the value of the keys in azure_ai_project is not none
-        if not all(azure_ai_project[key] for key in ["subscription_id", "resource_group_name", "project_name"]):
-            msg = "subscription_id, resource_group_name and project_name cannot be None"
-            raise EvaluationException(
-                message=msg,
-                internal_message=msg,
-                target=ErrorTarget.ADVERSARIAL_SIMULATOR,
-                category=ErrorCategory.MISSING_FIELD,
-                blame=ErrorBlame.USER_ERROR,
-            )
-        if "credential" not in azure_ai_project and not credential:
-            credential = DefaultAzureCredential()
-        elif "credential" in azure_ai_project:
-            credential = azure_ai_project["credential"]
-        self.azure_ai_project = azure_ai_project
+                category=e.category,
+                blame=e.blame,
+            ) from e
+
         self.token_manager = ManagedIdentityAPITokenManager(
             token_scope=TokenScope.DEFAULT_AZURE_MANAGEMENT,
             logger=logging.getLogger("AdversarialSimulator"),
-            credential=credential,
+            credential=cast(TokenCredential, credential),
         )
-        self.rai_client = RAIClient(azure_ai_project=azure_ai_project, token_manager=self.token_manager)
+        self.rai_client = RAIClient(azure_ai_project=self.azure_ai_project, token_manager=self.token_manager)
         self.adversarial_template_handler = AdversarialTemplateHandler(
-            azure_ai_project=azure_ai_project, rai_client=self.rai_client
+            azure_ai_project=self.azure_ai_project, rai_client=self.rai_client
         )
 
     def _ensure_service_dependencies(self):
@@ -92,7 +82,7 @@ class AdversarialSimulator:
                 blame=ErrorBlame.USER_ERROR,
             )
 
-    # @monitor_adversarial_scenario
+    # pylint: disable=too-many-locals
     async def __call__(
         self,
         *,
@@ -106,10 +96,10 @@ class AdversarialSimulator:
         api_call_retry_sleep_sec: int = 1,
         api_call_delay_sec: int = 0,
         concurrent_async_task: int = 3,
-        _jailbreak_type: Optional[str] = None,
         language: SupportedLanguages = SupportedLanguages.English,
         randomize_order: bool = True,
         randomization_seed: Optional[int] = None,
+        **kwargs,
     ):
         """
         Executes the adversarial simulation against a specified target function asynchronously.
@@ -216,6 +206,7 @@ class AdversarialSimulator:
                 total_tasks,
             )
         total_tasks = min(total_tasks, max_simulation_results)
+        _jailbreak_type = kwargs.get("_jailbreak_type", None)
         if _jailbreak_type:
             jailbreak_dataset = await self.rai_client.get_jailbreaks_dataset(type=_jailbreak_type)
         progress_bar = tqdm(
@@ -263,16 +254,21 @@ class AdversarialSimulator:
 
         return JsonLineList(sim_results)
 
-    def _to_chat_protocol(self, *, conversation_history, template_parameters: Dict = None):
+    def _to_chat_protocol(
+        self,
+        *,
+        conversation_history: List[ConversationTurn],
+        template_parameters: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
+    ):
         if template_parameters is None:
             template_parameters = {}
         messages = []
         for _, m in enumerate(conversation_history):
             message = {"content": m.message, "role": m.role.value}
-            if "context" in m.full_response:
+            if m.full_response is not None and "context" in m.full_response:
                 message["context"] = m.full_response["context"]
             messages.append(message)
-        conversation_category = template_parameters.pop("metadata", {}).get("Category")
+        conversation_category = cast(Dict[str, str], template_parameters.pop("metadata", {})).get("Category")
         template_parameters["metadata"] = {}
         for key in (
             "conversation_starter",
@@ -280,6 +276,9 @@ class AdversarialSimulator:
             "target_population",
             "topic",
             "ch_template_placeholder",
+            "chatbot_name",
+            "name",
+            "group",
         ):
             template_parameters.pop(key, None)
         if conversation_category:
@@ -294,14 +293,14 @@ class AdversarialSimulator:
         self,
         *,
         target: Callable,
-        template,
-        parameters,
-        max_conversation_turns,
-        api_call_retry_limit,
-        api_call_retry_sleep_sec,
-        api_call_delay_sec,
-        language,
-        semaphore,
+        template: AdversarialTemplate,
+        parameters: TemplateParameters,
+        max_conversation_turns: int,
+        api_call_retry_limit: int,
+        api_call_retry_sleep_sec: int,
+        api_call_delay_sec: int,
+        language: SupportedLanguages,
+        semaphore: asyncio.Semaphore,
     ) -> List[Dict]:
         user_bot = self._setup_bot(role=ConversationRole.USER, template=template, parameters=parameters)
         system_bot = self._setup_bot(
@@ -324,9 +323,15 @@ class AdversarialSimulator:
                 api_call_delay_sec=api_call_delay_sec,
                 language=language,
             )
-        return self._to_chat_protocol(conversation_history=conversation_history, template_parameters=parameters)
 
-    def _get_user_proxy_completion_model(self, template_key, template_parameters):
+        return self._to_chat_protocol(
+            conversation_history=conversation_history,
+            template_parameters=cast(Dict[str, Union[str, Dict[str, str]]], parameters),
+        )
+
+    def _get_user_proxy_completion_model(
+        self, template_key: str, template_parameters: TemplateParameters
+    ) -> ProxyChatCompletionsModel:
         return ProxyChatCompletionsModel(
             name="raisvc_proxy_model",
             template_key=template_key,
@@ -338,8 +343,15 @@ class AdversarialSimulator:
             temperature=0.0,
         )
 
-    def _setup_bot(self, *, role, template, parameters, target: Callable = None):
-        if role == ConversationRole.USER:
+    def _setup_bot(
+        self,
+        *,
+        role: ConversationRole,
+        template: AdversarialTemplate,
+        parameters: TemplateParameters,
+        target: Optional[Callable] = None,
+    ) -> ConversationBot:
+        if role is ConversationRole.USER:
             model = self._get_user_proxy_completion_model(
                 template_key=template.template_name, template_parameters=parameters
             )
@@ -350,30 +362,46 @@ class AdversarialSimulator:
                 instantiation_parameters=parameters,
             )
 
-        if role == ConversationRole.ASSISTANT:
+        if role is ConversationRole.ASSISTANT:
+            if target is None:
+                msg = "Cannot setup system bot. Target is None"
 
-            def dummy_model() -> None:
-                return None
+                raise EvaluationException(
+                    message=msg,
+                    internal_message=msg,
+                    target=ErrorTarget.ADVERSARIAL_SIMULATOR,
+                    error_category=ErrorCategory.INVALID_VALUE,
+                    blame=ErrorBlame.SYSTEM_ERROR,
+                )
+
+            class DummyModel:
+                def __init__(self):
+                    self.name = "dummy_model"
+
+                def __call__(self) -> None:
+                    pass
 
-            dummy_model.name = "dummy_model"
             return CallbackConversationBot(
                 callback=target,
                 role=role,
-                model=dummy_model,
+                model=DummyModel(),
                 user_template=str(template),
                 user_template_parameters=parameters,
                 conversation_template="",
                 instantiation_parameters={},
             )
-        return ConversationBot(
-            role=role,
-            model=model,
-            conversation_template=template,
-            instantiation_parameters=parameters,
+
+        msg = "Invalid value for enum ConversationRole. This should never happen."
+        raise EvaluationException(
+            message=msg,
+            internal_message=msg,
+            target=ErrorTarget.ADVERSARIAL_SIMULATOR,
+            category=ErrorCategory.INVALID_VALUE,
+            blame=ErrorBlame.SYSTEM_ERROR,
         )
 
-    def _join_conversation_starter(self, parameters, to_join):
-        key = "conversation_starter"
+    def _join_conversation_starter(self, parameters: TemplateParameters, to_join: str) -> TemplateParameters:
+        key: Literal["conversation_starter"] = "conversation_starter"
         if key in parameters.keys():
             parameters[key] = f"{to_join} {parameters[key]}"
         else:

azure/ai/evaluation/simulator/_conversation/__init__.py

@@ -7,7 +7,7 @@ import copy
 import logging
 import time
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
 
 import jinja2
 
@@ -15,6 +15,7 @@ from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarg
 from azure.ai.evaluation._http_utils import AsyncHttpPipeline
 
 from .._model_tools import LLMBase, OpenAIChatCompletionsModel
+from .._model_tools._template_handler import TemplateParameters
 from .constants import ConversationRole
 
 
@@ -40,7 +41,7 @@ class ConversationTurn:
     role: "ConversationRole"
     name: Optional[str] = None
     message: str = ""
-    full_response: Optional[Any] = None
+    full_response: Optional[Dict[str, Any]] = None
     request: Optional[Any] = None
 
     def to_openai_chat_format(self, reverse: bool = False) -> Dict[str, str]:
@@ -109,7 +110,7 @@ class ConversationBot:
         role: ConversationRole,
         model: Union[LLMBase, OpenAIChatCompletionsModel],
         conversation_template: str,
-        instantiation_parameters: Dict[str, str],
+        instantiation_parameters: TemplateParameters,
     ) -> None:
         self.role = role
         self.conversation_template_orig = conversation_template
@@ -118,13 +119,13 @@ class ConversationBot:
         )
         self.persona_template_args = instantiation_parameters
         if self.role == ConversationRole.USER:
-            self.name = self.persona_template_args.get("name", role.value)
+            self.name: str = cast(str, self.persona_template_args.get("name", role.value))
         else:
-            self.name = self.persona_template_args.get("chatbot_name", role.value) or model.name
+            self.name = cast(str, self.persona_template_args.get("chatbot_name", role.value)) or model.name
         self.model = model
 
         self.logger = logging.getLogger(repr(self))
-        self.conversation_starter = None  # can either be a dictionary or jinja template
+        self.conversation_starter: Optional[Union[str, jinja2.Template, Dict]] = None
         if role == ConversationRole.USER:
             if "conversation_starter" in self.persona_template_args:
                 conversation_starter_content = self.persona_template_args["conversation_starter"]
@@ -148,7 +149,7 @@ class ConversationBot:
         conversation_history: List[ConversationTurn],
         max_history: int,
         turn_number: int = 0,
-    ) -> Tuple[dict, dict, int, dict]:
+    ) -> Tuple[dict, dict, float, dict]:
         """
         Prompt the ConversationBot for a response.
 
@@ -161,7 +162,7 @@ class ConversationBot:
         :param turn_number: Parameters used to query GPT-4 model.
         :type turn_number: int
         :return: The response from the ConversationBot.
-        :rtype: Tuple[dict, dict, int, dict]
+        :rtype: Tuple[dict, dict, float, dict]
         """
 
         # check if this is the first turn and the conversation_starter is not None,
@@ -169,11 +170,11 @@ class ConversationBot:
         if turn_number == 0 and self.conversation_starter is not None:
             # if conversation_starter is a dictionary, pass it into samples as is
             if isinstance(self.conversation_starter, dict):
-                samples = [self.conversation_starter]
+                samples: List[Union[str, jinja2.Template, Dict]] = [self.conversation_starter]
             if isinstance(self.conversation_starter, jinja2.Template):
                 samples = [self.conversation_starter.render(**self.persona_template_args)]
             else:
-                samples = [self.conversation_starter]  # type: ignore[attr-defined]
+                samples = [self.conversation_starter]
             time_taken = 0
 
             finish_reason = ["stop"]
@@ -238,7 +239,7 @@ class CallbackConversationBot(ConversationBot):
         self,
         callback: Callable,
         user_template: str,
-        user_template_parameters: Dict,
+        user_template_parameters: TemplateParameters,
         *args,
         **kwargs,
     ) -> None:
@@ -254,7 +255,7 @@ class CallbackConversationBot(ConversationBot):
         conversation_history: List[Any],
         max_history: int,
         turn_number: int = 0,
-    ) -> Tuple[dict, dict, int, dict]:
+    ) -> Tuple[dict, dict, float, dict]:
         chat_protocol_message = self._to_chat_protocol(
             self.user_template, conversation_history, self.user_template_parameters
         )

azure/ai/evaluation/simulator/_conversation/_conversation.py

@@ -4,7 +4,7 @@
 
 import asyncio
 import logging
-from typing import Callable, Dict, List, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
 from azure.ai.evaluation.simulator._constants import SupportedLanguages
@@ -80,7 +80,7 @@ async def simulate_conversation(
     history_limit: int = 5,
     api_call_delay_sec: float = 0,
     logger: logging.Logger = logging.getLogger(__name__),
-) -> Tuple:
+) -> Tuple[Optional[str], List[ConversationTurn]]:
     """
     Simulate a conversation between the given bots.
 
@@ -99,7 +99,7 @@ async def simulate_conversation(
     :keyword logger: The logger to use for logging. Defaults to the logger named after the current module.
     :paramtype logger: logging.Logger
     :return: Simulation a conversation between the given bots.
-    :rtype: Tuple
+    :rtype: Tuple[Optional[str], List[ConversationTurn]]
     """
 
     # Read the first prompt.
@@ -110,7 +110,7 @@ async def simulate_conversation(
         turn_number=0,
     )
     if "id" in first_response:
-        conversation_id = first_response["id"]
+        conversation_id: Optional[str] = first_response["id"]
     else:
         conversation_id = None
     first_prompt = first_response["samples"][0]

azure/ai/evaluation/simulator/_data_sources/__init__.py

@@ -0,0 +1,3 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------