azure-ai-evaluation 1.10.0__py3-none-any.whl → 1.11.0__py3-none-any.whl

azure/ai/evaluation/red_team/_mlflow_integration.py

@@ -0,0 +1,322 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""
+MLflow integration module for Red Team Agent.
+
+This module handles MLflow run creation, logging, and tracking for red team evaluations.
+"""
+
+import json
+import os
+import tempfile
+import uuid
+from datetime import datetime
+from typing import Dict, Optional, cast
+from pathlib import Path
+
+# Azure AI Evaluation imports
+from azure.ai.evaluation._evaluate._eval_run import EvalRun
+from azure.ai.evaluation._evaluate._utils import _trace_destination_from_project_scope, _get_ai_studio_url
+from azure.ai.evaluation._evaluate._utils import extract_workspace_triad_from_trace_provider
+from azure.ai.evaluation._version import VERSION
+from azure.ai.evaluation._azure._clients import LiteMLClient
+from azure.ai.evaluation._constants import EvaluationRunProperties, DefaultOpenEncoding
+from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException
+from azure.ai.evaluation._common import RedTeamUpload, ResultType
+from azure.ai.evaluation._model_configurations import AzureAIProject
+
+# Local imports
+from ._red_team_result import RedTeamResult
+from ._utils.logging_utils import log_error
+
+
+class MLflowIntegration:
+    """Handles MLflow integration for red team evaluations."""
+
+    def __init__(self, logger, azure_ai_project, generated_rai_client, one_dp_project, scan_output_dir=None):
+        """Initialize the MLflow integration.
+
+        :param logger: Logger instance for logging
+        :param azure_ai_project: Azure AI project configuration
+        :param generated_rai_client: RAI client for service interactions
+        :param one_dp_project: Whether this is a OneDP project
+        :param scan_output_dir: Directory for scan outputs
+        """
+        self.logger = logger
+        self.azure_ai_project = azure_ai_project
+        self.generated_rai_client = generated_rai_client
+        self._one_dp_project = one_dp_project
+        self.scan_output_dir = scan_output_dir
+        self.ai_studio_url = None
+        self.trace_destination = None
+
+    def start_redteam_mlflow_run(
+        self,
+        azure_ai_project: Optional[AzureAIProject] = None,
+        run_name: Optional[str] = None,
+    ) -> EvalRun:
+        """Start an MLFlow run for the Red Team Agent evaluation.
+
+        :param azure_ai_project: Azure AI project details for logging
+        :type azure_ai_project: Optional[AzureAIProject]
+        :param run_name: Optional name for the MLFlow run
+        :type run_name: Optional[str]
+        :return: The MLFlow run object
+        :rtype: EvalRun
+        :raises EvaluationException: If no azure_ai_project is provided or trace destination cannot be determined
+        """
+        if not azure_ai_project:
+            log_error(self.logger, "No azure_ai_project provided, cannot upload run")
+            raise EvaluationException(
+                message="No azure_ai_project provided",
+                blame=ErrorBlame.USER_ERROR,
+                category=ErrorCategory.MISSING_FIELD,
+                target=ErrorTarget.RED_TEAM,
+            )
+
+        if self._one_dp_project:
+            response = self.generated_rai_client._evaluation_onedp_client.start_red_team_run(
+                red_team=RedTeamUpload(
+                    display_name=run_name or f"redteam-agent-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
+                )
+            )
+
+            self.ai_studio_url = response.properties.get("AiStudioEvaluationUri")
+            return response
+
+        else:
+            trace_destination = _trace_destination_from_project_scope(azure_ai_project)
+            if not trace_destination:
+                self.logger.warning("Could not determine trace destination from project scope")
+                raise EvaluationException(
+                    message="Could not determine trace destination",
+                    blame=ErrorBlame.SYSTEM_ERROR,
+                    category=ErrorCategory.UNKNOWN,
+                    target=ErrorTarget.RED_TEAM,
+                )
+
+            ws_triad = extract_workspace_triad_from_trace_provider(trace_destination)
+
+            management_client = LiteMLClient(
+                subscription_id=ws_triad.subscription_id,
+                resource_group=ws_triad.resource_group_name,
+                logger=self.logger,
+                credential=azure_ai_project.get("credential"),
+            )
+
+            tracking_uri = management_client.workspace_get_info(ws_triad.workspace_name).ml_flow_tracking_uri
+
+            run_display_name = run_name or f"redteam-agent-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+            self.logger.debug(f"Starting MLFlow run with name: {run_display_name}")
+            eval_run = EvalRun(
+                run_name=run_display_name,
+                tracking_uri=cast(str, tracking_uri),
+                subscription_id=ws_triad.subscription_id,
+                group_name=ws_triad.resource_group_name,
+                workspace_name=ws_triad.workspace_name,
+                management_client=management_client,
+            )
+            eval_run._start_run()
+            self.logger.debug(f"MLFlow run started successfully with ID: {eval_run.info.run_id}")
+
+            self.trace_destination = trace_destination
+            self.logger.debug(f"MLFlow run created successfully with ID: {eval_run}")
+
+            self.ai_studio_url = _get_ai_studio_url(
+                trace_destination=self.trace_destination,
+                evaluation_id=eval_run.info.run_id,
+            )
+
+            return eval_run
+
+    async def log_redteam_results_to_mlflow(
+        self,
+        redteam_result: RedTeamResult,
+        eval_run: EvalRun,
+        red_team_info: Dict,
+        _skip_evals: bool = False,
+    ) -> Optional[str]:
+        """Log the Red Team Agent results to MLFlow.
+
+        :param redteam_result: The output from the red team agent evaluation
+        :type redteam_result: RedTeamResult
+        :param eval_run: The MLFlow run object
+        :type eval_run: EvalRun
+        :param red_team_info: Red team tracking information
+        :type red_team_info: Dict
+        :param _skip_evals: Whether to log only data without evaluation results
+        :type _skip_evals: bool
+        :return: The URL to the run in Azure AI Studio, if available
+        :rtype: Optional[str]
+        """
+        self.logger.debug(f"Logging results to MLFlow, _skip_evals={_skip_evals}")
+        artifact_name = "instance_results.json"
+        eval_info_name = "redteam_info.json"
+        properties = {}
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            if self.scan_output_dir:
+                artifact_path = os.path.join(self.scan_output_dir, artifact_name)
+                self.logger.debug(f"Saving artifact to scan output directory: {artifact_path}")
+                with open(artifact_path, "w", encoding=DefaultOpenEncoding.WRITE) as f:
+                    if _skip_evals:
+                        # In _skip_evals mode, we write the conversations in conversation/messages format
+                        f.write(json.dumps({"conversations": redteam_result.attack_details or []}))
+                    elif redteam_result.scan_result:
+                        # Create a copy to avoid modifying the original scan result
+                        result_with_conversations = (
+                            redteam_result.scan_result.copy() if isinstance(redteam_result.scan_result, dict) else {}
+                        )
+
+                        # Preserve all original fields needed for scorecard generation
+                        result_with_conversations["scorecard"] = result_with_conversations.get("scorecard", {})
+                        result_with_conversations["parameters"] = result_with_conversations.get("parameters", {})
+
+                        # Add conversations field with all conversation data including user messages
+                        result_with_conversations["conversations"] = redteam_result.attack_details or []
+
+                        # Keep original attack_details field to preserve compatibility with existing code
+                        if (
+                            "attack_details" not in result_with_conversations
+                            and redteam_result.attack_details is not None
+                        ):
+                            result_with_conversations["attack_details"] = redteam_result.attack_details
+
+                        json.dump(result_with_conversations, f)
+
+                eval_info_path = os.path.join(self.scan_output_dir, eval_info_name)
+                self.logger.debug(f"Saving evaluation info to scan output directory: {eval_info_path}")
+                with open(eval_info_path, "w", encoding=DefaultOpenEncoding.WRITE) as f:
+                    # Remove evaluation_result from red_team_info before logging
+                    red_team_info_logged = {}
+                    for strategy, harms_dict in red_team_info.items():
+                        red_team_info_logged[strategy] = {}
+                        for harm, info_dict in harms_dict.items():
+                            # Create a copy to avoid modifying the original
+                            info_dict_copy = info_dict.copy()
+                            info_dict_copy.pop("evaluation_result", None)
+                            red_team_info_logged[strategy][harm] = info_dict_copy
+                    f.write(json.dumps(red_team_info_logged, indent=2))
+                self.logger.debug(f"Successfully wrote redteam_info.json to: {eval_info_path}")
+
+                # Also save a human-readable scorecard if available
+                if not _skip_evals and redteam_result.scan_result:
+                    from ._utils.formatting_utils import format_scorecard
+
+                    scorecard_path = os.path.join(self.scan_output_dir, "scorecard.txt")
+                    with open(scorecard_path, "w", encoding=DefaultOpenEncoding.WRITE) as f:
+                        f.write(format_scorecard(redteam_result.scan_result))
+                    self.logger.debug(f"Saved scorecard to: {scorecard_path}")
+
+                # Create a dedicated artifacts directory with proper structure for MLFlow
+                # First, create the main artifact file that MLFlow expects
+                with open(
+                    os.path.join(tmpdir, artifact_name),
+                    "w",
+                    encoding=DefaultOpenEncoding.WRITE,
+                ) as f:
+                    if _skip_evals:
+                        f.write(json.dumps({"conversations": redteam_result.attack_details or []}))
+                    elif redteam_result.scan_result:
+                        json.dump(redteam_result.scan_result, f)
+
+                # Copy all relevant files to the temp directory
+                import shutil
+
+                for file in os.listdir(self.scan_output_dir):
+                    file_path = os.path.join(self.scan_output_dir, file)
+
+                    # Skip directories and log files if not in debug mode
+                    if os.path.isdir(file_path):
+                        continue
+                    if file.endswith(".log") and not os.environ.get("DEBUG"):
+                        continue
+                    if file.endswith(".gitignore"):
+                        continue
+                    if file == artifact_name:
+                        continue
+
+                    try:
+                        shutil.copy(file_path, os.path.join(tmpdir, file))
+                        self.logger.debug(f"Copied file to artifact directory: {file}")
+                    except Exception as e:
+                        self.logger.warning(f"Failed to copy file {file} to artifact directory: {str(e)}")
+
+                properties.update({"scan_output_dir": str(self.scan_output_dir)})
+            else:
+                # Use temporary directory as before if no scan output directory exists
+                artifact_file = Path(tmpdir) / artifact_name
+                with open(artifact_file, "w", encoding=DefaultOpenEncoding.WRITE) as f:
+                    if _skip_evals:
+                        f.write(json.dumps({"conversations": redteam_result.attack_details or []}))
+                    elif redteam_result.scan_result:
+                        json.dump(redteam_result.scan_result, f)
+                self.logger.debug(f"Logged artifact: {artifact_name}")
+
+            properties.update(
+                {
+                    "redteaming": "asr",
+                    EvaluationRunProperties.EVALUATION_SDK: f"azure-ai-evaluation:{VERSION}",
+                }
+            )
+
+            metrics = {}
+            if redteam_result.scan_result:
+                scorecard = redteam_result.scan_result["scorecard"]
+                joint_attack_summary = scorecard["joint_risk_attack_summary"]
+
+                if joint_attack_summary:
+                    for risk_category_summary in joint_attack_summary:
+                        risk_category = risk_category_summary.get("risk_category").lower()
+                        for key, value in risk_category_summary.items():
+                            if key != "risk_category":
+                                metrics.update({f"{risk_category}_{key}": cast(float, value)})
+                                self.logger.debug(f"Logged metric: {risk_category}_{key} = {value}")
+
+            if self._one_dp_project:
+                try:
+                    create_evaluation_result_response = (
+                        self.generated_rai_client._evaluation_onedp_client.create_evaluation_result(
+                            name=uuid.uuid4(),
+                            path=tmpdir,
+                            metrics=metrics,
+                            result_type=ResultType.REDTEAM,
+                        )
+                    )
+
+                    update_run_response = self.generated_rai_client._evaluation_onedp_client.update_red_team_run(
+                        name=eval_run.id,
+                        red_team=RedTeamUpload(
+                            id=eval_run.id,
+                            display_name=eval_run.display_name
+                            or f"redteam-agent-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
+                            status="Completed",
+                            outputs={
+                                "evaluationResultId": create_evaluation_result_response.id,
+                            },
+                            properties=properties,
+                        ),
+                    )
+                    self.logger.debug(f"Updated UploadRun: {update_run_response.id}")
+                except Exception as e:
+                    self.logger.warning(f"Failed to upload red team results to AI Foundry: {str(e)}")
+            else:
+                # Log the entire directory to MLFlow
+                try:
+                    eval_run.log_artifact(tmpdir, artifact_name)
+                    if self.scan_output_dir:
+                        eval_run.log_artifact(tmpdir, eval_info_name)
+                    self.logger.debug(f"Successfully logged artifacts directory to AI Foundry")
+                except Exception as e:
+                    self.logger.warning(f"Failed to log artifacts to AI Foundry: {str(e)}")
+
+                for k, v in metrics.items():
+                    eval_run.log_metric(k, v)
+                    self.logger.debug(f"Logged metric: {k} = {v}")
+
+                eval_run.write_properties_to_run_history(properties)
+                eval_run._end_run("FINISHED")
+
+        self.logger.info("Successfully logged results to AI Foundry")
+        return None