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

azure/ai/evaluation/simulator/_simulator.py

@@ -1,5 +1,5 @@
 # flake8: noqa
-# pylint: disable=W0102,W0613,R0914,C0301,E0401,E0611,C0114,R0913,E0702,R0903,C0411
+# pylint: disable=W0102,W0613,R0914,C0301,E0401,E0611
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
@@ -9,19 +9,17 @@ import json
 import os
 import re
 import warnings
-from typing import Any, Callable, Dict, List, Optional, Union, Tuple
+from typing import Any, Callable, Dict, List, Optional, Union
 
-from promptflow.core import AsyncPrompty
+from promptflow.client import load_flow
+from promptflow.core import AzureOpenAIModelConfiguration, Flow
 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
+from ._helpers import ConversationHistory, Turn, experimental
+
+# from ._tracing import monitor_task_simulator
 from ._utils import JsonLineChatProtocol
 
 
@@ -29,77 +27,53 @@ from ._utils import JsonLineChatProtocol
 class Simulator:
     """
     Simulator for generating synthetic conversations.
-
-    :param model_config: A dictionary defining the configuration for the model. Acceptable types are AzureOpenAIModelConfiguration and OpenAIModelConfiguration.
-    :type model_config: Union[~azure.ai.evaluation.AzureOpenAIModelConfiguration, ~azure.ai.evaluation.OpenAIModelConfiguration]
-    :raises ValueError: If the model_config does not contain the required keys or any value is None.
-
-    .. admonition:: Example:
-
-        .. literalinclude:: ../samples/evaluation_samples_simulate.py
-            :start-after: [START nonadversarial_simulator]
-            :end-before: [END nonadversarial_simulator]
-            :language: python
-            :dedent: 8
-            :caption: Run a Simulator for 2 queries and 4 conversation turns.
     """
 
-    def __init__(self, model_config: Union[AzureOpenAIModelConfiguration, OpenAIModelConfiguration]):
-        self._validate_model_config(model_config)
-        self.model_config = model_config
-        if "api_version" not in self.model_config:
-            self.model_config["api_version"] = "2024-06-01"  # type: ignore
+    def __init__(self, azure_ai_project: Dict[str, Any], credential: Optional[Any] = None):
+        """
+        Initializes the task simulator with a project scope.
+
+        :param azure_ai_project: A dictionary defining the scope of the project, including keys such as
+                                "subscription_id", "resource_group_name", and "project_name".
+        :param credential: Azure credentials to authenticate the user. If None, the default credentials are used.
+        :paramtype credential: Optional[Any]
+        :raises ValueError: If the azure_ai_project does not contain the required keys or any value is None.
+        """
+        self._validate_project_config(azure_ai_project)
+        self.azure_ai_project = azure_ai_project
+        self.azure_ai_project["api_version"] = "2024-02-15-preview"
+        self.credential = credential
 
     @staticmethod
-    def _validate_model_config(model_config: Any):
+    def _validate_project_config(azure_ai_project: Dict[str, Any]):
         """
-        Validates the model_config to ensure all required keys are present and have non-None values.
-        If 'type' is not specified, it will attempt to infer the type based on the keys present.
+        Validates the azure_ai_project configuration to ensure all required keys are present and have non-None values.
 
-        :param model_config: The model configuration dictionary.
-        :type model_config: Dict[str, Any]
+        :param azure_ai_project: The Azure AI project configuration dictionary.
+        :type azure_ai_project: Dict[str, Any]
         :raises ValueError: If required keys are missing or any of the values are None.
         """
-        # Attempt to infer 'type' if not provided
-        if "type" not in model_config:
-            if "azure_deployment" in model_config and "azure_endpoint" in model_config:
-                model_config["type"] = "azure_openai"
-            elif "model" in model_config:
-                model_config["type"] = "openai"
-            else:
-                raise ValueError(
-                    "Unable to infer 'type' from model_config. Please specify 'type' as 'azure_openai' or 'openai'."
-                )
-
-        if model_config["type"] == "azure_openai":
-            required_keys = ["azure_deployment", "azure_endpoint"]
-        elif model_config["type"] == "openai":
-            required_keys = ["api_key", "model"]
-        else:
-            raise ValueError("model_config 'type' must be 'azure_openai' or 'openai'.")
-
-        missing_keys = [key for key in required_keys if key not in model_config]
-        if missing_keys:
-            raise ValueError(f"model_config is missing required keys: {', '.join(missing_keys)}")
-        none_keys = [key for key in required_keys if model_config.get(key) is None]
-        if none_keys:
-            raise ValueError(f"The following keys in model_config must not be None: {', '.join(none_keys)}")
+        required_keys = ["subscription_id", "resource_group_name", "project_name"]
+        if not all(key in azure_ai_project for key in required_keys):
+            raise ValueError(f"azure_ai_project must contain keys: {', '.join(required_keys)}")
+        if not all(azure_ai_project[key] for key in required_keys):
+            raise ValueError("subscription_id, resource_group_name, and project_name must not be None")
 
+    # @monitor_task_simulator
     async def __call__(
         self,
         *,
         target: Callable,
         max_conversation_turns: int = 5,
-        tasks: List[str] = [],
+        tasks: List[Dict] = [],
         text: str = "",
         num_queries: int = 5,
         query_response_generating_prompty: Optional[str] = None,
         user_simulator_prompty: Optional[str] = None,
         api_call_delay_sec: float = 1,
-        query_response_generating_prompty_options: Dict[str, Any] = {},
-        user_simulator_prompty_options: Dict[str, Any] = {},
-        conversation_turns: List[List[Union[str, Dict[str, Any]]]] = [],
-        concurrent_async_tasks: int = 5,
+        query_response_generating_prompty_kwargs: Dict[str, Any] = {},
+        user_simulator_prompty_kwargs: Dict[str, Any] = {},
+        conversation_turns: List[List[str]] = [],
         **kwargs,
     ) -> List[JsonLineChatProtocol]:
         """
@@ -121,15 +95,12 @@ class Simulator:
         :paramtype user_simulator_prompty: Optional[str]
         :keyword api_call_delay_sec: Delay in seconds between API calls.
         :paramtype api_call_delay_sec: float
-        :keyword query_response_generating_prompty_options: Additional keyword arguments for the query response generating prompty.
-        :paramtype query_response_generating_prompty_options: Dict[str, Any]
-        :keyword user_simulator_prompty_options: Additional keyword arguments for the user simulator prompty.
-        :paramtype user_simulator_prompty_options: Dict[str, Any]
+        :keyword query_response_generating_prompty_kwargs: Additional keyword arguments for the query response generating prompty.
+        :paramtype query_response_generating_prompty_kwargs: Dict[str, Any]
+        :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[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
+        :paramtype conversation_turns: List[List[str]]
         :return: A list of simulated conversations represented as JsonLineChatProtocol objects.
         :rtype: List[JsonLineChatProtocol]
 
@@ -138,18 +109,18 @@ class Simulator:
 
         Modes:
         - Task-Free Mode: When only num_queries is specified and tasks is not, the method generates num_queries x max_conversation_turns lines of simulated data grounded in the context of the text.
-        - Task-Specific Mode: When both num_queries and tasks are specified, the method generates lines of simulated data based on the tasks. If num_queries > len(tasks), the remaining lines will be simulated in task-free mode. If num_queries < len(tasks), only the first num_queries tasks are used.
+        - Task-Specific Mode: When both num_queries and tasks are specified, the method generates lines of simulated data based on the tasks. If num_queries > len(tasks), the remaining lines are simulated in task-free mode. If num_queries < len(tasks), only the first num_queries tasks are used.
         - Conversation Starter Mode: When conversation_turns are specified, the method starts each conversation with the user-specified queries and then follows the conversation history for the remaining turns.
         """
         if conversation_turns and (text or tasks):
             raise ValueError("Cannot specify both conversation_turns and text/tasks")
 
-        if text and num_queries > len(tasks):
+        if 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 text and num_queries < len(tasks):
+        elif 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."
@@ -157,49 +128,60 @@ class Simulator:
         num_queries = min(num_queries, len(tasks))
         max_conversation_turns *= 2  # account for both user and assistant turns
 
-        prompty_model_config = self.model_config
+        prompty_model_config = self._build_prompty_model_config()
+
         if conversation_turns:
             return await self._simulate_with_predefined_turns(
                 target=target,
                 max_conversation_turns=max_conversation_turns,
                 conversation_turns=conversation_turns,
                 user_simulator_prompty=user_simulator_prompty,
-                user_simulator_prompty_options=user_simulator_prompty_options,
+                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(
             text=text,
             num_queries=num_queries,
             query_response_generating_prompty=query_response_generating_prompty,
-            query_response_generating_prompty_options=query_response_generating_prompty_options,
+            query_response_generating_prompty_kwargs=query_response_generating_prompty_kwargs,
             prompty_model_config=prompty_model_config,
             **kwargs,
         )
+
         return await self._create_conversations_from_query_responses(
             query_responses=query_responses,
             max_conversation_turns=max_conversation_turns,
             tasks=tasks,
             user_simulator_prompty=user_simulator_prompty,
-            user_simulator_prompty_options=user_simulator_prompty_options,
+            user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
             target=target,
             api_call_delay_sec=api_call_delay_sec,
-            text=text,
         )
 
+    def _build_prompty_model_config(self) -> Dict[str, Any]:
+        """
+        Constructs the configuration for the prompty model.
+
+        :return: A dictionary containing the prompty model configuration, including API version and user agent headers if applicable.
+        :rtype: Dict[str, Any]
+        """
+        config = {"configuration": self.azure_ai_project}
+        if USER_AGENT and isinstance(self.azure_ai_project, AzureOpenAIModelConfiguration):
+            config.update({"parameters": {"extra_headers": {"x-ms-useragent": USER_AGENT}}})
+        return config
+
     async def _simulate_with_predefined_turns(
         self,
         *,
         target: Callable,
         max_conversation_turns: int,
-        conversation_turns: List[List[Union[str, Dict[str, Any]]]],
+        conversation_turns: List[List[str]],
         user_simulator_prompty: Optional[str],
-        user_simulator_prompty_options: Dict[str, Any],
+        user_simulator_prompty_kwargs: Dict[str, Any],
         api_call_delay_sec: float,
-        prompty_model_config: Any,
-        concurrent_async_tasks: int,
+        prompty_model_config: Dict[str, Any],
     ) -> List[JsonLineChatProtocol]:
         """
         Simulates conversations using predefined conversation turns.
@@ -209,81 +191,54 @@ 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[Union[str, Dict[str, Any]]]]
+        :paramtype conversation_turns: List[List[str]]
         :keyword user_simulator_prompty: Path to the user simulator prompty file.
         :paramtype user_simulator_prompty: Optional[str]
-        :keyword user_simulator_prompty_options: Additional keyword arguments for the user simulator prompty.
-        :paramtype user_simulator_prompty_options: Dict[str, Any]
+        :keyword user_simulator_prompty_kwargs: Additional keyword arguments for the user simulator prompty.
+        :paramtype user_simulator_prompty_kwargs: Dict[str, Any]
         :keyword api_call_delay_sec: Delay in seconds between API calls.
         :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
+        :paramtype prompty_model_config: Dict[str, Any]
         :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",
         )
-        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_options=user_simulator_prompty_options,
-                        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"],
-                        "context": {},
-                        "$schema": "http://azureml/sdk-2-0/ChatConversation.json",
-                    }
+
+        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(current_simulation.to_list())
 
-        tasks = [asyncio.create_task(run_simulation(simulation)) for simulation in conversation_turns]
-        results = await asyncio.gather(*tasks)
         progress_bar.close()
-        return results
+        return simulated_conversations
 
     async def _extend_conversation_with_simulator(
         self,
@@ -291,12 +246,11 @@ class Simulator:
         current_simulation: ConversationHistory,
         max_conversation_turns: int,
         user_simulator_prompty: Optional[str],
-        user_simulator_prompty_options: Dict[str, Any],
+        user_simulator_prompty_kwargs: Dict[str, Any],
         api_call_delay_sec: float,
         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.
@@ -307,8 +261,8 @@ class Simulator:
         :paramtype max_conversation_turns: int,
         :keyword user_simulator_prompty: Path to the user simulator prompty file.
         :paramtype user_simulator_prompty: Optional[str],
-        :keyword user_simulator_prompty_options: Additional keyword arguments for the user simulator prompty.
-        :paramtype user_simulator_prompty_options: Dict[str, Any],
+        :keyword user_simulator_prompty_kwargs: Additional keyword arguments for the user simulator prompty.
+        :paramtype user_simulator_prompty_kwargs: Dict[str, Any],
         :keyword api_call_delay_sec: Delay in seconds between API calls.
         :paramtype api_call_delay_sec: float,
         :keyword prompty_model_config: The configuration for the prompty model.
@@ -317,92 +271,68 @@ 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
+            user_simulator_prompty=user_simulator_prompty,
             prompty_model_config=prompty_model_config,
-            user_simulator_prompty_options=user_simulator_prompty_options,
+            user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
         )
 
         while len(current_simulation) < max_conversation_turns:
-            user_response_content = await user_flow(
-                task="Continue the conversation",
-                conversation_history=current_simulation.to_context_free_list(),
-                **user_simulator_prompty_options,
+            user_response_content = user_flow(
+                task="Continue the conversation", conversation_history=current_simulation.to_list()
             )
             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, assistant_context = await self._get_target_response(
+            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, context=assistant_context
-            )
+            assistant_turn = Turn(role=ConversationRole.ASSISTANT, content=assistant_response)
             current_simulation.add_to_history(assistant_turn)
-            async with progress_bar_lock:
-                progress_bar.update(1)
+            progress_bar.update(1)
 
     def _load_user_simulation_flow(
         self,
         *,
-        user_simulator_prompty: Optional[Union[str, os.PathLike]],
+        user_simulator_prompty: Union[str, os.PathLike],
         prompty_model_config: Dict[str, Any],
-        user_simulator_prompty_options: Dict[str, Any],
-    ) -> "AsyncPrompty":  # type: ignore
+        user_simulator_prompty_kwargs: Dict[str, Any],
+    ) -> Flow:
         """
         Loads the flow for simulating user interactions.
 
         :keyword user_simulator_prompty: Path to the user simulator prompty file.
-        :paramtype user_simulator_prompty: Optional[Union[str, os.PathLike]]
+        :paramtype user_simulator_prompty: Union[str, os.PathLike]
         :keyword prompty_model_config: The configuration for the prompty model.
         :paramtype prompty_model_config: Dict[str, Any]
-        :keyword user_simulator_prompty_options: Additional keyword arguments for the user simulator prompty.
-        :paramtype user_simulator_prompty_options: Dict[str, Any]
+        :keyword user_simulator_prompty_kwargs: Additional keyword arguments for the user simulator prompty.
+        :paramtype user_simulator_prompty_kwargs: Dict[str, Any]
         :return: The loaded flow for simulating user interactions.
-        :rtype: AsyncPrompty
+        :rtype: Flow
         """
         if not user_simulator_prompty:
             package = "azure.ai.evaluation.simulator._prompty"
             resource_name = "task_simulate.prompty"
             try:
                 # Access the resource as a file path
-                # pylint: disable=deprecated-method
                 with pkg_resources.path(package, resource_name) as prompty_path:
-                    prompty_model_config = construct_prompty_model_config(
-                        model_config=prompty_model_config,  # type: ignore
-                        default_api_version="2024-06-01",
-                        user_agent=USER_AGENT,
-                    )
-                    return AsyncPrompty.load(source=prompty_path, model=prompty_model_config)  # type: ignore
+                    return load_flow(source=str(prompty_path), model=prompty_model_config)
             except FileNotFoundError as e:
-                msg = f"Flow path for {resource_name} does not exist in package {package}."
-                raise EvaluationException(
-                    message=msg,
-                    internal_message=msg,
-                    error_category=ErrorCategory.FILE_OR_FOLDER_NOT_FOUND,
-                    blame=ErrorBlame.USER_ERROR,
-                ) from e
-        prompty_model_config = construct_prompty_model_config(
-            model_config=prompty_model_config,  # type: ignore
-            default_api_version="2024-06-01",
-            user_agent=USER_AGENT,
-        )
-        return AsyncPrompty.load(
+                raise f"Flow path for {resource_name} does not exist in package {package}." from e
+        return load_flow(
             source=user_simulator_prompty,
             model=prompty_model_config,
-            **user_simulator_prompty_options,
-        )  # type: ignore
+            **user_simulator_prompty_kwargs,
+        )
 
     def _parse_prompty_response(self, *, response: str) -> Dict[str, Any]:
         """
         Parses the response from the prompty execution.
 
         :keyword response: The raw response from the prompty.
-        :paramtype response: str
+        :paramtype str: str
         :return: A dictionary representing the parsed response content.
         :rtype: Dict[str, Any]
         :raises ValueError: If the response cannot be parsed.
@@ -442,8 +372,8 @@ class Simulator:
         text: str,
         num_queries: int,
         query_response_generating_prompty: Optional[str],
-        query_response_generating_prompty_options: Dict[str, Any],
-        prompty_model_config: Any,
+        query_response_generating_prompty_kwargs: Dict[str, Any],
+        prompty_model_config: Dict[str, Any],
         **kwargs,
     ) -> List[Dict[str, str]]:
         """
@@ -455,32 +385,25 @@ class Simulator:
         :paramtype num_queries: int
         :keyword query_response_generating_prompty: Path to the query response generating prompty file.
         :paramtype query_response_generating_prompty: Optional[str]
-        :keyword query_response_generating_prompty_options: Additional keyword arguments for the query response generating prompty.
-        :paramtype query_response_generating_prompty_options: Dict[str, Any]
+        :keyword query_response_generating_prompty_kwargs: Additional keyword arguments for the query response generating prompty.
+        :paramtype query_response_generating_prompty_kwargs: Dict[str, Any]
         :keyword prompty_model_config: The configuration for the prompty model.
-        :paramtype prompty_model_config: Any
+        :paramtype prompty_model_config: Dict[str, Any]
         :return: A list of query-response dictionaries.
         :rtype: List[Dict[str, str]]
         :raises RuntimeError: If an error occurs during query generation.
         """
         query_flow = self._load_query_generation_flow(
-            query_response_generating_prompty=query_response_generating_prompty,  # type: ignore
+            query_response_generating_prompty=query_response_generating_prompty,
             prompty_model_config=prompty_model_config,
-            query_response_generating_prompty_options=query_response_generating_prompty_options,
+            query_response_generating_prompty_kwargs=query_response_generating_prompty_kwargs,
         )
+
         try:
-            query_responses = await query_flow(text=text, num_queries=num_queries)
+            query_responses = query_flow(text=text, num_queries=num_queries)
             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
@@ -488,65 +411,47 @@ class Simulator:
     def _load_query_generation_flow(
         self,
         *,
-        query_response_generating_prompty: Optional[Union[str, os.PathLike]],
+        query_response_generating_prompty: Union[str, os.PathLike],
         prompty_model_config: Dict[str, Any],
-        query_response_generating_prompty_options: Dict[str, Any],
-    ) -> "AsyncPrompty":
+        query_response_generating_prompty_kwargs: Dict[str, Any],
+    ) -> Flow:
         """
         Loads the flow for generating query responses.
 
         :keyword query_response_generating_prompty: Path to the query response generating prompty file.
-        :paramtype query_response_generating_prompty: Optional[Union[str, os.PathLike]]
+        :paramtype query_response_generating_prompty: Union[str, os.PathLike]
         :keyword prompty_model_config: The configuration for the prompty model.
         :paramtype prompty_model_config: Dict[str, Any]
-        :keyword query_response_generating_prompty_options: Additional keyword arguments for the flow.
-        :paramtype query_response_generating_prompty_options: Dict[str, Any]
+        :keyword query_response_generating_prompty_kwargs: Additional keyword arguments for the flow.
+        :paramtype query_response_generating_prompty_kwargs: Dict[str, Any]
         :return: The loaded flow for generating query responses.
-        :rtype: AsyncPrompty
+        :rtype: Flow
         """
         if not query_response_generating_prompty:
             package = "azure.ai.evaluation.simulator._prompty"
             resource_name = "task_query_response.prompty"
             try:
                 # Access the resource as a file path
-                # pylint: disable=deprecated-method
                 with pkg_resources.path(package, resource_name) as prompty_path:
-                    prompty_model_config = construct_prompty_model_config(
-                        model_config=prompty_model_config,  # type: ignore
-                        default_api_version="2024-06-01",
-                        user_agent=USER_AGENT,
-                    )
-                    return AsyncPrompty.load(source=prompty_path, model=prompty_model_config)  # type: ignore
+                    return load_flow(source=str(prompty_path), model=prompty_model_config)
             except FileNotFoundError as e:
-                msg = f"Flow path for {resource_name} does not exist in package {package}."
-                raise EvaluationException(
-                    message=msg,
-                    internal_message=msg,
-                    error_category=ErrorCategory.FILE_OR_FOLDER_NOT_FOUND,
-                    blame=ErrorBlame.USER_ERROR,
-                ) from e
-        prompty_model_config = construct_prompty_model_config(
-            model_config=prompty_model_config,  # type: ignore
-            default_api_version="2024-06-01",
-            user_agent=USER_AGENT,
-        )
-        return AsyncPrompty.load(
+                raise f"Flow path for {resource_name} does not exist in package {package}." from e
+        return load_flow(
             source=query_response_generating_prompty,
             model=prompty_model_config,
-            **query_response_generating_prompty_options,
-        )  # type: ignore
+            **query_response_generating_prompty_kwargs,
+        )
 
     async def _create_conversations_from_query_responses(
         self,
         *,
         query_responses: List[Dict[str, str]],
         max_conversation_turns: int,
-        tasks: List[str],
+        tasks: List[Dict],
         user_simulator_prompty: Optional[str],
-        user_simulator_prompty_options: Dict[str, Any],
+        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.
@@ -556,17 +461,15 @@ class Simulator:
         :keyword max_conversation_turns: The maximum number of conversation turns.
         :paramtype max_conversation_turns: int
         :keyword tasks: A list of tasks for the simulation.
-        :paramtype tasks: List[str]
+        :paramtype tasks: List[Dict]
         :keyword user_simulator_prompty: Path to the user simulator prompty file.
         :paramtype user_simulator_prompty: Optional[str]
-        :keyword user_simulator_prompty_options: Additional keyword arguments for the user simulator prompty.
-        :paramtype user_simulator_prompty_options: Dict[str, Any]
+        :keyword user_simulator_prompty_kwargs: Additional keyword arguments for the user simulator prompty.
+        :paramtype user_simulator_prompty_kwargs: Dict[str, Any]
         :keyword target: The target function to call for responses.
         :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]
         """
@@ -588,9 +491,9 @@ class Simulator:
             conversation = await self._complete_conversation(
                 conversation_starter=query,
                 max_conversation_turns=max_conversation_turns,
-                task=task,  # type: ignore
+                task=task,
                 user_simulator_prompty=user_simulator_prompty,
-                user_simulator_prompty_options=user_simulator_prompty_options,
+                user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
                 target=target,
                 api_call_delay_sec=api_call_delay_sec,
                 progress_bar=progress_bar,
@@ -604,7 +507,6 @@ class Simulator:
                             "task": task,
                             "expected_response": response,
                             "query": query,
-                            "original_text": text,
                         },
                         "$schema": "http://azureml/sdk-2-0/ChatConversation.json",
                     }
@@ -620,11 +522,11 @@ class Simulator:
         max_conversation_turns: int,
         task: str,
         user_simulator_prompty: Optional[str],
-        user_simulator_prompty_options: Dict[str, Any],
+        user_simulator_prompty_kwargs: Dict[str, Any],
         target: Callable,
         api_call_delay_sec: float,
         progress_bar: tqdm,
-    ) -> List[Dict[str, Optional[str]]]:
+    ) -> List[Dict[str, str]]:
         """
         Completes a conversation with the target model based on the conversation starter.
 
@@ -636,8 +538,8 @@ class Simulator:
         :paramtype task: str
         :keyword user_simulator_prompty: Path to the user simulator prompty file.
         :paramtype user_simulator_prompty: Optional[str]
-        :keyword user_simulator_prompty_options: Additional keyword arguments for the user simulator prompty.
-        :paramtype user_simulator_prompty_options: Dict[str, Any]
+        :keyword user_simulator_prompty_kwargs: Additional keyword arguments for the user simulator prompty.
+        :paramtype user_simulator_prompty_kwargs: Dict[str, Any]
         :keyword target: The target function to call for responses.
         :paramtype target: Callable
         :keyword api_call_delay_sec: Delay in seconds between API calls.
@@ -645,43 +547,36 @@ class Simulator:
         :keyword progress_bar: Progress bar for tracking simulation progress.
         :paramtype progress_bar: tqdm
         :return: A list representing the conversation history with each turn's content.
-        :rtype: List[Dict[str, Optional[str]]]
+        :rtype: List[Dict[str, 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(
-                user_simulator_prompty=user_simulator_prompty,  # type: ignore
-                prompty_model_config=self.model_config,  # type: ignore
-                user_simulator_prompty_options=user_simulator_prompty_options,
+                user_simulator_prompty=user_simulator_prompty,
+                prompty_model_config=self._build_prompty_model_config(),
+                user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
+            )
+            conversation_starter_from_simulated_user = 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, assistant_context = await self._get_target_response(
+            assistant_response = 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, context=assistant_context
-            )
+            assistant_turn = Turn(role=ConversationRole.ASSISTANT, content=assistant_response)
             conversation_history.add_to_history(assistant_turn)
             progress_bar.update(1)
 
@@ -690,9 +585,44 @@ class Simulator:
 
         return conversation_history.to_list()
 
+    async def _build_user_simulation_response(
+        self,
+        task: str,
+        conversation_history: List[Dict[str, Any]],
+        user_simulator_prompty: Optional[str],
+        user_simulator_prompty_kwargs: Dict[str, Any],
+    ) -> str:
+        """
+        Builds a response from the user simulator based on the current conversation history.
+
+        :param task: A string representing the task details.
+        :type task: str
+        :param conversation_history: The current conversation history as a list of dictionaries.
+        :type conversation_history: List[Dict[str, Any]]
+        :param user_simulator_prompty: Path to the user simulator prompty file.
+        :type user_simulator_prompty: Optional[str]
+        :param user_simulator_prompty_kwargs: Additional keyword arguments for the user simulator prompty.
+        :type user_simulator_prompty_kwargs: Dict[str, Any]
+        :return: The generated response content from the user simulator.
+        :rtype: str
+        :raises RuntimeError: If an error occurs during response generation.
+        """
+        user_flow = self._load_user_simulation_flow(
+            user_simulator_prompty=user_simulator_prompty,
+            prompty_model_config=self._build_prompty_model_config(),
+            user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
+        )
+
+        try:
+            response_content = user_flow(task=task, conversation_history=conversation_history)
+            user_response = self._parse_prompty_response(response=response_content)
+            return user_response["content"]
+        except Exception as e:
+            raise RuntimeError("Error building user simulation response") from e
+
     async def _get_target_response(
         self, *, target: Callable, api_call_delay_sec: float, conversation_history: ConversationHistory
-    ) -> Tuple[str, Optional[str]]:
+    ) -> str:
         """
         Retrieves the response from the target callback based on the current conversation history.
 
@@ -702,8 +632,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 and an optional context.
-        :rtype: str, Optional[str]
+        :return: The content of the response from the target.
+        :rtype: str
         """
         response = await target(
             messages={"messages": conversation_history.to_list()},
@@ -713,4 +643,4 @@ class Simulator:
         )
         await asyncio.sleep(api_call_delay_sec)
         latest_message = response["messages"][-1]
-        return latest_message["content"], latest_message.get("context", "")  # type: ignore
+        return latest_message["content"]