azure-ai-evaluation 1.0.1__py3-none-any.whl → 1.13.3__py3-none-any.whl

azure/ai/evaluation/simulator/_adversarial_simulator.py

@@ -6,30 +6,38 @@
 import asyncio
 import logging
 import random
-from typing import Any, Callable, Dict, List, Literal, Optional, Union, cast
-from itertools import zip_longest
+from typing import Any, Callable, Dict, List, Optional, Union, cast
+import uuid
+import warnings
 
 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._common.utils import validate_azure_ai_project, is_onedp_project
+from azure.ai.evaluation._common.onedp._client import ProjectsClient as AIProjectClient
 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 import AdversarialScenario, AdversarialScenarioJailbreak
 from azure.ai.evaluation.simulator._adversarial_scenario import _UnstableAdversarialScenario
+from azure.ai.evaluation._constants import TokenScope
 from azure.core.credentials import TokenCredential
 from azure.core.pipeline.policies import AsyncRetryPolicy, RetryMode
 
 from ._constants import SupportedLanguages
-from ._conversation import CallbackConversationBot, ConversationBot, ConversationRole, ConversationTurn
+from ._conversation import (
+    CallbackConversationBot,
+    MultiModalConversationBot,
+    ConversationBot,
+    ConversationRole,
+    ConversationTurn,
+)
 from ._conversation._conversation import simulate_conversation
 from ._model_tools import (
     AdversarialTemplateHandler,
     ManagedIdentityAPITokenManager,
     ProxyChatCompletionsModel,
     RAIClient,
-    TokenScope,
 )
 from ._model_tools._template_handler import AdversarialTemplate, TemplateParameters
 from ._utils import JsonLineList
@@ -42,9 +50,9 @@ class AdversarialSimulator:
     """
     Initializes the adversarial simulator with a project scope.
 
-    :param azure_ai_project: The scope of the Azure AI project. It contains subscription id, resource group, and project
-        name.
-    :type azure_ai_project: ~azure.ai.evaluation.AzureAIProject
+    :param azure_ai_project: The Azure AI project, which can either be a string representing the project endpoint
+        or an instance of AzureAIProject. It contains subscription id, resource group, and project name.
+    :type azure_ai_project: Union[str, AzureAIProject]
     :param credential: The credential for connecting to Azure AI project.
     :type credential: ~azure.core.credentials.TokenCredential
 
@@ -59,26 +67,45 @@ class AdversarialSimulator:
                 2 conversation turns each (4 messages per result).
     """
 
-    def __init__(self, *, azure_ai_project: AzureAIProject, credential: TokenCredential):
+    def __init__(self, *, azure_ai_project: Union[str, AzureAIProject], credential: TokenCredential):
         """Constructor."""
-
-        try:
-            self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
-        except EvaluationException as e:
-            raise EvaluationException(
-                message=e.message,
-                internal_message=e.internal_message,
-                target=ErrorTarget.ADVERSARIAL_SIMULATOR,
-                category=e.category,
-                blame=e.blame,
-            ) from e
-
-        self.token_manager = ManagedIdentityAPITokenManager(
-            token_scope=TokenScope.DEFAULT_AZURE_MANAGEMENT,
-            logger=logging.getLogger("AdversarialSimulator"),
-            credential=cast(TokenCredential, credential),
+        warnings.warn(
+            "DEPRECATION NOTE: Azure AI Evaluation SDK has discontinued active development on the AdversarialSimulator class."
+            + " While existing functionality remains available in preview, it is no longer recommended for production workloads or future integration. "
+            + "We recommend users migrate to the AI Red Teaming Agent for future use as it supports full parity of functionality."
+            + " See https://aka.ms/airedteamingagent-sample for details on AI Red Teaming Agent.",
+            DeprecationWarning,
+            stacklevel=2,
         )
-        self.rai_client = RAIClient(azure_ai_project=self.azure_ai_project, token_manager=self.token_manager)
+
+        if is_onedp_project(azure_ai_project):
+            self.azure_ai_project = azure_ai_project
+            self.credential = cast(TokenCredential, credential)
+            self.token_manager = ManagedIdentityAPITokenManager(
+                token_scope=TokenScope.COGNITIVE_SERVICES_MANAGEMENT,
+                logger=logging.getLogger("AdversarialSimulator"),
+                credential=self.credential,
+            )
+            self.rai_client = AIProjectClient(endpoint=azure_ai_project, credential=credential)
+        else:
+            try:
+                self.azure_ai_project = validate_azure_ai_project(azure_ai_project)
+            except EvaluationException as e:
+                raise EvaluationException(
+                    message=e.message,
+                    internal_message=e.internal_message,
+                    target=ErrorTarget.ADVERSARIAL_SIMULATOR,
+                    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=self.credential,
+            )
+            self.rai_client = RAIClient(azure_ai_project=self.azure_ai_project, token_manager=self.token_manager)
+
         self.adversarial_template_handler = AdversarialTemplateHandler(
             azure_ai_project=self.azure_ai_project, rai_client=self.rai_client
         )
@@ -182,6 +209,14 @@ class AdversarialSimulator:
             )
         self._ensure_service_dependencies()
         templates = await self.adversarial_template_handler._get_content_harm_template_collections(scenario.value)
+        if len(templates) == 0:
+            raise EvaluationException(
+                message="Templates not found. Please check https://aka.ms/azureaiadvsimulator-regionsupport for region support.",
+                internal_message="Please check https://aka.ms/azureaiadvsimulator-regionsupport for region support.",
+                target=ErrorTarget.ADVERSARIAL_SIMULATOR,
+            )
+        simulation_id = str(uuid.uuid4())
+        logger.warning("Use simulation_id to help debug the issue: %s", str(simulation_id))
         concurrent_async_task = min(concurrent_async_task, 1000)
         semaphore = asyncio.Semaphore(concurrent_async_task)
         sim_results = []
@@ -198,46 +233,83 @@ class AdversarialSimulator:
         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)
+            if isinstance(self.rai_client, RAIClient):
+                jailbreak_dataset = await self.rai_client.get_jailbreaks_dataset(type=_jailbreak_type)
+            elif isinstance(self.rai_client, AIProjectClient):
+                jailbreak_dataset = self.rai_client.red_teams.get_jail_break_dataset_with_type(type=_jailbreak_type)
         progress_bar = tqdm(
             total=total_tasks,
             desc="generating jailbreak simulations" if _jailbreak_type else "generating simulations",
             ncols=100,
             unit="simulations",
         )
-
         if randomize_order:
             # The template parameter lists are persistent across sim runs within a session,
             # So randomize a the selection instead of the parameter list directly,
             # or a potentially large deep copy.
             if randomization_seed is not None:
-                random.seed(randomization_seed)
-            random.shuffle(templates)
-        parameter_lists = [t.template_parameters for t in templates]
-        zipped_parameters = list(zip_longest(*parameter_lists))
-        for param_group in zipped_parameters:
-            for template, parameter in zip(templates, param_group):
-                if _jailbreak_type == "upia":
-                    parameter = self._join_conversation_starter(parameter, random.choice(jailbreak_dataset))
-                tasks.append(
-                    asyncio.create_task(
-                        self._simulate_async(
-                            target=target,
-                            template=template,
-                            parameters=parameter,
-                            max_conversation_turns=max_conversation_turns,
-                            api_call_retry_limit=api_call_retry_limit,
-                            api_call_retry_sleep_sec=api_call_retry_sleep_sec,
-                            api_call_delay_sec=api_call_delay_sec,
-                            language=language,
-                            semaphore=semaphore,
-                        )
+                # Create a local random instance to avoid polluting global state
+                local_random = random.Random(randomization_seed)
+                local_random.shuffle(templates)
+            else:
+                random.shuffle(templates)
+
+        # Prepare task parameters based on scenario - but use a single append call for all scenarios
+        tasks = []
+        template_parameter_pairs = []
+
+        if scenario == AdversarialScenario.ADVERSARIAL_CONVERSATION:
+            # For ADVERSARIAL_CONVERSATION, flatten the parameters
+            for i, template in enumerate(templates):
+                if not template.template_parameters:
+                    continue
+                for parameter in template.template_parameters:
+                    template_parameter_pairs.append((template, parameter))
+        else:
+            # Use original logic for other scenarios - zip parameters
+            parameter_lists = [t.template_parameters for t in templates]
+            zipped_parameters = list(zip(*parameter_lists))
+
+            for param_group in zipped_parameters:
+                for template, parameter in zip(templates, param_group):
+                    template_parameter_pairs.append((template, parameter))
+
+        # Limit to max_simulation_results if needed
+        if len(template_parameter_pairs) > max_simulation_results:
+            template_parameter_pairs = template_parameter_pairs[
+                :max_simulation_results
+            ]  # Create a seeded random instance for jailbreak selection if randomization_seed is provided
+        jailbreak_random = None
+        if _jailbreak_type == "upia" and randomization_seed is not None:
+            jailbreak_random = random.Random(randomization_seed)
+
+        # Single task append loop for all scenarios
+        for template, parameter in template_parameter_pairs:
+            if _jailbreak_type == "upia":
+                if jailbreak_random is not None:
+                    selected_jailbreak = jailbreak_random.choice(jailbreak_dataset)
+                else:
+                    selected_jailbreak = random.choice(jailbreak_dataset)
+                parameter = self._add_jailbreak_parameter(parameter, selected_jailbreak)
+
+            tasks.append(
+                asyncio.create_task(
+                    self._simulate_async(
+                        target=target,
+                        template=template,
+                        parameters=parameter,
+                        max_conversation_turns=max_conversation_turns,
+                        api_call_retry_limit=api_call_retry_limit,
+                        api_call_retry_sleep_sec=api_call_retry_sleep_sec,
+                        api_call_delay_sec=api_call_delay_sec,
+                        language=language,
+                        semaphore=semaphore,
+                        scenario=scenario,
+                        simulation_id=simulation_id,
                     )
                 )
-                if len(tasks) >= max_simulation_results:
-                    break
-            if len(tasks) >= max_simulation_results:
-                break
+            )
+
         for task in asyncio.as_completed(tasks):
             sim_results.append(await task)
             progress_bar.update(1)
@@ -292,28 +364,43 @@ class AdversarialSimulator:
         api_call_delay_sec: int,
         language: SupportedLanguages,
         semaphore: asyncio.Semaphore,
+        scenario: Union[AdversarialScenario, AdversarialScenarioJailbreak],
+        simulation_id: str = "",
     ) -> List[Dict]:
-        user_bot = self._setup_bot(role=ConversationRole.USER, template=template, parameters=parameters)
+        user_bot = self._setup_bot(
+            role=ConversationRole.USER,
+            template=template,
+            parameters=parameters,
+            scenario=scenario,
+            simulation_id=simulation_id,
+        )
         system_bot = self._setup_bot(
-            target=target, role=ConversationRole.ASSISTANT, template=template, parameters=parameters
+            target=target, role=ConversationRole.ASSISTANT, template=template, parameters=parameters, scenario=scenario
         )
         bots = [user_bot, system_bot]
-        session = get_async_http_client().with_policies(
-            retry_policy=AsyncRetryPolicy(
-                retry_total=api_call_retry_limit,
-                retry_backoff_factor=api_call_retry_sleep_sec,
-                retry_mode=RetryMode.Fixed,
-            )
-        )
 
-        async with semaphore, session:
-            _, conversation_history = await simulate_conversation(
-                bots=bots,
-                session=session,
-                turn_limit=max_conversation_turns,
-                api_call_delay_sec=api_call_delay_sec,
-                language=language,
+        async def run_simulation(session_obj):
+            async with semaphore:
+                _, conversation_history = await simulate_conversation(
+                    bots=bots,
+                    session=session_obj,
+                    turn_limit=max_conversation_turns,
+                    api_call_delay_sec=api_call_delay_sec,
+                    language=language,
+                )
+            return conversation_history
+
+        if isinstance(self.rai_client, AIProjectClient):
+            session = self.rai_client
+        else:
+            session = get_async_http_client().with_policies(
+                retry_policy=AsyncRetryPolicy(
+                    retry_total=api_call_retry_limit,
+                    retry_backoff_factor=api_call_retry_sleep_sec,
+                    retry_mode=RetryMode.Fixed,
+                )
             )
+        conversation_history = await run_simulation(session)
 
         return self._to_chat_protocol(
             conversation_history=conversation_history,
@@ -321,17 +408,23 @@ class AdversarialSimulator:
         )
 
     def _get_user_proxy_completion_model(
-        self, template_key: str, template_parameters: TemplateParameters
+        self, template_key: str, template_parameters: TemplateParameters, simulation_id: str = ""
     ) -> ProxyChatCompletionsModel:
+        endpoint_url = (
+            self.rai_client._config.endpoint + "/redTeams/simulation/chat/completions/submit"
+            if isinstance(self.rai_client, AIProjectClient)
+            else self.rai_client.simulation_submit_endpoint
+        )
         return ProxyChatCompletionsModel(
             name="raisvc_proxy_model",
             template_key=template_key,
             template_parameters=template_parameters,
-            endpoint_url=self.rai_client.simulation_submit_endpoint,
+            endpoint_url=endpoint_url,
             token_manager=self.token_manager,
             api_version="2023-07-01-preview",
             max_tokens=1200,
             temperature=0.0,
+            simulation_id=simulation_id,
         )
 
     def _setup_bot(
@@ -341,10 +434,14 @@ class AdversarialSimulator:
         template: AdversarialTemplate,
         parameters: TemplateParameters,
         target: Optional[Callable] = None,
+        scenario: Union[AdversarialScenario, AdversarialScenarioJailbreak],
+        simulation_id: str = "",
     ) -> ConversationBot:
         if role is ConversationRole.USER:
             model = self._get_user_proxy_completion_model(
-                template_key=template.template_name, template_parameters=parameters
+                template_key=template.template_name,
+                template_parameters=parameters,
+                simulation_id=simulation_id,
             )
             return ConversationBot(
                 role=role,
@@ -372,6 +469,21 @@ class AdversarialSimulator:
                 def __call__(self) -> None:
                     pass
 
+            if scenario in [
+                _UnstableAdversarialScenario.ADVERSARIAL_IMAGE_GEN,
+                _UnstableAdversarialScenario.ADVERSARIAL_IMAGE_MULTIMODAL,
+            ]:
+                return MultiModalConversationBot(
+                    callback=target,
+                    role=role,
+                    model=DummyModel(),
+                    user_template=str(template),
+                    user_template_parameters=parameters,
+                    rai_client=self.rai_client,
+                    conversation_template="",
+                    instantiation_parameters={},
+                )
+
             return CallbackConversationBot(
                 callback=target,
                 role=role,
@@ -391,13 +503,8 @@ class AdversarialSimulator:
             blame=ErrorBlame.SYSTEM_ERROR,
         )
 
-    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:
-            parameters[key] = to_join
-
+    def _add_jailbreak_parameter(self, parameters: TemplateParameters, to_join: str) -> TemplateParameters:
+        parameters["jailbreak_string"] = to_join
         return parameters
 
     def call_sync(

azure/ai/evaluation/simulator/_constants.py

@@ -25,3 +25,4 @@ class SupportedLanguages(Enum):
     Portuguese = "pt"
     Japanese = "ja"
     English = "en"
+    Korean = "ko"

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

@@ -8,13 +8,14 @@ import logging
 import time
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
-
+import base64
+import re
 import jinja2
 
 from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
 from azure.ai.evaluation._http_utils import AsyncHttpPipeline
-
-from .._model_tools import LLMBase, OpenAIChatCompletionsModel
+from .._model_tools import LLMBase, OpenAIChatCompletionsModel, RAIClient
+from azure.ai.evaluation._common.onedp._client import ProjectsClient as AIProjectClient
 from .._model_tools._template_handler import TemplateParameters
 from .constants import ConversationRole
 
@@ -136,7 +137,7 @@ class ConversationBot:
                         self.conversation_starter = jinja2.Template(
                             conversation_starter_content, undefined=jinja2.StrictUndefined
                         )
-                    except jinja2.exceptions.TemplateSyntaxError:  # noqa: F841
+                    except jinja2.exceptions.TemplateSyntaxError as e:  # noqa: F841
                         self.conversation_starter = conversation_starter_content
             else:
                 self.logger.info(
@@ -145,10 +146,11 @@ class ConversationBot:
 
     async def generate_response(
         self,
-        session: AsyncHttpPipeline,
+        session: Union[AsyncHttpPipeline, AIProjectClient],
         conversation_history: List[ConversationTurn],
         max_history: int,
         turn_number: int = 0,
+        session_state: Optional[Dict[str, Any]] = None,
     ) -> Tuple[dict, dict, float, dict]:
         """
         Prompt the ConversationBot for a response.
@@ -175,6 +177,9 @@ class ConversationBot:
                 samples = [self.conversation_starter.render(**self.persona_template_args)]
             else:
                 samples = [self.conversation_starter]
+            jailbreak_string = self.persona_template_args.get("jailbreak_string", None)
+            if jailbreak_string:
+                samples = [f"{jailbreak_string} {samples[0]}"]
             time_taken = 0
 
             finish_reason = ["stop"]
@@ -251,10 +256,11 @@ class CallbackConversationBot(ConversationBot):
 
     async def generate_response(
         self,
-        session: AsyncHttpPipeline,
+        session: Union[AsyncHttpPipeline, AIProjectClient],
         conversation_history: List[Any],
         max_history: int,
         turn_number: int = 0,
+        session_state: Optional[Dict[str, Any]] = None,
     ) -> Tuple[dict, dict, float, dict]:
         chat_protocol_message = self._to_chat_protocol(
             self.user_template, conversation_history, self.user_template_parameters
@@ -262,7 +268,7 @@ class CallbackConversationBot(ConversationBot):
         msg_copy = copy.deepcopy(chat_protocol_message)
         result = {}
         start_time = time.time()
-        result = await self.callback(msg_copy)
+        result = await self.callback(msg_copy, session_state=session_state)
         end_time = time.time()
         if not result:
             result = {
@@ -271,8 +277,6 @@ class CallbackConversationBot(ConversationBot):
                 "id": None,
                 "template_parameters": {},
             }
-        self.logger.info("Using user provided callback returning response.")
-
         time_taken = end_time - start_time
         try:
             response = {
@@ -290,8 +294,6 @@ class CallbackConversationBot(ConversationBot):
                 blame=ErrorBlame.USER_ERROR,
             ) from exc
 
-        self.logger.info("Parsed callback response")
-
         return response, {}, time_taken, result
 
     # Bug 3354264: template is unused in the method - is this intentional?
@@ -308,9 +310,134 @@ class CallbackConversationBot(ConversationBot):
         }
 
 
+class MultiModalConversationBot(ConversationBot):
+    """MultiModal Conversation bot that uses a user provided callback to generate responses.
+
+    :param callback: The callback function to use to generate responses.
+    :type callback: Callable
+    :param user_template: The template to use for the request.
+    :type user_template: str
+    :param user_template_parameters: The template parameters to use for the request.
+    :type user_template_parameters: Dict
+    :param args: Optional arguments to pass to the parent class.
+    :type args: Any
+    :param kwargs: Optional keyword arguments to pass to the parent class.
+    :type kwargs: Any
+    """
+
+    def __init__(
+        self,
+        callback: Callable,
+        user_template: str,
+        user_template_parameters: TemplateParameters,
+        rai_client: Union[RAIClient, AIProjectClient],
+        *args,
+        **kwargs,
+    ) -> None:
+        self.callback = callback
+        self.user_template = user_template
+        self.user_template_parameters = user_template_parameters
+        self.rai_client = rai_client
+
+        super().__init__(*args, **kwargs)
+
+    async def generate_response(
+        self,
+        session: Union[AsyncHttpPipeline, AIProjectClient],
+        conversation_history: List[Any],
+        max_history: int,
+        turn_number: int = 0,
+        session_state: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[dict, dict, float, dict]:
+        previous_prompt = conversation_history[-1]
+        chat_protocol_message = await self._to_chat_protocol(conversation_history, self.user_template_parameters)
+
+        # replace prompt with {image.jpg} tags with image content data.
+        conversation_history.pop()
+        conversation_history.append(
+            ConversationTurn(
+                role=previous_prompt.role,
+                name=previous_prompt.name,
+                message=chat_protocol_message["messages"][0]["content"],
+                full_response=previous_prompt.full_response,
+                request=chat_protocol_message,
+            )
+        )
+        msg_copy = copy.deepcopy(chat_protocol_message)
+        result = {}
+        start_time = time.time()
+        result = await self.callback(msg_copy)
+        end_time = time.time()
+        if not result:
+            result = {
+                "messages": [{"content": "Callback did not return a response.", "role": "assistant"}],
+                "finish_reason": ["stop"],
+                "id": None,
+                "template_parameters": {},
+            }
+
+        time_taken = end_time - start_time
+        try:
+            response = {
+                "samples": [result["messages"][-1]["content"]],
+                "finish_reason": ["stop"],
+                "id": None,
+            }
+        except Exception as exc:
+            msg = "User provided callback does not conform to chat protocol standard."
+            raise EvaluationException(
+                message=msg,
+                internal_message=msg,
+                target=ErrorTarget.CALLBACK_CONVERSATION_BOT,
+                category=ErrorCategory.INVALID_VALUE,
+                blame=ErrorBlame.USER_ERROR,
+            ) from exc
+
+        return response, chat_protocol_message, time_taken, result
+
+    async def _to_chat_protocol(self, conversation_history, template_parameters):  # pylint: disable=unused-argument
+        messages = []
+
+        for _, m in enumerate(conversation_history):
+            if "image:" in m.message:
+                content = await self._to_multi_modal_content(m.message)
+                messages.append({"content": content, "role": m.role.value})
+            else:
+                messages.append({"content": m.message, "role": m.role.value})
+
+        return {
+            "template_parameters": template_parameters,
+            "messages": messages,
+            "$schema": "http://azureml/sdk-2-0/ChatConversation.json",
+        }
+
+    async def _to_multi_modal_content(self, text: str) -> list:
+        split_text = re.findall(r"[^{}]+|\{[^{}]*\}", text)
+        messages = [
+            text.strip("{}").replace("image:", "").strip() if text.startswith("{") else text for text in split_text
+        ]
+        contents = []
+        for msg in messages:
+            if msg.startswith("image_understanding/"):
+                if isinstance(self.rai_client, RAIClient):
+                    encoded_image = await self.rai_client.get_image_data(msg)
+                else:
+                    response = self.rai_client.red_teams.get_template_parameters_image(path=msg, stream="true")
+                    image_data = b"".join(response)
+                    encoded_image = base64.b64encode(image_data).decode("utf-8")
+
+                contents.append(
+                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{encoded_image}"}},
+                )
+            else:
+                contents.append({"type": "text", "text": msg})
+        return contents
+
+
 __all__ = [
     "ConversationRole",
     "ConversationBot",
     "CallbackConversationBot",
+    "MultiModalConversationBot",
     "ConversationTurn",
 ]

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

@@ -9,9 +9,9 @@ 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
 from azure.ai.evaluation.simulator._helpers._language_suffix_mapping import SUPPORTED_LANGUAGES_MAPPING
-
 from ..._http_utils import AsyncHttpPipeline
 from . import ConversationBot, ConversationTurn
+from azure.ai.evaluation._common.onedp._client import ProjectsClient as AIProjectClient
 
 
 def is_closing_message(response: Union[Dict, str], recursion_depth: int = 0) -> bool:
@@ -73,7 +73,7 @@ def is_closing_message_helper(response: str) -> bool:
 async def simulate_conversation(
     *,
     bots: List[ConversationBot],
-    session: AsyncHttpPipeline,
+    session: Union[AsyncHttpPipeline, AIProjectClient],
     language: SupportedLanguages,
     stopping_criteria: Callable[[str], bool] = is_closing_message,
     turn_limit: int = 10,
@@ -102,6 +102,7 @@ async def simulate_conversation(
     :rtype: Tuple[Optional[str], List[ConversationTurn]]
     """
 
+    session_state = {}
     # Read the first prompt.
     (first_response, request, _, full_response) = await bots[0].generate_response(
         session=session,
@@ -150,7 +151,10 @@ async def simulate_conversation(
                 conversation_history=conversation_history,
                 max_history=history_limit,
                 turn_number=current_turn,
+                session_state=session_state,
             )
+            if "session_state" in full_response and full_response["session_state"] is not None:
+                session_state.update(full_response["session_state"])
 
             # check if conversation id is null, which means conversation starter was used. use id from next turn
             if conversation_id is None and "id" in response:

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

@@ -12,7 +12,7 @@ OUTPUT_FILE = "openai_api_response.jsonl"
 
 # Azure endpoint constants
 AZUREML_TOKEN_SCOPE = "https://ml.azure.com"
-COGNITIVE_SERVICES_TOKEN_SCOPE = "https://cognitiveservices.azure.com/"
+COGNITIVE_SERVICES_TOKEN_SCOPE = "https://ai.azure.com/"
 AZURE_TOKEN_REFRESH_INTERVAL = 600  # seconds
 AZURE_ENDPOINT_DOMAIN_VALID_PATTERN_RE = (
     r"^(?=.{1,255}$)(?!-)[a-zA-Z0-9-]{1,63}(?<!-)"