azure-ai-evaluation 1.0.0b2__py3-none-any.whl → 1.0.0b4__py3-none-any.whl

azure/ai/evaluation/simulator/_adversarial_simulator.py

@@ -6,21 +6,22 @@
 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.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 ._helpers import experimental
 from ._model_tools import (
     AdversarialTemplateHandler,
     ManagedIdentityAPITokenManager,
@@ -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",
@@ -294,14 +290,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 +320,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 +340,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 +359,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/_direct_attack_simulator.py

@@ -1,54 +1,25 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
+# pylint: disable=C0301,C0114,R0913,R0903
 # noqa: E501
-import functools
 import logging
 from random import randint
-from typing import Callable, Optional
-
-from promptflow._sdk._telemetry import ActivityType, monitor_operation
+from typing import Callable, Optional, cast
 
+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._model_configurations import AzureAIProject
 from azure.ai.evaluation.simulator import AdversarialScenario
-from azure.identity import DefaultAzureCredential
+from azure.core.credentials import TokenCredential
 
 from ._adversarial_simulator import AdversarialSimulator
+from ._helpers import experimental
 from ._model_tools import AdversarialTemplateHandler, ManagedIdentityAPITokenManager, RAIClient, TokenScope
 
 logger = logging.getLogger(__name__)
 
 
-def monitor_adversarial_scenario(func) -> Callable:
-    """Decorator to monitor adversarial scenario.
-
-    :param func: The function to be decorated.
-    :type func: Callable
-    :return: The decorated function.
-    :rtype: Callable
-    """
-
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        scenario = str(kwargs.get("scenario", None))
-        max_conversation_turns = kwargs.get("max_conversation_turns", None)
-        max_simulation_results = kwargs.get("max_simulation_results", None)
-        decorated_func = monitor_operation(
-            activity_name="jailbreak.adversarial.simulator.call",
-            activity_type=ActivityType.PUBLICAPI,
-            custom_dimensions={
-                "scenario": scenario,
-                "max_conversation_turns": max_conversation_turns,
-                "max_simulation_results": max_simulation_results,
-            },
-        )(func)
-
-        return decorated_func(*args, **kwargs)
-
-    return wrapper
-
-
+@experimental
 class DirectAttackSimulator:
     """
     Initialize a UPIA (user prompt injected attack) jailbreak adversarial simulator with a project scope.
@@ -61,42 +32,28 @@ class DirectAttackSimulator:
     :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, project_name, credential
-        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 and project_name"
-            raise EvaluationException(
-                message=msg,
-                internal_message=msg,
-                target=ErrorTarget.DIRECT_ATTACK_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 keys cannot be None"
+
+        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.DIRECT_ATTACK_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.credential = credential
-        self.azure_ai_project = azure_ai_project
+                category=e.category,
+                blame=e.blame,
+            ) from e
+        self.credential = cast(TokenCredential, credential)
         self.token_manager = ManagedIdentityAPITokenManager(
             token_scope=TokenScope.DEFAULT_AZURE_MANAGEMENT,
             logger=logging.getLogger("AdversarialSimulator"),
-            credential=credential,
+            credential=self.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):
@@ -110,7 +67,6 @@ class DirectAttackSimulator:
                 blame=ErrorBlame.USER_ERROR,
             )
 
-    # @monitor_adversarial_scenario
     async def __call__(
         self,
         *,
@@ -222,7 +178,9 @@ class DirectAttackSimulator:
         if not randomization_seed:
             randomization_seed = randint(0, 1000000)
 
-        regular_sim = AdversarialSimulator(azure_ai_project=self.azure_ai_project, credential=self.credential)
+        regular_sim = AdversarialSimulator(
+            azure_ai_project=cast(dict, self.azure_ai_project), credential=self.credential
+        )
         regular_sim_results = await regular_sim(
             scenario=scenario,
             target=target,
@@ -235,7 +193,7 @@ class DirectAttackSimulator:
             randomize_order=True,
             randomization_seed=randomization_seed,
         )
-        jb_sim = AdversarialSimulator(azure_ai_project=self.azure_ai_project, credential=self.credential)
+        jb_sim = AdversarialSimulator(azure_ai_project=cast(dict, self.azure_ai_project), credential=self.credential)
         jb_sim_results = await jb_sim(
             scenario=scenario,
             target=target,

azure/ai/evaluation/simulator/_helpers/_experimental.py

@@ -6,9 +6,9 @@ import functools
 import inspect
 import logging
 import sys
-from typing import Callable, Type, TypeVar, Union
+from typing import Callable, Type, TypeVar, Union, overload
 
-from typing_extensions import ParamSpec
+from typing_extensions import ParamSpec, TypeGuard
 
 DOCSTRING_TEMPLATE = ".. note::    {0} {1}\n\n"
 DOCSTRING_DEFAULT_INDENTATION = 8
@@ -22,20 +22,31 @@ EXPERIMENTAL_LINK_MESSAGE = (
 _warning_cache = set()
 module_logger = logging.getLogger(__name__)
 
-TExperimental = TypeVar("TExperimental", bound=Union[Type, Callable])
 P = ParamSpec("P")
 T = TypeVar("T")
 
 
-def experimental(wrapped: TExperimental) -> TExperimental:
+@overload
+def experimental(wrapped: Type[T]) -> Type[T]: ...
+
+
+@overload
+def experimental(wrapped: Callable[P, T]) -> Callable[P, T]: ...
+
+
+def experimental(wrapped: Union[Type[T], Callable[P, T]]) -> Union[Type[T], Callable[P, T]]:
     """Add experimental tag to a class or a method.
 
     :param wrapped: Either a Class or Function to mark as experimental
-    :type wrapped: TExperimental
+    :type wrapped: Union[Type[T], Callable[P, T]]
     :return: The wrapped class or method
-    :rtype: TExperimental
+    :rtype: Union[Type[T], Callable[P, T]]
     """
-    if inspect.isclass(wrapped):
+
+    def is_class(t: Union[Type[T], Callable[P, T]]) -> TypeGuard[Type[T]]:
+        return isinstance(t, type)
+
+    if is_class(wrapped):
         return _add_class_docstring(wrapped)
     if inspect.isfunction(wrapped):
         return _add_method_docstring(wrapped)
@@ -74,11 +85,11 @@ def _add_class_docstring(cls: Type[T]) -> Type[T]:
         cls.__doc__ = _add_note_to_docstring(cls.__doc__, doc_string)
     else:
         cls.__doc__ = doc_string + ">"
-    cls.__init__ = _add_class_warning(cls.__init__)
+    cls.__init__ = _add_class_warning(cls.__init__)  # type: ignore[method-assign]
     return cls
 
 
-def _add_method_docstring(func: Callable[P, T] = None) -> Callable[P, T]:
+def _add_method_docstring(func: Callable[P, T]) -> Callable[P, T]:
     """Add experimental tag to the method doc string.
 
     :param func: The function to update

azure/ai/evaluation/simulator/_helpers/_simulator_data_classes.py

@@ -18,7 +18,7 @@ class Turn:
 
     role: Union[str, ConversationRole]
     content: str
-    context: str = None
+    context: Optional[str] = None
 
     def to_dict(self) -> Dict[str, Optional[str]]:
         """
@@ -42,13 +42,13 @@ class ConversationHistory:
     Conversation history class to keep track of the conversation turns in a conversation.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """
         Initializes the conversation history with an empty list of turns.
         """
         self.history: List[Turn] = []
 
-    def add_to_history(self, turn: Turn):
+    def add_to_history(self, turn: Turn) -> None:
         """
         Adds a turn to the conversation history.
 
@@ -57,7 +57,7 @@ class ConversationHistory:
         """
         self.history.append(turn)
 
-    def to_list(self) -> List[Dict[str, str]]:
+    def to_list(self) -> List[Dict[str, Optional[str]]]:
         """
         Converts the conversation history to a list of dictionaries.