azure-ai-evaluation 1.0.0b4__py3-none-any.whl → 1.0.1__py3-none-any.whl

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
 
 
@@ -26,16 +29,22 @@ 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]):
-        """
-        Initializes the task simulator with the model configuration.
-
-        :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.
-        """
         self._validate_model_config(model_config)
         self.model_config = model_config
         if "api_version" not in self.model_config:
@@ -87,9 +96,10 @@ class Simulator:
         query_response_generating_prompty: Optional[str] = None,
         user_simulator_prompty: Optional[str] = None,
         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]] = [],
+        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,
         **kwargs,
     ) -> List[JsonLineChatProtocol]:
         """
@@ -111,12 +121,15 @@ 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_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 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 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 +144,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."
@@ -151,16 +164,17 @@ class Simulator:
                 max_conversation_turns=max_conversation_turns,
                 conversation_turns=conversation_turns,
                 user_simulator_prompty=user_simulator_prompty,
-                user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
+                user_simulator_prompty_options=user_simulator_prompty_options,
                 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_kwargs=query_response_generating_prompty_kwargs,
+            query_response_generating_prompty_options=query_response_generating_prompty_options,
             prompty_model_config=prompty_model_config,
             **kwargs,
         )
@@ -169,9 +183,10 @@ class Simulator:
             max_conversation_turns=max_conversation_turns,
             tasks=tasks,
             user_simulator_prompty=user_simulator_prompty,
-            user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
+            user_simulator_prompty_options=user_simulator_prompty_options,
             target=target,
             api_call_delay_sec=api_call_delay_sec,
+            text=text,
         )
 
     async def _simulate_with_predefined_turns(
@@ -179,11 +194,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],
+        user_simulator_prompty_options: 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,51 +209,69 @@ 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.
-        :paramtype user_simulator_prompty_kwargs: 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 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
         :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_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"],
@@ -245,10 +279,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,
@@ -256,11 +291,12 @@ class Simulator:
         current_simulation: ConversationHistory,
         max_conversation_turns: int,
         user_simulator_prompty: Optional[str],
-        user_simulator_prompty_kwargs: Dict[str, Any],
+        user_simulator_prompty_options: 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.
@@ -271,8 +307,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_kwargs: Additional keyword arguments for the user simulator prompty.
-        :paramtype user_simulator_prompty_kwargs: 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 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.
@@ -281,36 +317,41 @@ 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
             prompty_model_config=prompty_model_config,
-            user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
+            user_simulator_prompty_options=user_simulator_prompty_options,
         )
 
         while len(current_simulation) < max_conversation_turns:
             user_response_content = await user_flow(
                 task="Continue the conversation",
-                conversation_history=current_simulation.to_list(),
-                **user_simulator_prompty_kwargs,
+                conversation_history=current_simulation.to_context_free_list(),
+                **user_simulator_prompty_options,
             )
             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,
         *,
         user_simulator_prompty: Optional[Union[str, os.PathLike]],
         prompty_model_config: Dict[str, Any],
-        user_simulator_prompty_kwargs: Dict[str, Any],
+        user_simulator_prompty_options: Dict[str, Any],
     ) -> "AsyncPrompty":  # type: ignore
         """
         Loads the flow for simulating user interactions.
@@ -319,8 +360,8 @@ class Simulator:
         :paramtype user_simulator_prompty: Optional[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_kwargs: Additional keyword arguments for the user simulator prompty.
-        :paramtype user_simulator_prompty_kwargs: Dict[str, Any]
+        :keyword user_simulator_prompty_options: Additional keyword arguments for the user simulator prompty.
+        :paramtype user_simulator_prompty_options: Dict[str, Any]
         :return: The loaded flow for simulating user interactions.
         :rtype: AsyncPrompty
         """
@@ -353,7 +394,7 @@ class Simulator:
         return AsyncPrompty.load(
             source=user_simulator_prompty,
             model=prompty_model_config,
-            **user_simulator_prompty_kwargs,
+            **user_simulator_prompty_options,
         )  # type: ignore
 
     def _parse_prompty_response(self, *, response: str) -> Dict[str, Any]:
@@ -401,7 +442,7 @@ class Simulator:
         text: str,
         num_queries: int,
         query_response_generating_prompty: Optional[str],
-        query_response_generating_prompty_kwargs: Dict[str, Any],
+        query_response_generating_prompty_options: Dict[str, Any],
         prompty_model_config: Any,
         **kwargs,
     ) -> List[Dict[str, str]]:
@@ -414,8 +455,8 @@ 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_kwargs: Additional keyword arguments for the query response generating prompty.
-        :paramtype query_response_generating_prompty_kwargs: Dict[str, Any]
+        :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 prompty_model_config: The configuration for the prompty model.
         :paramtype prompty_model_config: Any
         :return: A list of query-response dictionaries.
@@ -425,13 +466,21 @@ class Simulator:
         query_flow = self._load_query_generation_flow(
             query_response_generating_prompty=query_response_generating_prompty,  # type: ignore
             prompty_model_config=prompty_model_config,
-            query_response_generating_prompty_kwargs=query_response_generating_prompty_kwargs,
+            query_response_generating_prompty_options=query_response_generating_prompty_options,
         )
         try:
             query_responses = await 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
@@ -441,7 +490,7 @@ class Simulator:
         *,
         query_response_generating_prompty: Optional[Union[str, os.PathLike]],
         prompty_model_config: Dict[str, Any],
-        query_response_generating_prompty_kwargs: Dict[str, Any],
+        query_response_generating_prompty_options: Dict[str, Any],
     ) -> "AsyncPrompty":
         """
         Loads the flow for generating query responses.
@@ -450,8 +499,8 @@ class Simulator:
         :paramtype query_response_generating_prompty: Optional[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_kwargs: Additional keyword arguments for the flow.
-        :paramtype query_response_generating_prompty_kwargs: Dict[str, Any]
+        :keyword query_response_generating_prompty_options: Additional keyword arguments for the flow.
+        :paramtype query_response_generating_prompty_options: Dict[str, Any]
         :return: The loaded flow for generating query responses.
         :rtype: AsyncPrompty
         """
@@ -484,7 +533,7 @@ class Simulator:
         return AsyncPrompty.load(
             source=query_response_generating_prompty,
             model=prompty_model_config,
-            **query_response_generating_prompty_kwargs,
+            **query_response_generating_prompty_options,
         )  # type: ignore
 
     async def _create_conversations_from_query_responses(
@@ -494,9 +543,10 @@ class Simulator:
         max_conversation_turns: int,
         tasks: List[str],
         user_simulator_prompty: Optional[str],
-        user_simulator_prompty_kwargs: Dict[str, Any],
+        user_simulator_prompty_options: Dict[str, Any],
         target: Callable,
         api_call_delay_sec: float,
+        text: str,
     ) -> List[JsonLineChatProtocol]:
         """
         Creates full conversations from query-response pairs.
@@ -509,12 +559,14 @@ class Simulator:
         :paramtype tasks: List[str]
         :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.
-        :paramtype user_simulator_prompty_kwargs: 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 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]
         """
@@ -538,7 +590,7 @@ class Simulator:
                 max_conversation_turns=max_conversation_turns,
                 task=task,  # type: ignore
                 user_simulator_prompty=user_simulator_prompty,
-                user_simulator_prompty_kwargs=user_simulator_prompty_kwargs,
+                user_simulator_prompty_options=user_simulator_prompty_options,
                 target=target,
                 api_call_delay_sec=api_call_delay_sec,
                 progress_bar=progress_bar,
@@ -552,6 +604,7 @@ class Simulator:
                             "task": task,
                             "expected_response": response,
                             "query": query,
+                            "original_text": text,
                         },
                         "$schema": "http://azureml/sdk-2-0/ChatConversation.json",
                     }
@@ -567,7 +620,7 @@ class Simulator:
         max_conversation_turns: int,
         task: str,
         user_simulator_prompty: Optional[str],
-        user_simulator_prompty_kwargs: Dict[str, Any],
+        user_simulator_prompty_options: Dict[str, Any],
         target: Callable,
         api_call_delay_sec: float,
         progress_bar: tqdm,
@@ -583,8 +636,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_kwargs: Additional keyword arguments for the user simulator prompty.
-        :paramtype user_simulator_prompty_kwargs: 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 target: The target function to call for responses.
         :paramtype target: Callable
         :keyword api_call_delay_sec: Delay in seconds between API calls.
@@ -595,33 +648,40 @@ 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(
                 user_simulator_prompty=user_simulator_prompty,  # type: ignore
                 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.",
-                    }
-                ],
+                user_simulator_prompty_options=user_simulator_prompty_options,
             )
+            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 +692,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 +702,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 +713,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

azure/ai/evaluation/simulator/_utils.py

@@ -26,9 +26,9 @@ class JsonLineList(list):
             json_lines += json.dumps(item) + "\n"
         return json_lines
 
-    def to_eval_qa_json_lines(self):
+    def to_eval_qr_json_lines(self):
         """
-        Converts the list to a string of JSON lines suitable for evaluation in a Q&A format.
+        Converts the list to a string of JSON lines suitable for evaluation in a query & response format.
         Each item in the list is expected to be a dictionary with
         'messages' key. The 'messages' value is a list of
         dictionaries, each with a 'role' key and a 'content' key.
@@ -44,23 +44,41 @@ class JsonLineList(list):
         for item in self:
             user_message = None
             assistant_message = None
-            context = None
+            user_context = None
+            assistant_context = None
+            template_parameters = item.get("template_parameters", {})
+            category = template_parameters.get("category", None)
             for message in item["messages"]:
                 if message["role"] == "user":
                     user_message = message["content"]
+                    user_context = message.get("context", "")
                 elif message["role"] == "assistant":
                     assistant_message = message["content"]
-                if "context" in message:
-                    context = message.get("context", None)
+                    assistant_context = message.get("context", "")
                 if user_message and assistant_message:
-                    if context:
+                    if user_context or assistant_context:
                         json_lines += (
-                            json.dumps({"query": user_message, "response": assistant_message, "context": context})
+                            json.dumps(
+                                {
+                                    "query": user_message,
+                                    "response": assistant_message,
+                                    "context": str(
+                                        {
+                                            "user_context": user_context,
+                                            "assistant_context": assistant_context,
+                                        }
+                                    ),
+                                    "category": category,
+                                }
+                            )
                             + "\n"
                         )
-                        user_message = assistant_message = context = None
+                        user_message = assistant_message = None
                     else:
-                        json_lines += json.dumps({"query": user_message, "response": assistant_message}) + "\n"
+                        json_lines += (
+                            json.dumps({"query": user_message, "response": assistant_message, "category": category})
+                            + "\n"
+                        )
                         user_message = assistant_message = None
 
         return json_lines
@@ -80,9 +98,9 @@ class JsonLineChatProtocol(dict):
         """
         return json.dumps(self)
 
-    def to_eval_qa_json_lines(self) -> str:
+    def to_eval_qr_json_lines(self) -> str:
         """
-        Converts the object to a string of JSON lines suitable for evaluation in a Q&A format.
+        Converts the object to a string of JSON lines suitable for evaluation in a query and response format.
         The object is expected to be a dictionary with 'messages' key.
 
         :returns: A json lines document
@@ -105,10 +123,10 @@ class JsonLineChatProtocol(dict):
             if user_message and assistant_message:
                 if context:
                     json_lines += (
-                        json.dumps({"question": user_message, "answer": assistant_message, "context": context}) + "\n"
+                        json.dumps({"query": user_message, "response": assistant_message, "context": context}) + "\n"
                     )
                     user_message = assistant_message = None
                 else:
-                    json_lines += json.dumps({"question": user_message, "answer": assistant_message}) + "\n"
+                    json_lines += json.dumps({"query": user_message, "response": assistant_message}) + "\n"
                     user_message = assistant_message = None
         return json_lines