ibm-watsonx-orchestrate-evaluation-framework 1.0.3__py3-none-any.whl → 1.1.8b0__py3-none-any.whl

wxo_agentic_evaluation/red_teaming/attack_runner.py

@@ -0,0 +1,204 @@
+import dataclasses
+import glob
+import json
+import os
+import traceback
+from concurrent.futures import ThreadPoolExecutor
+
+import rich
+import yaml
+from jsonargparse import CLI
+from rich.progress import Progress
+
+from wxo_agentic_evaluation.arg_configs import AttackConfig
+from wxo_agentic_evaluation.evaluation_controller.evaluation_controller import (
+    AttackEvaluationController,
+    EvaluationController,
+)
+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,
+    evaluate_all_attacks,
+)
+from wxo_agentic_evaluation.resource_map import ResourceMap
+from wxo_agentic_evaluation.runtime_adapter.wxo_runtime_adapter import (
+    WXORuntimeAdapter,
+)
+from wxo_agentic_evaluation.service_provider import (
+    USE_GATEWAY_MODEL_PROVIDER,
+    get_provider,
+)
+from wxo_agentic_evaluation.type import AttackData
+from wxo_agentic_evaluation.utils import json_dump
+from wxo_agentic_evaluation.wxo_client import get_wxo_client
+
+
+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))
+
+    attack_evaluator = AttackEvaluator(config, resource_map, attack_filename)
+
+    evaluation_controller = AttackEvaluationController(
+        runtime=inference_backend,
+        llm_user=llm_user,
+        config=config,
+        attack_data=attack,
+        attack_evaluator=attack_evaluator,
+    )
+    rich.print(
+        f"[bold magenta]Running attack: {attack_filename}[/bold magenta]"
+    )
+    history, _, _ = evaluation_controller.run(
+        task_n,
+        attack.story,
+        agent_name=attack.agent,
+        starting_user_input=attack.starting_sentence,
+        attack_instructions=attack.attack_data.attack_instructions,
+    )
+    result = list()
+    for message in history:
+        result.append(message.model_dump())
+
+    json_dump(
+        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
+
+
+def run_attacks(config: AttackConfig):
+    executor = ThreadPoolExecutor(max_workers=config.num_workers)
+    wxo_client = get_wxo_client(
+        config.auth_config.url,
+        config.auth_config.tenant_name,
+        config.auth_config.token,
+    )
+    resource_map = ResourceMap(wxo_client)
+    inference_backend = WXORuntimeAdapter(wxo_client=wxo_client)
+    provider_kwargs = {}
+    if USE_GATEWAY_MODEL_PROVIDER:
+        provider_kwargs.update(
+            instance_url=wxo_client.service_url,
+            token=wxo_client.api_key,
+        )
+    llm_user = LLMUser(
+        wai_client=get_provider(
+            model_id=config.llm_user_config.model_id, **provider_kwargs
+        ),
+        template=LlamaUserTemplateRenderer(
+            config.llm_user_config.prompt_config
+        ),
+        user_response_style=config.llm_user_config.user_response_style,
+    )
+
+    print(
+        f"Running red teaming attacks with tenant {config.auth_config.tenant_name}"
+    )
+    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 = []
+    for path in config.attack_paths:
+        if os.path.isdir(path):
+            path = os.path.join(path, "*.json")
+        attack_paths.extend(sorted(glob.glob(path)))
+
+    futures = []
+    task_n = 0
+
+    for attack_path in attack_paths:
+        if not attack_path.endswith(".json") or attack_path.endswith(
+            "agent.json"
+        ):
+            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,
+            attack_path,
+            config,
+            inference_backend,
+            llm_user,
+            resource_map,
+        )
+
+        futures.append((attack_path, future))
+        task_n += 1
+
+    if futures:
+        with Progress() as progress:
+            task1 = progress.add_task(
+                f"[purple]Running {len(futures)} attacks...", total=len(futures)
+            )
+            for attack_path, future in futures:
+                try:
+                    results_list.extend(future.result())
+                except Exception as e:
+                    rich.print(f"Attack {attack_path} fails with {e}")
+                    traceback.print_exc()
+                finally:
+                    progress.update(task1, advance=1)
+
+    attack_results = evaluate_all_attacks(config, resource_map)
+
+    with open(
+        os.path.join(config.output_dir, "config.yml"), "w", encoding="utf-8"
+    ) as f:
+        yaml.safe_dump(dataclasses.asdict(config), f)
+
+    with open(
+        os.path.join(config.output_dir, "attacks_results.json"), "w"
+    ) as f:
+        json.dump(attack_results, f, indent=2)
+
+    print(f"Attack results saved to {config.output_dir}")
+
+
+if __name__ == "__main__":
+    run_attacks(CLI(AttackConfig, as_positional=False))

wxo_agentic_evaluation/referenceless_eval/__init__.py

@@ -0,0 +1,3 @@
+from wxo_agentic_evaluation.referenceless_eval.referenceless_eval import (
+    ReferencelessEvaluation,
+)

wxo_agentic_evaluation/referenceless_eval/function_calling/consts.py

@@ -0,0 +1,28 @@
+### Metric name constants
+## General metrics
+METRIC_GENERAL_HALLUCINATION_CHECK = "general_hallucination_check"
+METRIC_GENERAL_VALUE_FORMAT_ALIGNMENT = "general_value_format_alignment"
+
+## Function selection metrics
+METRIC_FUNCTION_SELECTION_APPROPRIATENESS = "function_selection_appropriateness"
+METRIC_AGENTIC_CONSTRAINTS_SATISFACTION = "agentic_constraints_satisfaction"
+
+## Parameter metrics
+METRIC_PARAMETER_VALUE_FORMAT_ALIGNMENT = "parameter_value_format_alignment"
+METRIC_PARAMETER_HALLUCINATION_CHECK = "parameter_hallucination_check"
+
+## Metric category mapping
+GENERAL_METRICS = [
+    METRIC_GENERAL_HALLUCINATION_CHECK,
+    METRIC_GENERAL_VALUE_FORMAT_ALIGNMENT,
+]
+
+FUNCTION_SELECTION_METRICS = [
+    METRIC_FUNCTION_SELECTION_APPROPRIATENESS,
+    METRIC_AGENTIC_CONSTRAINTS_SATISFACTION,
+]
+
+PARAMETER_METRICS = [
+    METRIC_PARAMETER_VALUE_FORMAT_ALIGNMENT,
+    METRIC_PARAMETER_HALLUCINATION_CHECK,
+]

wxo_agentic_evaluation/referenceless_eval/function_calling/metrics/base.py

@@ -0,0 +1,29 @@
+from abc import ABC
+
+from wxo_agentic_evaluation.referenceless_eval.metrics.metric import Metric
+from wxo_agentic_evaluation.referenceless_eval.metrics.prompt import (
+    MetricPrompt,
+)
+
+
+class FunctionMetricsPrompt(MetricPrompt, ABC):
+    """
+    Abstract base for function-calling metric prompts.
+    Subclasses must define class attrs:
+      - system_template: str
+      - user_template: str
+    """
+
+    system_template: str
+    user_template: str
+
+    def __init__(self, metric: Metric, task_description: str) -> None:
+        super().__init__(
+            metric=metric,
+            system_template=self.system_template,
+            user_template=self.user_template,
+            system_kwargs_defaults={
+                "task_description": task_description,
+                "metric_jsonschema": metric.to_jsonschema(),
+            },
+        )

wxo_agentic_evaluation/referenceless_eval/function_calling/metrics/function_call/general.py

@@ -0,0 +1,49 @@
+from typing import Any, Dict, List, Union
+
+from wxo_agentic_evaluation.referenceless_eval.function_calling.metrics.base import (
+    FunctionMetricsPrompt,
+)
+
+_general_system = (
+    "### Task Description and Role:\n\n"
+    "{{ task_description }}\n\n"
+    "Your output must conform to the following JSON schema, in the same order as the fields appear in the schema:\n"
+    "{{ metric_jsonschema }}"
+)
+
+_general_user: str = (
+    "Conversation context:\n"
+    "{{ conversation_context }}\n\n"
+    "Tool Specification:\n"
+    "{{ tool_inventory }}\n\n"
+    "Proposed tool call:\n"
+    "{{ tool_call }}\n\n"
+    "Return a JSON object as specified in the system prompt. You MUST keep the same order of fields in the JSON object as provided in the JSON schema and examples."
+)
+
+
+class GeneralMetricsPrompt(FunctionMetricsPrompt):
+    """Prompt builder for general tool-call semantic metrics."""
+
+    system_template = _general_system
+    user_template = _general_user
+
+
+def get_general_metrics_prompt(
+    prompt: GeneralMetricsPrompt,
+    conversation_context: Union[str, List[Dict[str, str]]],
+    tool_inventory: List[Dict[str, Any]],
+    tool_call: Dict[str, Any],
+) -> List[Dict[str, str]]:
+    """
+    Build the messages for a general semantic evaluation.
+
+    Returns the list of chat messages (system -> [few-shot] -> user).
+    """
+    return prompt.build_messages(
+        user_kwargs={
+            "conversation_context": conversation_context,
+            "tool_inventory": tool_inventory,
+            "tool_call": tool_call,
+        }
+    )