ibm-watsonx-orchestrate-evaluation-framework 1.1.5__py3-none-any.whl → 1.1.6__py3-none-any.whl

wxo_agentic_evaluation/red_teaming/attack_runner.py

@@ -12,32 +12,36 @@ from rich.progress import Progress
 
 from wxo_agentic_evaluation.arg_configs import AttackConfig
 from wxo_agentic_evaluation.inference_backend import (
-    EvaluationController,
+    AttackEvaluationController,
     WXOInferenceBackend,
-    get_wxo_client,
 )
+from wxo_agentic_evaluation.wxo_client import get_wxo_client
 from wxo_agentic_evaluation.llm_user import LLMUser
 from wxo_agentic_evaluation.prompt.template_render import (
     LlamaUserTemplateRenderer,
 )
-from wxo_agentic_evaluation.red_teaming.attack_evaluator import AttackEvaluator
+from wxo_agentic_evaluation.red_teaming.attack_evaluator import AttackEvaluator, evaluate_all_attacks
 from wxo_agentic_evaluation.resource_map import ResourceMap
 from wxo_agentic_evaluation.service_provider import get_provider
 from wxo_agentic_evaluation.type import AttackData
 from wxo_agentic_evaluation.utils import json_dump
 
 
-def process_attack(task_n, attack_path, config, inference_backend, llm_user):
-    tc_name = os.path.basename(attack_path).replace(".json", "")
+def process_attack(task_n, attack_path, config, inference_backend, llm_user, resource_map):
+    attack_filename = os.path.basename(attack_path).replace(".json", "")
     with open(attack_path, "r") as f:
         attack: AttackData = AttackData.model_validate(json.load(f))
 
-    evaluation_controller = EvaluationController(
+    attack_evaluator = AttackEvaluator(config, resource_map, attack_filename)
+    
+    evaluation_controller = AttackEvaluationController(
         wxo_inference_backend=inference_backend,
         llm_user=llm_user,
         config=config,
+        attack_data=attack,
+        attack_evaluator=attack_evaluator
     )
-    rich.print(f"[bold magenta]Running attack: {tc_name}[/bold magenta]")
+    rich.print(f"[bold magenta]Running attack: {attack_filename}[/bold magenta]")
     history, _, _ = evaluation_controller.run(
         task_n,
         attack.story,
@@ -50,10 +54,25 @@ def process_attack(task_n, attack_path, config, inference_backend, llm_user):
         result.append(message.model_dump())
 
     json_dump(
-        os.path.join(config.output_dir, "messages", tc_name + ".messages.json"),
+        os.path.join(config.output_dir, "messages", attack_filename + ".messages.json"),
         result,
     )
 
+    # Ensure there's a results dir and write a failure record if none was created
+    result_path = os.path.join(config.output_dir, "results", attack_filename + ".result.json")
+    if not os.path.exists(result_path):
+        # attack evaluator should have written a success on early termination; if not, mark as failed
+        json_dump(
+            result_path, 
+            {
+                "attack_filename": attack_filename, 
+                "success": False,
+                "attack_category": str(attack.attack_data.attack_category),
+                "attack_name": getattr(attack.attack_data, "attack_name", ""),
+                "attack_type": getattr(attack.attack_data, "attack_type", ""),
+            }
+        )
+
     return result
 
 
@@ -80,7 +99,21 @@ def run_attacks(config: AttackConfig):
     print(
         f"Running red teaming attacks with tenant {config.auth_config.tenant_name}"
     )
-    os.makedirs(os.path.join(config.output_dir, "messages"), exist_ok=True)
+    for folder in ["messages", "results", "evaluations"]:
+        os.makedirs(os.path.join(config.output_dir, folder), exist_ok=True)
+    
+    available_res = set()
+    if config.skip_available_results:
+        available_res = set(
+            [
+                os.path.basename(f).replace(".result", "")
+                for f in glob.glob(
+                    os.path.join(
+                        config.output_dir, "results", "*.result.json"
+                    )
+                )
+            ]
+        )
 
     results_list = []
     attack_paths = []
@@ -98,6 +131,13 @@ def run_attacks(config: AttackConfig):
         ):
             continue
 
+        if config.skip_available_results:
+            if os.path.basename(attack_path) in available_res:
+                print(
+                    f"Skipping attack {os.path.basename(attack_path)} as results already exist."
+                )
+                continue
+
         future = executor.submit(
             process_attack,
             task_n,
@@ -105,6 +145,7 @@ def run_attacks(config: AttackConfig):
             config,
             inference_backend,
             llm_user,
+            resource_map
         )
 
         futures.append((attack_path, future))
@@ -124,8 +165,7 @@ def run_attacks(config: AttackConfig):
                 finally:
                     progress.update(task1, advance=1)
 
-    attack_evaluator = AttackEvaluator(config, resource_map)
-    attack_results = attack_evaluator.evaluate_attacks()
+    attack_results = evaluate_all_attacks(config, resource_map)
 
     with open(
         os.path.join(config.output_dir, "config.yml"), "w", encoding="utf-8"