azure-ai-evaluation 1.11.1__py3-none-any.whl → 1.12.0__py3-none-any.whl

azure/ai/evaluation/red_team/_mlflow_integration.py

@@ -12,7 +12,7 @@ import os
 import tempfile
 import uuid
 from datetime import datetime
-from typing import Dict, Optional, cast
+from typing import Any, Dict, List, Optional, Set, cast
 from pathlib import Path
 
 # Azure AI Evaluation imports
@@ -27,7 +27,14 @@ 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 ._red_team_result import (
+    RedTeamResult,
+    RedTeamRun,
+    ResultCount,
+    PerTestingCriteriaResult,
+    DataSource,
+    OutputItemsList,
+)
 from ._utils.logging_utils import log_error
 
 
@@ -50,6 +57,32 @@ class MLflowIntegration:
         self.scan_output_dir = scan_output_dir
         self.ai_studio_url = None
         self.trace_destination = None
+        self._run_id_override: Optional[str] = None
+        self._eval_id_override: Optional[str] = None
+        self._created_at_override: Optional[int] = None
+
+    def set_run_identity_overrides(
+        self,
+        *,
+        run_id: Optional[str] = None,
+        eval_id: Optional[str] = None,
+        created_at: Optional[Any] = None,
+    ) -> None:
+        """Allow callers to supply pre-existing identifiers for the run payload."""
+
+        self._run_id_override = str(run_id).strip() if run_id else None
+        self._eval_id_override = str(eval_id).strip() if eval_id else None
+
+        if created_at is None or created_at == "":
+            self._created_at_override = None
+        else:
+            if isinstance(created_at, datetime):
+                self._created_at_override = int(created_at.timestamp())
+            else:
+                try:
+                    self._created_at_override = int(created_at)
+                except (TypeError, ValueError):
+                    self._created_at_override = None
 
     def start_redteam_mlflow_run(
         self,
@@ -152,38 +185,36 @@ class MLflowIntegration:
         """
         self.logger.debug(f"Logging results to MLFlow, _skip_evals={_skip_evals}")
         artifact_name = "instance_results.json"
+        results_name = "results.json"
         eval_info_name = "redteam_info.json"
         properties = {}
 
         with tempfile.TemporaryDirectory() as tmpdir:
             if self.scan_output_dir:
+                # Save new format as results.json
+                results_path = os.path.join(self.scan_output_dir, results_name)
+                self.logger.debug(f"Saving results to scan output directory: {results_path}")
+                with open(results_path, "w", encoding=DefaultOpenEncoding.WRITE) as f:
+                    payload = self._build_results_payload(
+                        redteam_result=redteam_result,
+                        eval_run=eval_run,
+                        red_team_info=red_team_info,
+                        include_conversations=True,
+                        scan_name=getattr(eval_run, "display_name", None),
+                    )
+                    json.dump(payload, f)
+
+                # Save legacy format as instance_results.json
                 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)
+                    legacy_payload = self._build_instance_results_payload(
+                        redteam_result=redteam_result,
+                        eval_run=eval_run,
+                        red_team_info=red_team_info,
+                        scan_name=getattr(eval_run, "display_name", None),
+                    )
+                    json.dump(legacy_payload, 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}")
@@ -210,16 +241,34 @@ class MLflowIntegration:
                     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
+                # First, create the main artifact file that MLFlow expects (new format)
+                with open(
+                    os.path.join(tmpdir, results_name),
+                    "w",
+                    encoding=DefaultOpenEncoding.WRITE,
+                ) as f:
+                    payload = self._build_results_payload(
+                        redteam_result=redteam_result,
+                        eval_run=eval_run,
+                        red_team_info=red_team_info,
+                        include_conversations=False,
+                        scan_name=getattr(eval_run, "display_name", None),
+                    )
+                    json.dump(payload, f)
+
+                # Also create legacy instance_results.json for compatibility
                 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)
+                    legacy_payload = self._build_instance_results_payload(
+                        redteam_result=redteam_result,
+                        eval_run=eval_run,
+                        red_team_info=red_team_info,
+                        scan_name=getattr(eval_run, "display_name", None),
+                    )
+                    json.dump(legacy_payload, f)
 
                 # Copy all relevant files to the temp directory
                 import shutil
@@ -246,12 +295,28 @@ class MLflowIntegration:
                 properties.update({"scan_output_dir": str(self.scan_output_dir)})
             else:
                 # Use temporary directory as before if no scan output directory exists
+                results_file = Path(tmpdir) / results_name
+                with open(results_file, "w", encoding=DefaultOpenEncoding.WRITE) as f:
+                    payload = self._build_results_payload(
+                        redteam_result=redteam_result,
+                        eval_run=eval_run,
+                        red_team_info=red_team_info,
+                        include_conversations=_skip_evals,
+                        scan_name=getattr(eval_run, "display_name", None),
+                    )
+                    json.dump(payload, f)
+                self.logger.debug(f"Logged artifact: {results_name}")
+
+                # Also create legacy instance_results.json
                 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)
+                    legacy_payload = self._build_instance_results_payload(
+                        redteam_result=redteam_result,
+                        eval_run=eval_run,
+                        red_team_info=red_team_info,
+                        scan_name=getattr(eval_run, "display_name", None),
+                    )
+                    json.dump(legacy_payload, f)
                 self.logger.debug(f"Logged artifact: {artifact_name}")
 
             properties.update(
@@ -320,3 +385,357 @@ class MLflowIntegration:
 
         self.logger.info("Successfully logged results to AI Foundry")
         return None
+
+    @staticmethod
+    def _compute_result_count(output_items: List[Dict[str, Any]]) -> Dict[str, int]:
+        """Aggregate run-level pass/fail counts from individual output items."""
+
+        total = len(output_items)
+        passed = failed = errored = 0
+
+        for item in output_items:
+            item_status: Optional[bool] = None
+            for result in item.get("results", []):
+                result_properties = result.get("properties", {}) if isinstance(result, dict) else {}
+                attack_success = result_properties.get("attack_success")
+                if attack_success is True:
+                    item_status = False
+                    break
+                if attack_success is False:
+                    item_status = True
+                elif item_status is None and result.get("passed") is not None:
+                    item_status = bool(result.get("passed"))
+
+            if item_status is True:
+                passed += 1
+            elif item_status is False:
+                failed += 1
+            else:
+                errored += 1
+
+        return {
+            "total": total,
+            "passed": passed,
+            "failed": failed,
+            "errored": errored,
+        }
+
+    @staticmethod
+    def _compute_per_testing_criteria(output_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Build aggregated pass/fail counts per testing criteria (risk category)."""
+
+        criteria: Dict[str, Dict[str, int]] = {}
+
+        for item in output_items:
+            for result in item.get("results", []):
+                if not isinstance(result, dict):
+                    continue
+                name = result.get("name")
+                if not name:
+                    continue
+                passed_value = result.get("passed")
+                if passed_value is None:
+                    continue
+
+                bucket = criteria.setdefault(str(name), {"passed": 0, "failed": 0})
+                if passed_value:
+                    bucket["passed"] += 1
+                else:
+                    bucket["failed"] += 1
+
+        return [
+            {
+                "testing_criteria": criteria_name,
+                "passed": counts["passed"],
+                "failed": counts["failed"],
+            }
+            for criteria_name, counts in sorted(criteria.items())
+        ]
+
+    @staticmethod
+    def _build_data_source_section(parameters: Dict[str, Any], red_team_info: Optional[Dict]) -> Dict[str, Any]:
+        """Build the data_source portion of the run payload for red-team scans."""
+
+        attack_strategies: List[str] = []
+        if isinstance(red_team_info, dict):
+            attack_strategies = sorted(str(strategy) for strategy in red_team_info.keys())
+
+        item_generation_params: Dict[str, Any] = {"type": "red_team"}
+        if attack_strategies:
+            item_generation_params["attack_strategies"] = attack_strategies
+
+        # Attempt to infer turns from parameters if available
+        num_turns = parameters.get("max_turns") if isinstance(parameters, dict) else None
+        if isinstance(num_turns, int) and num_turns > 0:
+            item_generation_params["num_turns"] = num_turns
+
+        data_source: Dict[str, Any] = {"type": "azure_ai_red_team", "target": {}}
+        if item_generation_params:
+            data_source["item_generation_params"] = item_generation_params
+
+        return data_source
+
+    def _determine_run_status(
+        self,
+        scan_result: Dict[str, Any],
+        red_team_info: Optional[Dict],
+        output_items: List[Dict[str, Any]],
+    ) -> str:
+        """Determine the run-level status based on red team info status values."""
+
+        # Check if any tasks are still incomplete/failed
+        if isinstance(red_team_info, dict):
+            for risk_data in red_team_info.values():
+                if not isinstance(risk_data, dict):
+                    continue
+                for details in risk_data.values():
+                    if not isinstance(details, dict):
+                        continue
+                    status = details.get("status", "").lower()
+                    if status in ("incomplete", "failed", "timeout"):
+                        return "failed"
+                    elif status in ("running", "pending"):
+                        return "in_progress"
+
+        return "completed"
+
+    def _build_results_payload(
+        self,
+        redteam_result: RedTeamResult,
+        eval_run: Optional[Any] = None,
+        red_team_info: Optional[Dict] = None,
+        include_conversations: bool = False,
+        scan_name: Optional[str] = None,
+    ) -> RedTeamRun:
+        """Assemble the new structure for results.json with eval.run format."""
+
+        scan_result = cast(Dict[str, Any], redteam_result.scan_result or {})
+        output_items = cast(List[Dict[str, Any]], scan_result.get("output_items") or [])
+        scorecard = cast(Dict[str, Any], scan_result.get("scorecard") or {})
+        parameters = cast(Dict[str, Any], scan_result.get("parameters") or {})
+
+        run_id = self._run_id_override
+        eval_id = self._eval_id_override
+        run_name: Optional[str] = None
+        created_at = self._created_at_override
+
+        if eval_run is not None:
+            run_info = getattr(eval_run, "info", None)
+
+            if run_id is None:
+                candidate_run_id = (
+                    getattr(run_info, "run_id", None)
+                    or getattr(eval_run, "run_id", None)
+                    or getattr(eval_run, "id", None)
+                )
+                if candidate_run_id is not None:
+                    run_id = str(candidate_run_id)
+
+            if eval_id is None:
+                candidate_eval_id = (
+                    getattr(run_info, "experiment_id", None)
+                    or getattr(eval_run, "experiment_id", None)
+                    or getattr(eval_run, "eval_id", None)
+                )
+                if candidate_eval_id is not None:
+                    eval_id = str(candidate_eval_id)
+
+            if run_name is None:
+                candidate_run_name = (
+                    getattr(run_info, "run_name", None)
+                    or getattr(eval_run, "run_name", None)
+                    or getattr(eval_run, "display_name", None)
+                    or getattr(eval_run, "name", None)
+                )
+                if candidate_run_name is not None:
+                    run_name = str(candidate_run_name)
+
+            if created_at is None:
+                raw_created = (
+                    getattr(run_info, "created_time", None)
+                    or getattr(eval_run, "created_at", None)
+                    or getattr(eval_run, "created_time", None)
+                )
+                if isinstance(raw_created, datetime):
+                    created_at = int(raw_created.timestamp())
+                elif isinstance(raw_created, (int, float)):
+                    created_at = int(raw_created)
+                elif isinstance(raw_created, str):
+                    try:
+                        created_at = int(float(raw_created))
+                    except ValueError:
+                        created_at = None
+
+        if run_id is None:
+            run_id = str(uuid.uuid4())
+        if eval_id is None:
+            eval_id = str(uuid.uuid4())
+        if created_at is None:
+            created_at = int(datetime.now().timestamp())
+        if run_name is None:
+            run_name = scan_name or f"redteam-run-{run_id[:8]}"
+
+        result_count = self._compute_result_count(output_items)
+        per_testing_results = self._compute_per_testing_criteria(output_items)
+        data_source = self._build_data_source_section(parameters, red_team_info)
+        status = self._determine_run_status(scan_result, red_team_info, output_items)
+
+        list_wrapper: OutputItemsList = {
+            "object": "list",
+            "data": output_items,
+        }
+
+        run_payload: RedTeamRun = {
+            "object": "eval.run",
+            "id": run_id,
+            "eval_id": eval_id,
+            "created_at": created_at,
+            "status": status,
+            "name": run_name,
+            "report_url": scan_result.get("studio_url") or self.ai_studio_url,
+            "data_source": data_source,
+            "metadata": {},
+            "result_count": result_count,
+            "per_model_usage": [],
+            "per_testing_criteria_results": per_testing_results,
+            "output_items": list_wrapper,
+        }
+
+        if include_conversations:
+            run_payload["conversations"] = redteam_result.attack_details or scan_result.get("attack_details") or []
+
+        return run_payload
+
+    def _build_results_payload(
+        self,
+        redteam_result: RedTeamResult,
+        eval_run: Optional[Any] = None,
+        red_team_info: Optional[Dict] = None,
+        include_conversations: bool = False,
+        scan_name: Optional[str] = None,
+    ) -> RedTeamRun:
+        """Assemble the new structure for results.json with eval.run format."""
+
+        scan_result = cast(Dict[str, Any], redteam_result.scan_result or {})
+        output_items = cast(List[Dict[str, Any]], scan_result.get("output_items") or [])
+        scorecard = cast(Dict[str, Any], scan_result.get("scorecard") or {})
+        parameters = cast(Dict[str, Any], scan_result.get("parameters") or {})
+
+        run_id = self._run_id_override
+        eval_id = self._eval_id_override
+        run_name: Optional[str] = None
+        created_at = self._created_at_override
+
+        if eval_run is not None:
+            run_info = getattr(eval_run, "info", None)
+
+            if run_id is None:
+                candidate_run_id = (
+                    getattr(run_info, "run_id", None)
+                    or getattr(eval_run, "run_id", None)
+                    or getattr(eval_run, "id", None)
+                )
+                if candidate_run_id is not None:
+                    run_id = str(candidate_run_id)
+
+            if eval_id is None:
+                candidate_eval_id = (
+                    getattr(run_info, "experiment_id", None)
+                    or getattr(eval_run, "experiment_id", None)
+                    or getattr(eval_run, "eval_id", None)
+                )
+                if candidate_eval_id is not None:
+                    eval_id = str(candidate_eval_id)
+
+            if run_name is None:
+                candidate_run_name = (
+                    getattr(run_info, "run_name", None)
+                    or getattr(eval_run, "run_name", None)
+                    or getattr(eval_run, "display_name", None)
+                    or getattr(eval_run, "name", None)
+                )
+                if candidate_run_name is not None:
+                    run_name = str(candidate_run_name)
+
+            if created_at is None:
+                raw_created = (
+                    getattr(run_info, "created_time", None)
+                    or getattr(eval_run, "created_at", None)
+                    or getattr(eval_run, "created_time", None)
+                )
+                if isinstance(raw_created, datetime):
+                    created_at = int(raw_created.timestamp())
+                elif isinstance(raw_created, (int, float)):
+                    created_at = int(raw_created)
+                elif isinstance(raw_created, str):
+                    try:
+                        created_at = int(float(raw_created))
+                    except ValueError:
+                        created_at = None
+
+        if run_id is None:
+            run_id = str(uuid.uuid4())
+        if eval_id is None:
+            eval_id = str(uuid.uuid4())
+        if created_at is None:
+            created_at = int(datetime.now().timestamp())
+        if run_name is None:
+            run_name = scan_name or f"redteam-run-{run_id[:8]}"
+
+        result_count = self._compute_result_count(output_items)
+        per_testing_results = self._compute_per_testing_criteria(output_items)
+        data_source = self._build_data_source_section(parameters, red_team_info)
+        status = self._determine_run_status(scan_result, red_team_info, output_items)
+
+        list_wrapper: OutputItemsList = {
+            "object": "list",
+            "data": output_items,
+        }
+
+        run_payload: RedTeamRun = {
+            "object": "eval.run",
+            "id": run_id,
+            "eval_id": eval_id,
+            "created_at": created_at,
+            "status": status,
+            "name": run_name,
+            "report_url": scan_result.get("studio_url") or self.ai_studio_url,
+            "data_source": data_source,
+            "metadata": {},
+            "result_count": result_count,
+            "per_model_usage": [],
+            "per_testing_criteria_results": per_testing_results,
+            "output_items": list_wrapper,
+        }
+
+        if include_conversations:
+            run_payload["conversations"] = redteam_result.attack_details or scan_result.get("attack_details") or []
+
+        return run_payload
+
+    def _build_instance_results_payload(
+        self,
+        redteam_result: RedTeamResult,
+        eval_run: Optional[Any] = None,
+        red_team_info: Optional[Dict] = None,
+        scan_name: Optional[str] = None,
+    ) -> Dict:
+        """Assemble the legacy structure for instance_results.json (scan_result format)."""
+
+        scan_result = cast(Dict[str, Any], redteam_result.scan_result or {})
+
+        # Return the scan_result directly for legacy compatibility
+        # This maintains the old format that was expected previously
+        legacy_payload = scan_result.copy() if scan_result else {}
+
+        # Ensure we have the basic required fields
+        if "scorecard" not in legacy_payload:
+            legacy_payload["scorecard"] = {}
+        if "parameters" not in legacy_payload:
+            legacy_payload["parameters"] = {}
+        if "output_items" not in legacy_payload:
+            legacy_payload["output_items"] = []
+        if "attack_details" not in legacy_payload:
+            legacy_payload["attack_details"] = redteam_result.attack_details or []
+
+        return legacy_payload

azure/ai/evaluation/red_team/_red_team.py

@@ -783,6 +783,9 @@ class RedTeam:
         :rtype: RedTeamResult
         """
         user_agent: Optional[str] = kwargs.get("user_agent", "(type=redteam; subtype=RedTeam)")
+        run_id_override = kwargs.get("run_id") or kwargs.get("runId")
+        eval_id_override = kwargs.get("eval_id") or kwargs.get("evalId")
+        created_at_override = kwargs.get("created_at") or kwargs.get("createdAt")
         with UserAgentSingleton().add_useragent_product(user_agent):
             # Initialize scan
             self._initialize_scan(scan_name, application_scenario)
@@ -802,6 +805,12 @@ class RedTeam:
             self.mlflow_integration.logger = self.logger
             self.result_processor.logger = self.logger
 
+            self.mlflow_integration.set_run_identity_overrides(
+                run_id=run_id_override,
+                eval_id=eval_id_override,
+                created_at=created_at_override,
+            )
+
             # Validate attack objective generator
             if not self.attack_objective_generator:
                 raise EvaluationException(