diff --git a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_conversation/conversation.py b/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_conversation/conversation.py index 63f2eefb0508..814da640c76b 100644 --- a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_conversation/conversation.py +++ b/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)) diff --git a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_conversation/conversation_bot.py b/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_conversation/conversation_bot.py index 2dcb48c5b98f..705c653e75b6 100644 --- a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_conversation/conversation_bot.py +++ b/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 diff --git a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_model_tools/models.py b/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_model_tools/models.py index 246eccb89f35..3923ec070935 100644 --- a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/_model_tools/models.py +++ b/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) diff --git a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/_callback_conversation_bot.py b/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/_callback_conversation_bot.py index 8cdf2f45ec8c..03460ef8800f 100644 --- a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/_callback_conversation_bot.py +++ b/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/_callback_conversation_bot.py @@ -46,6 +46,17 @@ async def generate_response( "id": None, "template_parameters": {} } + if not result: + result = { + "messages": [{ + "content": "Callback did not return a response.", + "role": "assistant" + }], + "finish_reason": ["stop"], + "id": None, + "template_parameters": {} + } + self.logger.info("Using user provided callback returning response.") time_taken = 0 diff --git a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/_utils.py b/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/_utils.py index 3f4db78249d5..a50e3562efbd 100644 --- a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/_utils.py +++ b/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/_utils.py @@ -1,7 +1,7 @@ # --------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -# pylint: disable=C0303 +# pylint: skip-file """ This module contains a utility class for managing a list of JSON lines. """ @@ -52,14 +52,17 @@ def to_eval_qa_json_lines(self): assistant_message = message['content'] if 'context' in message: context = message.get("context", None) - if user_message and assistant_message: - if context: - json_lines += json.dumps({ - 'question': user_message, - 'answer': assistant_message, - 'context': context}) + "\n" - else: - json_lines += json.dumps({ - 'question': user_message, - 'answer': assistant_message}) + "\n" + if user_message and assistant_message: + if context: + json_lines += json.dumps({ + 'question': user_message, + 'answer': assistant_message, + 'context': context}) + "\n" + user_message = assistant_message = context = None + else: + json_lines += json.dumps({ + 'question': user_message, + 'answer': assistant_message}) + "\n" + user_message = assistant_message = None + return json_lines diff --git a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/simulator.py b/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/simulator.py index 2554ad75e489..2f653bfac2a3 100644 --- a/sdk/ai/azure-ai-generative/azure/ai/generative/synthetic/simulator/simulator/simulator.py +++ b/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,17 +257,23 @@ 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] @@ -254,12 +281,32 @@ async def simulate_async( 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): diff --git a/sdk/ai/azure-ai-generative/cspell.json b/sdk/ai/azure-ai-generative/cspell.json index 0f49bbd407e5..687054a7bdd2 100644 --- a/sdk/ai/azure-ai-generative/cspell.json +++ b/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/**/*"] }