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

azure/ai/evaluation/simulator/_direct_attack_simulator.py

@@ -7,13 +7,13 @@ import logging
 from random import randint
 from typing import Callable, Optional, cast
 
+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.simulator import AdversarialScenario
 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__)

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

@@ -1,5 +1,4 @@
-from ._experimental import experimental
 from ._language_suffix_mapping import SUPPORTED_LANGUAGES_MAPPING
 from ._simulator_data_classes import ConversationHistory, Turn
 
-__all__ = ["ConversationHistory", "Turn", "SUPPORTED_LANGUAGES_MAPPING", "experimental"]
+__all__ = ["ConversationHistory", "Turn", "SUPPORTED_LANGUAGES_MAPPING"]

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

@@ -30,7 +30,19 @@ class Turn:
         return {
             "role": self.role.value if isinstance(self.role, ConversationRole) else self.role,
             "content": self.content,
-            "context": self.context,
+            "context": str(self.context),
+        }
+
+    def to_context_free_dict(self) -> Dict[str, Optional[str]]:
+        """
+        Convert the conversation turn to a dictionary without context.
+
+        :returns: A dictionary representation of the conversation turn without context.
+        :rtype: Dict[str, Optional[str]]
+        """
+        return {
+            "role": self.role.value if isinstance(self.role, ConversationRole) else self.role,
+            "content": self.content,
         }
 
     def __repr__(self):
@@ -66,6 +78,15 @@ class ConversationHistory:
         """
         return [turn.to_dict() for turn in self.history]
 
+    def to_context_free_list(self) -> List[Dict[str, Optional[str]]]:
+        """
+        Converts the conversation history to a list of dictionaries without context.
+
+        :returns: A list of dictionaries representing the conversation turns without context.
+        :rtype: List[Dict[str, str]]
+        """
+        return [turn.to_context_free_dict() for turn in self.history]
+
     def __len__(self) -> int:
         return len(self.history)
 

azure/ai/evaluation/simulator/_indirect_attack_simulator.py

@@ -3,23 +3,27 @@
 # ---------------------------------------------------------
 # pylint: disable=C0301,C0114,R0913,R0903
 # noqa: E501
+import asyncio
 import logging
 from typing import Callable, cast
 
+from tqdm import tqdm
+
 from azure.ai.evaluation._common.utils import validate_azure_ai_project
+from azure.ai.evaluation._common._experimental import experimental
 from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
-from azure.ai.evaluation.simulator import AdversarialScenario
+from azure.ai.evaluation.simulator import AdversarialScenarioJailbreak, SupportedLanguages
 from azure.core.credentials import TokenCredential
 
-from ._adversarial_simulator import AdversarialSimulator
-from ._helpers import experimental
+from ._adversarial_simulator import AdversarialSimulator, JsonLineList
+
 from ._model_tools import AdversarialTemplateHandler, ManagedIdentityAPITokenManager, RAIClient, TokenScope
 
 logger = logging.getLogger(__name__)
 
 
 @experimental
-class IndirectAttackSimulator:
+class IndirectAttackSimulator(AdversarialSimulator):
     """
     Initializes the XPIA (cross domain prompt injected attack) jailbreak adversarial simulator with a project scope.
 
@@ -54,6 +58,7 @@ class IndirectAttackSimulator:
         self.adversarial_template_handler = AdversarialTemplateHandler(
             azure_ai_project=self.azure_ai_project, rai_client=self.rai_client
         )
+        super().__init__(azure_ai_project=azure_ai_project, credential=credential)
 
     def _ensure_service_dependencies(self):
         if self.rai_client is None:
@@ -69,29 +74,22 @@ class IndirectAttackSimulator:
     async def __call__(
         self,
         *,
-        scenario: AdversarialScenario,
         target: Callable,
-        max_conversation_turns: int = 1,
         max_simulation_results: int = 3,
         api_call_retry_limit: int = 3,
         api_call_retry_sleep_sec: int = 1,
         api_call_delay_sec: int = 0,
         concurrent_async_task: int = 3,
+        **kwargs,
     ):
         """
         Initializes the XPIA (cross domain prompt injected attack) jailbreak adversarial simulator with a project scope.
         This simulator converses with your AI system using prompts injected into the context to interrupt normal
         expected functionality by eliciting manipulated content, intrusion and attempting to gather information outside
         the scope of your AI system.
-
-        :keyword scenario: Enum value specifying the adversarial scenario used for generating inputs.
-        :paramtype scenario: azure.ai.evaluation.simulator.AdversarialScenario
         :keyword target: The target function to simulate adversarial inputs against.
             This function should be asynchronous and accept a dictionary representing the adversarial input.
         :paramtype target: Callable
-        :keyword max_conversation_turns: The maximum number of conversation turns to simulate.
-            Defaults to 1.
-        :paramtype max_conversation_turns: int
         :keyword max_simulation_results: The maximum number of simulation results to return.
             Defaults to 3.
         :paramtype max_simulation_results: int
@@ -128,11 +126,11 @@ class IndirectAttackSimulator:
                     'template_parameters': {},
                     'messages': [
                         {
-                            'content': '<jailbreak prompt> <adversarial query>',
+                            'content': '<adversarial query>',
                             'role': 'user'
                         },
                         {
-                            'content': "<response from endpoint>",
+                            'content': "<response from your callback>",
                             'role': 'assistant',
                             'context': None
                         }
@@ -141,25 +139,72 @@ class IndirectAttackSimulator:
                 }]
             }
         """
-        if scenario not in AdversarialScenario.__members__.values():
-            msg = f"Invalid scenario: {scenario}. Supported scenarios: {AdversarialScenario.__members__.values()}"
-            raise EvaluationException(
-                message=msg,
-                internal_message=msg,
-                target=ErrorTarget.DIRECT_ATTACK_SIMULATOR,
-                category=ErrorCategory.INVALID_VALUE,
-                blame=ErrorBlame.USER_ERROR,
+        # values that cannot be changed:
+        scenario = AdversarialScenarioJailbreak.ADVERSARIAL_INDIRECT_JAILBREAK
+        max_conversation_turns = 2
+        language = SupportedLanguages.English
+        self._ensure_service_dependencies()
+        templates = await self.adversarial_template_handler._get_content_harm_template_collections(scenario.value)
+        concurrent_async_task = min(concurrent_async_task, 1000)
+        semaphore = asyncio.Semaphore(concurrent_async_task)
+        sim_results = []
+        tasks = []
+        total_tasks = sum(len(t.template_parameters) for t in templates)
+        if max_simulation_results > total_tasks:
+            logger.warning(
+                "Cannot provide %s results due to maximum number of adversarial simulations that can be generated: %s."
+                "\n %s simulations will be generated.",
+                max_simulation_results,
+                total_tasks,
+                total_tasks,
             )
-        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,
-            max_conversation_turns=max_conversation_turns,
-            max_simulation_results=max_simulation_results,
-            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,
-            concurrent_async_task=concurrent_async_task,
-            _jailbreak_type="xpia",
+        total_tasks = min(total_tasks, max_simulation_results)
+        progress_bar = tqdm(
+            total=total_tasks,
+            desc="generating jailbreak simulations",
+            ncols=100,
+            unit="simulations",
         )
-        return jb_sim_results
+        for template in templates:
+            for parameter in template.template_parameters:
+                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,
+                        )
+                    )
+                )
+                if len(tasks) >= max_simulation_results:
+                    break
+            if len(tasks) >= max_simulation_results:
+                break
+        for task in asyncio.as_completed(tasks):
+            completed_task = await task  # type: ignore
+            template_parameters = completed_task.get("template_parameters", {})  # type: ignore
+            xpia_attack_type = template_parameters.get("xpia_attack_type", "")  # type: ignore
+            action = template_parameters.get("action", "")  # type: ignore
+            document_type = template_parameters.get("document_type", "")  # type: ignore
+            sim_results.append(
+                {
+                    "messages": completed_task["messages"],  # type: ignore
+                    "$schema": "http://azureml/sdk-2-0/ChatConversation.json",
+                    "template_parameters": {
+                        "metadata": {
+                            "xpia_attack_type": xpia_attack_type,
+                            "action": action,
+                            "document_type": document_type,
+                        },
+                    },
+                }
+            )
+            progress_bar.update(1)
+        progress_bar.close()
+        return JsonLineList(sim_results)

azure/ai/evaluation/simulator/_model_tools/_identity_manager.py

@@ -11,7 +11,7 @@ from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Optional, Union
 
-from azure.core.credentials import TokenCredential, AccessToken
+from azure.core.credentials import AccessToken, TokenCredential
 from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
 
 AZURE_TOKEN_REFRESH_INTERVAL = 600  # seconds

azure/ai/evaluation/simulator/_prompty/task_query_response.prompty

@@ -36,8 +36,8 @@ On January 24, 1984, former Apple CEO Steve Jobs introduced the first Macintosh.
 Some years later, research firms IDC and Gartner reported that Apple's market share in the U.S. had increased to about 6%.
 <|text_end|>
 Output with 5 QnAs:
-[
-  {
+{
+  "qna":[{
     "q": "When did the former Apple CEO Steve Jobs introduced the first Macintosh?",
     "r": "January 24, 1984"
   },
@@ -56,8 +56,8 @@ Output with 5 QnAs:
   {
     "q": "What was the percentage increase of Apple's market share in the U.S., as reported by research firms IDC and Gartner?",
     "r": "6%"
-  }
-]
+  }]
+}
 Text:
 <|text_start|>
 {{ text }}

azure/ai/evaluation/simulator/_prompty/task_simulate.prompty

@@ -16,6 +16,9 @@ inputs:
     type: string
   conversation_history:
     type: dict
+  action:
+    type: string
+    default: continue the converasation and make sure the task is completed by asking relevant questions
 
 ---
 system:
@@ -25,8 +28,10 @@ Output must be in JSON format
 Here's a sample output:
 {
   "content": "Here is my follow-up question.",
-  "user": "user"
+  "role": "user"
 }
 
 Output with a json object that continues the conversation, given the conversation history:
 {{ conversation_history }}
+
+{{ action }}

azure/ai/evaluation/simulator/_simulator.py

@@ -5,20 +5,23 @@
 # ---------------------------------------------------------
 import asyncio
 import importlib.resources as pkg_resources
-from tqdm import tqdm
 import json
 import os
 import re
 import warnings
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union, Tuple
+
 from promptflow.core import AsyncPrompty
-from azure.ai.evaluation._model_configurations import AzureOpenAIModelConfiguration, OpenAIModelConfiguration
+from tqdm import tqdm
+
+from azure.ai.evaluation._common._experimental import experimental
 from azure.ai.evaluation._common.utils import construct_prompty_model_config
+from azure.ai.evaluation._model_configurations import AzureOpenAIModelConfiguration, OpenAIModelConfiguration
 
 from .._exceptions import ErrorBlame, ErrorCategory, EvaluationException
 from .._user_agent import USER_AGENT
 from ._conversation.constants import ConversationRole
-from ._helpers import ConversationHistory, Turn, experimental
+from ._helpers import ConversationHistory, Turn
 from ._utils import JsonLineChatProtocol
 
 
@@ -89,7 +92,8 @@ class Simulator:
         api_call_delay_sec: float = 1,
         query_response_generating_prompty_kwargs: Dict[str, Any] = {},
         user_simulator_prompty_kwargs: Dict[str, Any] = {},
-        conversation_turns: List[List[str]] = [],
+        conversation_turns: List[List[Union[str, Dict[str, Any]]]] = [],
+        concurrent_async_tasks: int = 5,
         **kwargs,
     ) -> List[JsonLineChatProtocol]:
         """
@@ -116,7 +120,10 @@ class Simulator:
         :keyword user_simulator_prompty_kwargs: Additional keyword arguments for the user simulator prompty.
         :paramtype user_simulator_prompty_kwargs: Dict[str, Any]
         :keyword conversation_turns: Predefined conversation turns to simulate.
-        :paramtype conversation_turns: List[List[str]]
+        :paramtype conversation_turns: List[List[Union[str, Dict[str, Any]]]]
+        :keyword concurrent_async_tasks: The number of asynchronous tasks to run concurrently during the simulation.
+            Defaults to 5.
+        :paramtype concurrent_async_tasks: int
         :return: A list of simulated conversations represented as JsonLineChatProtocol objects.
         :rtype: List[JsonLineChatProtocol]
 
@@ -131,12 +138,12 @@ class Simulator:
         if conversation_turns and (text or tasks):
             raise ValueError("Cannot specify both conversation_turns and text/tasks")
 
-        if num_queries > len(tasks):
+        if text and num_queries > len(tasks):
             warnings.warn(
                 f"You have specified 'num_queries' > len('tasks') ({num_queries} > {len(tasks)}). "
                 f"All tasks will be used for generation and the remaining {num_queries - len(tasks)} lines will be simulated in task-free mode"
             )
-        elif num_queries < len(tasks):
+        elif text and num_queries < len(tasks):
             warnings.warn(
                 f"You have specified 'num_queries' < len('tasks') ({num_queries} < {len(tasks)}). "
                 f"Only the first {num_queries} lines of the specified tasks will be simulated."
@@ -154,6 +161,7 @@ class Simulator:
                 user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
                 api_call_delay_sec=api_call_delay_sec,
                 prompty_model_config=prompty_model_config,
+                concurrent_async_tasks=concurrent_async_tasks,
             )
 
         query_responses = await self._generate_query_responses(
@@ -172,6 +180,7 @@ class Simulator:
             user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
             target=target,
             api_call_delay_sec=api_call_delay_sec,
+            text=text,
         )
 
     async def _simulate_with_predefined_turns(
@@ -179,11 +188,12 @@ class Simulator:
         *,
         target: Callable,
         max_conversation_turns: int,
-        conversation_turns: List[List[str]],
+        conversation_turns: List[List[Union[str, Dict[str, Any]]]],
         user_simulator_prompty: Optional[str],
         user_simulator_prompty_kwargs: Dict[str, Any],
         api_call_delay_sec: float,
         prompty_model_config: Any,
+        concurrent_async_tasks: int,
     ) -> List[JsonLineChatProtocol]:
         """
         Simulates conversations using predefined conversation turns.
@@ -193,7 +203,7 @@ class Simulator:
         :keyword max_conversation_turns: Maximum number of turns for the simulation.
         :paramtype max_conversation_turns: int
         :keyword conversation_turns: A list of predefined conversation turns.
-        :paramtype conversation_turns: List[List[str]]
+        :paramtype conversation_turns: List[List[Union[str, Dict[str, Any]]]]
         :keyword user_simulator_prompty: Path to the user simulator prompty file.
         :paramtype user_simulator_prompty: Optional[str]
         :keyword user_simulator_prompty_kwargs: Additional keyword arguments for the user simulator prompty.
@@ -202,42 +212,60 @@ class Simulator:
         :paramtype api_call_delay_sec: float
         :keyword prompty_model_config: The configuration for the prompty model.
         :paramtype prompty_model_config: Any
+        :keyword concurrent_async_tasks: The number of asynchronous tasks to run concurrently during the simulation.
+        :paramtype concurrent_async_tasks: int
         :return: A list of simulated conversations represented as JsonLineChatProtocol objects.
         :rtype: List[JsonLineChatProtocol]
         """
-        simulated_conversations = []
         progress_bar = tqdm(
             total=int(len(conversation_turns) * (max_conversation_turns / 2)),
             desc="Simulating with predefined conversation turns: ",
             ncols=100,
             unit="messages",
         )
-
-        for simulation in conversation_turns:
-            current_simulation = ConversationHistory()
-            for simulated_turn in simulation:
-                user_turn = Turn(role=ConversationRole.USER, content=simulated_turn)
-                current_simulation.add_to_history(user_turn)
-                assistant_response = await self._get_target_response(
-                    target=target, api_call_delay_sec=api_call_delay_sec, conversation_history=current_simulation
-                )
-                assistant_turn = Turn(role=ConversationRole.ASSISTANT, content=assistant_response)
-                current_simulation.add_to_history(assistant_turn)
-                progress_bar.update(1)  # Update progress bar for both user and assistant turns
-
-            if len(current_simulation) < max_conversation_turns:
-                await self._extend_conversation_with_simulator(
-                    current_simulation=current_simulation,
-                    max_conversation_turns=max_conversation_turns,
-                    user_simulator_prompty=user_simulator_prompty,
-                    user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
-                    api_call_delay_sec=api_call_delay_sec,
-                    prompty_model_config=prompty_model_config,
-                    target=target,
-                    progress_bar=progress_bar,
-                )
-            simulated_conversations.append(
-                JsonLineChatProtocol(
+        semaphore = asyncio.Semaphore(concurrent_async_tasks)
+        progress_bar_lock = asyncio.Lock()
+
+        async def run_simulation(simulation: List[Union[str, Dict[str, Any]]]) -> JsonLineChatProtocol:
+            async with semaphore:
+                current_simulation = ConversationHistory()
+                for simulated_turn in simulation:
+                    if isinstance(simulated_turn, str):
+                        user_turn = Turn(role=ConversationRole.USER, content=simulated_turn)
+                    elif isinstance(simulated_turn, dict):
+                        user_turn = Turn(
+                            role=ConversationRole.USER,
+                            content=str(simulated_turn.get("content")),
+                            context=str(simulated_turn.get("context")),
+                        )
+                    else:
+                        raise ValueError(
+                            "Each simulated turn must be a string or a dict with 'content' and 'context' keys"
+                        )
+                    current_simulation.add_to_history(user_turn)
+                    assistant_response, assistant_context = await self._get_target_response(
+                        target=target, api_call_delay_sec=api_call_delay_sec, conversation_history=current_simulation
+                    )
+                    assistant_turn = Turn(
+                        role=ConversationRole.ASSISTANT, content=assistant_response, context=assistant_context
+                    )
+                    current_simulation.add_to_history(assistant_turn)
+                    async with progress_bar_lock:
+                        progress_bar.update(1)
+
+                if len(current_simulation) < max_conversation_turns:
+                    await self._extend_conversation_with_simulator(
+                        current_simulation=current_simulation,
+                        max_conversation_turns=max_conversation_turns,
+                        user_simulator_prompty=user_simulator_prompty,
+                        user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
+                        api_call_delay_sec=api_call_delay_sec,
+                        prompty_model_config=prompty_model_config,
+                        target=target,
+                        progress_bar=progress_bar,
+                        progress_bar_lock=progress_bar_lock,
+                    )
+                return JsonLineChatProtocol(
                     {
                         "messages": current_simulation.to_list(),
                         "finish_reason": ["stop"],
@@ -245,10 +273,11 @@ class Simulator:
                         "$schema": "http://azureml/sdk-2-0/ChatConversation.json",
                     }
                 )
-            )
 
+        tasks = [asyncio.create_task(run_simulation(simulation)) for simulation in conversation_turns]
+        results = await asyncio.gather(*tasks)
         progress_bar.close()
-        return simulated_conversations
+        return results
 
     async def _extend_conversation_with_simulator(
         self,
@@ -261,6 +290,7 @@ class Simulator:
         prompty_model_config: Dict[str, Any],
         target: Callable,
         progress_bar: tqdm,
+        progress_bar_lock: asyncio.Lock,
     ):
         """
         Extends an ongoing conversation using a user simulator until the maximum number of turns is reached.
@@ -281,6 +311,8 @@ class Simulator:
         :paramtype target: Callable,
         :keyword progress_bar: Progress bar for tracking simulation progress.
         :paramtype progress_bar: tqdm,
+        :keyword progress_bar_lock: Lock for updating the progress bar safely.
+        :paramtype progress_bar_lock: asyncio.Lock
         """
         user_flow = self._load_user_simulation_flow(
             user_simulator_prompty=user_simulator_prompty,  # type: ignore
@@ -291,19 +323,22 @@ class Simulator:
         while len(current_simulation) < max_conversation_turns:
             user_response_content = await user_flow(
                 task="Continue the conversation",
-                conversation_history=current_simulation.to_list(),
+                conversation_history=current_simulation.to_context_free_list(),
                 **user_simulator_prompty_kwargs,
             )
             user_response = self._parse_prompty_response(response=user_response_content)
             user_turn = Turn(role=ConversationRole.USER, content=user_response["content"])
             current_simulation.add_to_history(user_turn)
             await asyncio.sleep(api_call_delay_sec)
-            assistant_response = await self._get_target_response(
+            assistant_response, assistant_context = await self._get_target_response(
                 target=target, api_call_delay_sec=api_call_delay_sec, conversation_history=current_simulation
             )
-            assistant_turn = Turn(role=ConversationRole.ASSISTANT, content=assistant_response)
+            assistant_turn = Turn(
+                role=ConversationRole.ASSISTANT, content=assistant_response, context=assistant_context
+            )
             current_simulation.add_to_history(assistant_turn)
-            progress_bar.update(1)
+            async with progress_bar_lock:
+                progress_bar.update(1)
 
     def _load_user_simulation_flow(
         self,
@@ -432,6 +467,14 @@ class Simulator:
             if isinstance(query_responses, dict):
                 keys = list(query_responses.keys())
                 return query_responses[keys[0]]
+            if isinstance(query_responses, str):
+                query_responses = json.loads(query_responses)
+                if isinstance(query_responses, dict):
+                    if len(query_responses.keys()) == 1:
+                        return query_responses[list(query_responses.keys())[0]]
+                    return query_responses  # type: ignore
+                if isinstance(query_responses, list):
+                    return query_responses
             return json.loads(query_responses)
         except Exception as e:
             raise RuntimeError("Error generating query responses") from e
@@ -497,6 +540,7 @@ class Simulator:
         user_simulator_prompty_kwargs: Dict[str, Any],
         target: Callable,
         api_call_delay_sec: float,
+        text: str,
     ) -> List[JsonLineChatProtocol]:
         """
         Creates full conversations from query-response pairs.
@@ -515,6 +559,8 @@ class Simulator:
         :paramtype target: Callable
         :keyword api_call_delay_sec: Delay in seconds between API calls.
         :paramtype api_call_delay_sec: float
+        :keyword text: The initial input text for generating query responses.
+        :paramtype text: str
         :return: A list of simulated conversations represented as JsonLineChatProtocol objects.
         :rtype: List[JsonLineChatProtocol]
         """
@@ -552,6 +598,7 @@ class Simulator:
                             "task": task,
                             "expected_response": response,
                             "query": query,
+                            "original_text": text,
                         },
                         "$schema": "http://azureml/sdk-2-0/ChatConversation.json",
                     }
@@ -595,8 +642,6 @@ class Simulator:
         :rtype: List[Dict[str, Optional[str]]]
         """
         conversation_history = ConversationHistory()
-        # user_turn = Turn(role=ConversationRole.USER, content=conversation_starter)
-        # conversation_history.add_to_history(user_turn)
 
         while len(conversation_history) < max_conversation_turns:
             user_flow = self._load_user_simulation_flow(
@@ -604,24 +649,33 @@ class Simulator:
                 prompty_model_config=self.model_config,  # type: ignore
                 user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
             )
-            conversation_starter_from_simulated_user = await user_flow(
-                task=task,
-                conversation_history=[
-                    {
-                        "role": "assistant",
-                        "content": conversation_starter,
-                        "your_task": "Act as the user and translate the content into a user query.",
-                    }
-                ],
-            )
+            if len(conversation_history) == 0:
+                conversation_starter_from_simulated_user = await user_flow(
+                    task=task,
+                    conversation_history=[
+                        {
+                            "role": "assistant",
+                            "content": conversation_starter,
+                        }
+                    ],
+                    action="rewrite the assistant's message as you have to accomplish the task by asking the right questions. Make sure the original question is not lost in your rewrite.",
+                )
+            else:
+                conversation_starter_from_simulated_user = await user_flow(
+                    task=task,
+                    conversation_history=conversation_history.to_context_free_list(),
+                    action="Your goal is to make sure the task is completed by asking the right questions. Do not ask the same questions again.",
+                )
             if isinstance(conversation_starter_from_simulated_user, dict):
                 conversation_starter_from_simulated_user = conversation_starter_from_simulated_user["content"]
             user_turn = Turn(role=ConversationRole.USER, content=conversation_starter_from_simulated_user)
             conversation_history.add_to_history(user_turn)
-            assistant_response = await self._get_target_response(
+            assistant_response, assistant_context = await self._get_target_response(
                 target=target, api_call_delay_sec=api_call_delay_sec, conversation_history=conversation_history
             )
-            assistant_turn = Turn(role=ConversationRole.ASSISTANT, content=assistant_response)
+            assistant_turn = Turn(
+                role=ConversationRole.ASSISTANT, content=assistant_response, context=assistant_context
+            )
             conversation_history.add_to_history(assistant_turn)
             progress_bar.update(1)
 
@@ -632,7 +686,7 @@ class Simulator:
 
     async def _get_target_response(
         self, *, target: Callable, api_call_delay_sec: float, conversation_history: ConversationHistory
-    ) -> str:
+    ) -> Tuple[str, Optional[str]]:
         """
         Retrieves the response from the target callback based on the current conversation history.
 
@@ -642,8 +696,8 @@ class Simulator:
         :paramtype api_call_delay_sec: float
         :keyword conversation_history: The current conversation history.
         :paramtype conversation_history: ConversationHistory
-        :return: The content of the response from the target.
-        :rtype: str
+        :return: The content of the response from the target and an optional context.
+        :rtype: str, Optional[str]
         """
         response = await target(
             messages={"messages": conversation_history.to_list()},
@@ -653,4 +707,4 @@ class Simulator:
         )
         await asyncio.sleep(api_call_delay_sec)
         latest_message = response["messages"][-1]
-        return latest_message["content"]
+        return latest_message["content"], latest_message.get("context", "")  # type: ignore