Simulator quiet run and progress bar (#34809)

* Reverting models to make sure calls to the simulator work

* quotes

* Spellcheck fixes

* ignore the models for doc generation

* Fixed the quotes on f strings

* pylint skip file

* Support for summarization

* Adding a limit of 2 conversation turns for all but conversation simulators

* exclude synthetic from mypy

* Another lint fix

* Skip the file causing linting issues

* Bugfix on output to json_qa_lines and empty response from callbacks

* Skip pylint

* Add if/else on message to eval json util

* Remove the verbose logs, add a progress bar and display a message when requested responses > rai service responses

* Added back missing param

* Spellcheck fix for pbar

* Adding parameter to the spellcheck

* Simulator logging initialized on the top of the page and reused

* Clean up the warning log message

* Removing logs from AsyncHttpClientWithRetry

* Using the right param

Author: nagkumar91

sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_conversation/conversation.py

@@ -90,8 +90,6 @@ async def simulate_conversation(
     else:
         conversation_id = None
     first_prompt = first_response["samples"][0]
-    logger.info(f"First turn: {first_prompt}")
-
     # Add all generated turns into array to pass for each bot while generating
     # new responses. We add generated response and the person generating it.
     # in the case of the first turn, it is supposed to be the user search query
@@ -115,7 +113,6 @@ async def simulate_conversation(
             current_character_idx = current_turn % len(bots)
             current_bot = bots[current_character_idx]
             # invoke Bot to generate response given the input request
-            logger.info("-- Sending to %s", current_bot.role.value)
             # pass only the last generated turn without passing the bot name.
             response, request, time_taken, full_response = await current_bot.generate_response(
                 session=session,
@@ -137,7 +134,6 @@ async def simulate_conversation(
                     request=request,
                 )
             )
-            logger.info("Last turn: %s", conversation_history[-1])
             if mlflow_logger is not None:
                 logger_tasks.append(  # schedule logging but don't get blocked by it
                     asyncio.create_task(mlflow_logger.log_successful_response(time_taken))

sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_conversation/conversation_bot.py

@@ -57,21 +57,13 @@ def __init__(
         self.model = model
 
         self.logger = logging.getLogger(repr(self))
-
         self.conversation_starter = None  # can either be a dictionary or jinja template
         if role == ConversationRole.USER:
             if "conversation_starter" in self.persona_template_args:
                 conversation_starter_content = self.persona_template_args["conversation_starter"]
                 if isinstance(conversation_starter_content, dict):
-                    msg = f"{conversation_starter_content} instead of generating a turn using a LLM"
-                    self.logger.info(
-                        "This simulated bot will use the provided conversation starter (passed in as dictionary): %s",
-                        msg,
-                    )
                     self.conversation_starter = conversation_starter_content
                 else:
-                    msg = f"{repr(conversation_starter_content)[:400]} instead of generating a turn using a LLM"
-                    self.logger.info("This simulated bot will use the provided conversation starter %s", msg)
                     self.conversation_starter = jinja2.Template(
                         conversation_starter_content, undefined=jinja2.StrictUndefined
                     )
@@ -107,12 +99,8 @@ async def generate_response(
         if turn_number == 0 and self.conversation_starter is not None:
             # if conversation_starter is a dictionary, pass it into samples as is
             if isinstance(self.conversation_starter, dict):
-                self.logger.info("Returning conversation starter: %s", self.conversation_starter)
                 samples = [self.conversation_starter]
             else:
-                self.logger.info(
-                    "Returning conversation starter: %s", repr(self.persona_template_args["conversation_starter"])[:400]
-                )
                 samples = [self.conversation_starter.render(**self.persona_template_args)]  # type: ignore[attr-defined]
             time_taken = 0
 

sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_model_tools/models.py

@@ -47,8 +47,9 @@ def __init__(self, n_retry, retry_timeout, logger, retry_options=None):
         # Set up async HTTP client with retry
 
         trace_config = TraceConfig()  # set up request logging
-        trace_config.on_request_start.append(self.on_request_start)
-        trace_config.on_request_end.append(self.on_request_end)
+        trace_config.on_request_end.append(self.delete_auth_header)
+        # trace_config.on_request_start.append(self.on_request_start)
+        # trace_config.on_request_end.append(self.on_request_end) 
         if retry_options is None:
             retry_options = RandomRetry(  # set up retry configuration
                 statuses=[104, 408, 409, 424, 429, 500, 502,
@@ -67,6 +68,13 @@ async def on_request_start(self, session, trace_config_ctx, params):
             current_attempt, params.method, params.url
         ))
 
+    async def delete_auth_header(self, session, trace_config_ctx, params):
+        request_headers = dict(params.response.request_info.headers)
+        if "Authorization" in request_headers:
+            del request_headers["Authorization"]
+        if "api-key" in request_headers:
+            del request_headers["api-key"]
+
     async def on_request_end(self, session, trace_config_ctx, params):
         current_attempt = trace_config_ctx.trace_request_ctx["current_attempt"]
         request_headers = dict(params.response.request_info.headers)

sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/simulator.py

@@ -1,7 +1,7 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
-# pylint: disable=E0401
+# pylint: skip-file
 # needed for 'list' type annotations on 3.8
 from __future__ import annotations
 
@@ -12,6 +12,9 @@
 import threading
 import json
 import random
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
 
 from azure.ai.generative.synthetic.simulator._conversation import (
     ConversationBot,
@@ -78,7 +81,9 @@ def __init__(
         if (ai_client is None and simulator_connection is None) or (
             ai_client is not None and simulator_connection is not None
         ):
-            raise ValueError("One and only one of the parameters [ai_client, simulator_connection] has to be set.")
+            raise ValueError(
+                "One and only one of the parameters [ai_client, simulator_connection] has to be set."
+            )
 
         if simulate_callback is None:
             raise ValueError("Callback cannot be None.")
@@ -87,7 +92,9 @@ def __init__(
             raise ValueError("Callback has to be an async function.")
 
         self.ai_client = ai_client
-        self.simulator_connection = self._to_openai_chat_completion_model(simulator_connection)
+        self.simulator_connection = self._to_openai_chat_completion_model(
+            simulator_connection
+        )
         self.adversarial = False
         self.rai_client = None
         if ai_client:
@@ -168,19 +175,33 @@ def _create_bot(
             instantiation_parameters=instantiation_parameters,
         )
 
-    def _setup_bot(self, role: Union[str, ConversationRole], template: "Template", parameters: dict):
+    def _setup_bot(
+        self,
+        role: Union[str, ConversationRole],
+        template: "Template",
+        parameters: dict,
+    ):
         if role == ConversationRole.ASSISTANT:
             return self._create_bot(role, str(template), parameters)
         if role == ConversationRole.USER:
             if template.content_harm:
-                return self._create_bot(role, str(template), parameters, template.template_name)
+                return self._create_bot(
+                    role, str(template), parameters, template.template_name
+                )
 
-            return self._create_bot(role, str(template), parameters, model=self.simulator_connection)
+            return self._create_bot(
+                role,
+                str(template),
+                parameters,
+                model=self.simulator_connection,
+            )
         return None
 
     def _ensure_service_dependencies(self):
         if self.rai_client is None:
-            raise ValueError("Simulation options require rai services but ai client is not provided.")
+            raise ValueError(
+                "Simulation options require rai services but ai client is not provided."
+            )
 
     def _join_conversation_starter(self, parameters, to_join):
         key = "conversation_starter"
@@ -236,30 +257,56 @@ async def simulate_async(
         if parameters is None:
             parameters = []
         if not isinstance(template, Template):
-            raise ValueError(f"Please use simulator to construct template. Found {type(template)}")
+            raise ValueError(
+                f"Please use simulator to construct template. Found {type(template)}"
+            )
 
         if not isinstance(parameters, list):
-            raise ValueError(f"Expect parameters to be a list of dictionary, but found {type(parameters)}")
+            raise ValueError(
+                f"Expect parameters to be a list of dictionary, but found {type(parameters)}"
+            )
         if "conversation" not in template.template_name:
             max_conversation_turns = 2
         if template.content_harm:
             self._ensure_service_dependencies()
             self.adversarial = True
             # pylint: disable=protected-access
-            templates = await self.template_handler._get_ch_template_collections(template.template_name)
+            templates = await self.template_handler._get_ch_template_collections(
+                template.template_name
+            )
         else:
             template.template_parameters = parameters
             templates = [template]
 
         semaphore = asyncio.Semaphore(concurrent_async_task)
         sim_results = []
         tasks = []
+        total_tasks = sum(len(t.template_parameters) for t in templates)
+
+        if simulation_result_limit > total_tasks and self.adversarial:
+            logger.warning(
+                "Cannot provide %s results due to maximum number of adversarial simulations that can be generated: %s."
+                "\n %s simulations will be generated.",
+                simulation_result_limit,
+                total_tasks,
+                total_tasks,
+            )
+        total_tasks = min(total_tasks, simulation_result_limit)
+        progress_bar = tqdm(
+            total=total_tasks,
+            desc="generating simulations",
+            ncols=100,
+            unit="simulations",
+        )
+
         for t in templates:
             for p in t.template_parameters:
                 if jailbreak:
                     self._ensure_service_dependencies()
                     jailbreak_dataset = await self.rai_client.get_jailbreaks_dataset()  # type: ignore[union-attr]
-                    p = self._join_conversation_starter(p, random.choice(jailbreak_dataset))
+                    p = self._join_conversation_starter(
+                        p, random.choice(jailbreak_dataset)
+                    )
 
                 tasks.append(
                     asyncio.create_task(
@@ -280,7 +327,15 @@ async def simulate_async(
             if len(tasks) >= simulation_result_limit:
                 break
 
-        sim_results = await asyncio.gather(*tasks)
+        sim_results = []
+
+        # Use asyncio.as_completed to update the progress bar when a task is complete
+        for task in asyncio.as_completed(tasks):
+            result = await task
+            sim_results.append(result)  # Store the result
+            progress_bar.update(1)
+
+        progress_bar.close()
 
         return JsonLineList(sim_results)
 
@@ -319,7 +374,9 @@ async def _simulate_async(
             parameters = {}
         # create user bot
         user_bot = self._setup_bot(ConversationRole.USER, template, parameters)
-        system_bot = self._setup_bot(ConversationRole.ASSISTANT, template, parameters)
+        system_bot = self._setup_bot(
+            ConversationRole.ASSISTANT, template, parameters
+        )
 
         bots = [user_bot, system_bot]
 
@@ -328,7 +385,7 @@ async def _simulate_async(
         asyncHttpClient = AsyncHTTPClientWithRetry(
             n_retry=api_call_retry_limit,
             retry_timeout=api_call_retry_sleep_sec,
-            logger=logging.getLogger(),
+            logger=logger,
         )
         async with sem:
             async with asyncHttpClient.client as session:
@@ -357,7 +414,9 @@ def _get_citations(self, parameters, context_keys, turn_num=None):
                 else:
                     for k, v in parameters[c_key].items():
                         if k not in ["callback_citations", "callback_citation_key"]:
-                            citations.append({"id": k, "content": self._to_citation_content(v)})
+                            citations.append(
+                                {"id": k, "content": self._to_citation_content(v)}
+                            )
             else:
                 citations.append(
                     {
@@ -373,7 +432,9 @@ def _to_citation_content(self, obj):
             return obj
         return json.dumps(obj)
 
-    def _get_callback_citations(self, callback_citations: dict, turn_num: Optional[int] = None):
+    def _get_callback_citations(
+        self, callback_citations: dict, turn_num: Optional[int] = None
+    ):
         if turn_num is None:
             return []
         current_turn_citations = []
@@ -382,7 +443,9 @@ def _get_callback_citations(self, callback_citations: dict, turn_num: Optional[i
             citations = callback_citations[current_turn_str]
             if isinstance(citations, dict):
                 for k, v in citations.items():
-                    current_turn_citations.append({"id": k, "content": self._to_citation_content(v)})
+                    current_turn_citations.append(
+                        {"id": k, "content": self._to_citation_content(v)}
+                    )
             else:
                 current_turn_citations.append(
                     {
@@ -397,13 +460,15 @@ def _to_chat_protocol(self, template, conversation_history, template_parameters)
         for i, m in enumerate(conversation_history):
             message = {"content": m.message, "role": m.role.value}
             if len(template.context_key) > 0:
-                citations = self._get_citations(template_parameters, template.context_key, i)
+                citations = self._get_citations(
+                    template_parameters, template.context_key, i
+                )
                 message["context"] = citations
             elif "context" in m.full_response:
                 # adding context for adv_qa
                 message["context"] = m.full_response["context"]
             messages.append(message)
-        template_parameters['metadata'] = {}
+        template_parameters["metadata"] = {}
         if "ch_template_placeholder" in template_parameters:
             del template_parameters["ch_template_placeholder"]
 
@@ -524,8 +589,13 @@ def from_fn(
         if hasattr(fn, "__wrapped__"):
             func_module = fn.__wrapped__.__module__
             func_name = fn.__wrapped__.__name__
-            if func_module == "openai.resources.chat.completions" and func_name == "create":
-                return Simulator._from_openai_chat_completions(fn, simulator_connection, ai_client, **kwargs)
+            if (
+                func_module == "openai.resources.chat.completions"
+                and func_name == "create"
+            ):
+                return Simulator._from_openai_chat_completions(
+                    fn, simulator_connection, ai_client, **kwargs
+                )
 
         return Simulator(
             simulator_connection=simulator_connection,
@@ -534,7 +604,9 @@ def from_fn(
         )
 
     @staticmethod
-    def _from_openai_chat_completions(fn: Callable[[Any], dict], simulator_connection=None, ai_client=None, **kwargs):
+    def _from_openai_chat_completions(
+        fn: Callable[[Any], dict], simulator_connection=None, ai_client=None, **kwargs
+    ):
         return Simulator(
             simulator_connection=simulator_connection,
             ai_client=ai_client,
@@ -625,7 +697,9 @@ async def callback(chat_protocol_message):
                 input_data[chat_history_key] = all_messages
 
             response = flow.invoke(input_data).output
-            chat_protocol_message["messages"].append({"role": "assistant", "content": response[chat_output_key]})
+            chat_protocol_message["messages"].append(
+                {"role": "assistant", "content": response[chat_output_key]}
+            )
 
             return chat_protocol_message
 
@@ -657,8 +731,12 @@ def create_template(
         One of 'template' or 'template_path' must be provided to create a template. If 'template' is provided,
         it is used directly; if 'template_path' is provided, the content is read from the file at that path.
         """
-        if (template is None and template_path is None) or (template is not None and template_path is not None):
-            raise ValueError("One and only one of the parameters [template, template_path] has to be set.")
+        if (template is None and template_path is None) or (
+            template is not None and template_path is not None
+        ):
+            raise ValueError(
+                "One and only one of the parameters [template, template_path] has to be set."
+            )
 
         if template is not None:
             return Template(template_name=name, text=template, context_key=context_key)
@@ -669,7 +747,9 @@ def create_template(
 
             return Template(template_name=name, text=tc, context_key=context_key)
 
-        raise ValueError("Condition not met for creating template, please check examples and parameter list.")
+        raise ValueError(
+            "Condition not met for creating template, please check examples and parameter list."
+        )
 
     @staticmethod
     def get_template(template_name: str):

sdk/ai/azure-ai-generative/cspell.json

@@ -1,4 +1,4 @@
 {
-    "ignoreWords": ["cmpl", "uqkvl", "redef", "datas", "unbatched", "endofprompt", "unlabel", "pydash", "raisvc", "tkey", "tparam", "punc"],
+    "ignoreWords": ["cmpl", "uqkvl", "redef", "datas", "unbatched", "endofprompt", "unlabel", "pydash", "raisvc", "tkey", "tparam", "punc", "ncols"],
     "ignorePaths": ["sdk/ai/azure-ai-generative/azure/ai/generative/evaluate/pf_templates/**/*"]
 }