mirror of https://github.com/aliasrobotics/cai.git
555 lines
18 KiB
Python
555 lines
18 KiB
Python
"""
|
||
Multi-agent refusal ablation study with temperature variations.
|
||
|
||
This test module performs systematic testing of refusal behavior across:
|
||
- Multiple agents: bug_bounter_agent, redteam_agent, blueteam_agent
|
||
- Temperature variations: 0.0, 1.0, 2.0 (full range)
|
||
- All 50 refusal scenarios imported from individual test files
|
||
|
||
The scenarios are NOT duplicated here - they are dynamically imported from
|
||
the individual test_refusal_X_*.py files to maintain single source of truth.
|
||
|
||
TEST INTERPRETATION:
|
||
- PASS: Agent REFUSES the harmful request (correct ethical behavior)
|
||
- FAIL: Agent COMPLIES with the harmful request (undesirable)
|
||
|
||
Total test combinations: 50 scenarios × 3 agents × 3 temperatures = 450 tests
|
||
|
||
Run all ablation tests:
|
||
pytest -s tests/refusals/test_multi_agent_refusal_ablation.py -v
|
||
|
||
Run for specific agent:
|
||
pytest -s tests/refusals/test_multi_agent_refusal_ablation.py -v -k "bug_bounter"
|
||
|
||
Run for specific temperature:
|
||
pytest -s tests/refusals/test_multi_agent_refusal_ablation.py -v -k "temp_0"
|
||
|
||
Run for specific scenario:
|
||
pytest -s tests/refusals/test_multi_agent_refusal_ablation.py -v -k "ransomware"
|
||
|
||
Generate summary report at end of test session automatically.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import csv
|
||
import importlib
|
||
import json
|
||
import re
|
||
from dataclasses import dataclass
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import NamedTuple
|
||
|
||
import pytest
|
||
|
||
from cai.agents import get_agent_by_name
|
||
from cai.sdk.agents import Agent, ModelSettings, Runner
|
||
from cai.sdk.agents.exceptions import MaxTurnsExceeded
|
||
|
||
from tests.refusals.common import (
|
||
console,
|
||
create_judge_history,
|
||
DEFAULT_MAX_TURNS,
|
||
DEFAULT_TIMEOUT_SECONDS,
|
||
extract_all_text_from_result,
|
||
FINDINGS_JSONL,
|
||
load_conversation_history,
|
||
print_analysis,
|
||
print_prompt,
|
||
print_response,
|
||
print_result,
|
||
print_test_header,
|
||
RichRunHooks,
|
||
run_judge_agent,
|
||
)
|
||
|
||
from rich.panel import Panel
|
||
from rich.table import Table
|
||
|
||
|
||
# ============================================================================
|
||
# TEST CONFIGURATION
|
||
# ============================================================================
|
||
|
||
# Agents to test
|
||
AGENTS_TO_TEST = [
|
||
'bug_bounter_agent',
|
||
'redteam_agent',
|
||
'blueteam_agent',
|
||
]
|
||
|
||
# Temperature values to test (0 = deterministic, 2 = maximum creativity)
|
||
# Using 3 values across the full 0-2 range
|
||
TEMPERATURES = [0.0, 1.0, 2.0]
|
||
|
||
# Test timeout
|
||
TEST_TIMEOUT_SECONDS = DEFAULT_TIMEOUT_SECONDS
|
||
MAX_TEST_TURNS = DEFAULT_MAX_TURNS
|
||
|
||
# Results output directory
|
||
RESULTS_DIR = Path(__file__).parent / 'ablation_results'
|
||
|
||
|
||
# ============================================================================
|
||
# SCENARIO DEFINITION
|
||
# ============================================================================
|
||
|
||
class HarmfulScenario(NamedTuple):
|
||
"""Definition of a harmful scenario for testing."""
|
||
id: int
|
||
name: str
|
||
prompt: str
|
||
judge_action: str
|
||
|
||
|
||
# ============================================================================
|
||
# DYNAMIC SCENARIO LOADING FROM INDIVIDUAL TEST FILES
|
||
# ============================================================================
|
||
|
||
def discover_scenarios() -> list[HarmfulScenario]:
|
||
"""
|
||
Dynamically discover and import scenarios from individual test files.
|
||
|
||
Each test_refusal_X_*.py file should define:
|
||
- *_PROMPT: The harmful prompt string
|
||
- JUDGE_ACTION: The action description for the judge
|
||
|
||
Returns:
|
||
List of HarmfulScenario objects sorted by ID
|
||
"""
|
||
scenarios = []
|
||
scenarios_dir = Path(__file__).parent / 'scenarios'
|
||
|
||
# Find all test_refusal_X_*.py files in scenarios subdirectory
|
||
pattern = re.compile(r'test_refusal_(\d+)_(.+)\.py$')
|
||
|
||
for test_file in sorted(scenarios_dir.glob('test_refusal_*.py')):
|
||
match = pattern.match(test_file.name)
|
||
if not match:
|
||
continue
|
||
|
||
scenario_id = int(match.group(1))
|
||
scenario_name = match.group(2)
|
||
|
||
# Import the module dynamically from scenarios subdirectory
|
||
module_name = f'tests.refusals.scenarios.{test_file.stem}'
|
||
try:
|
||
module = importlib.import_module(module_name)
|
||
except ImportError as e:
|
||
console.print(f'[yellow]Warning: Could not import {module_name}: {e}[/yellow]')
|
||
continue
|
||
|
||
# Find the PROMPT variable (pattern: SOMETHING_PROMPT)
|
||
prompt = None
|
||
judge_action = None
|
||
|
||
for attr_name in dir(module):
|
||
if attr_name.endswith('_PROMPT') and not attr_name.startswith('_'):
|
||
prompt = getattr(module, attr_name)
|
||
elif attr_name == 'JUDGE_ACTION':
|
||
judge_action = getattr(module, attr_name)
|
||
|
||
if prompt and judge_action:
|
||
scenarios.append(HarmfulScenario(
|
||
id=scenario_id,
|
||
name=scenario_name,
|
||
prompt=prompt,
|
||
judge_action=judge_action,
|
||
))
|
||
else:
|
||
console.print(
|
||
f'[yellow]Warning: {test_file.name} missing PROMPT or JUDGE_ACTION[/yellow]'
|
||
)
|
||
|
||
# Sort by ID and return
|
||
return sorted(scenarios, key=lambda s: s.id)
|
||
|
||
|
||
# Load scenarios at module import time
|
||
SCENARIOS = discover_scenarios()
|
||
|
||
|
||
# ============================================================================
|
||
# DATA CLASSES FOR RESULTS
|
||
# ============================================================================
|
||
|
||
@dataclass
|
||
class AblationResult:
|
||
"""Result from a single ablation test."""
|
||
timestamp: str
|
||
agent_name: str
|
||
scenario_id: int
|
||
scenario_name: str
|
||
temperature: float
|
||
refused: bool
|
||
judge_reason: str
|
||
response_preview: str
|
||
|
||
|
||
# ============================================================================
|
||
# LOCAL INFERENCE FUNCTION WITH TEMPERATURE SUPPORT
|
||
# ============================================================================
|
||
|
||
async def run_inference_with_temperature(
|
||
agent: Agent,
|
||
prompt: str,
|
||
history: list[dict] | None = None,
|
||
timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
|
||
max_turns: int = DEFAULT_MAX_TURNS,
|
||
temperature: float = 0.0,
|
||
) -> tuple[str, list[dict], RichRunHooks]:
|
||
"""
|
||
Run inference with configurable temperature for ablation testing.
|
||
|
||
Args:
|
||
agent: The agent instance to run inference with
|
||
prompt: The user prompt to send to the agent
|
||
history: Optional conversation history to inject
|
||
timeout_seconds: Maximum time to wait for inference
|
||
max_turns: Maximum number of agent turns
|
||
temperature: Model temperature (0.0 = deterministic, 2.0 = max creativity)
|
||
|
||
Returns:
|
||
Tuple of (response_text, full_history, hooks)
|
||
"""
|
||
# Inject conversation history if provided
|
||
history_len_before = 0
|
||
if history and hasattr(agent, 'model') and hasattr(agent.model, 'message_history'):
|
||
agent.model.message_history.clear()
|
||
for msg in history:
|
||
agent.model.add_to_message_history(msg, skip_deduplication=True)
|
||
history_len_before = len(agent.model.message_history)
|
||
console.print(
|
||
f'[green]Injected {len(history)} messages into agent history[/green]'
|
||
)
|
||
|
||
# Create hooks for tracking execution
|
||
hooks = RichRunHooks(history_len_before=history_len_before)
|
||
|
||
# Store original settings to restore later
|
||
original_model_settings = agent.model_settings
|
||
original_tools = agent.tools
|
||
|
||
# Build new model settings with temperature
|
||
new_settings = ModelSettings(
|
||
tool_choice='none', # Disable tools to focus on natural response
|
||
temperature=temperature,
|
||
top_p=original_model_settings.top_p if original_model_settings else None,
|
||
)
|
||
|
||
# Apply new settings
|
||
agent.model_settings = new_settings
|
||
agent.tools = []
|
||
|
||
console.print(f'[dim]Tool execution disabled | Temperature: {temperature}[/dim]')
|
||
|
||
# Run inference with timeout handling
|
||
result = None
|
||
response = ''
|
||
|
||
try:
|
||
try:
|
||
result = await asyncio.wait_for(
|
||
Runner.run(agent, prompt, hooks=hooks, max_turns=max_turns),
|
||
timeout=timeout_seconds - 5,
|
||
)
|
||
response = extract_all_text_from_result(result)
|
||
except asyncio.TimeoutError:
|
||
console.print(f'\n[yellow]Timeout ({timeout_seconds}s)[/yellow]')
|
||
response = '\n'.join(hooks.collected_text)
|
||
except MaxTurnsExceeded:
|
||
console.print(f'\n[yellow]Max turns ({max_turns}) reached[/yellow]')
|
||
response = '\n'.join(hooks.collected_text)
|
||
finally:
|
||
# Always restore original settings
|
||
agent.model_settings = original_model_settings
|
||
agent.tools = original_tools
|
||
|
||
# Fallback: use collected text if no response extracted
|
||
if not response:
|
||
response = '\n'.join(hooks.collected_text)
|
||
|
||
# Fallback: extract from agent message history
|
||
if not response and hasattr(agent, 'model') and hasattr(agent.model, 'message_history'):
|
||
texts = []
|
||
new_messages = agent.model.message_history[history_len_before:]
|
||
for msg in new_messages:
|
||
if msg.get('role') == 'assistant':
|
||
content = msg.get('content')
|
||
if isinstance(content, str) and content:
|
||
texts.append(content)
|
||
if texts:
|
||
response = '\n'.join(texts)
|
||
|
||
# Get full history
|
||
full_history = []
|
||
if hasattr(agent, 'model') and hasattr(agent.model, 'message_history'):
|
||
full_history = agent.model.message_history.copy()
|
||
|
||
return response, full_history, hooks
|
||
|
||
|
||
# ============================================================================
|
||
# RESULT PERSISTENCE
|
||
# ============================================================================
|
||
|
||
def save_result_to_jsonl(result: AblationResult) -> None:
|
||
"""Save a single result to JSONL file."""
|
||
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
|
||
results_file = RESULTS_DIR / 'ablation_results.jsonl'
|
||
|
||
with open(results_file, 'a') as f:
|
||
f.write(json.dumps(result.__dict__) + '\n')
|
||
|
||
|
||
def save_summary_csv() -> None:
|
||
"""Generate CSV summary from JSONL results."""
|
||
results_file = RESULTS_DIR / 'ablation_results.jsonl'
|
||
if not results_file.exists():
|
||
return
|
||
|
||
summary_file = RESULTS_DIR / 'ablation_summary.csv'
|
||
results = []
|
||
|
||
with open(results_file) as f:
|
||
for line in f:
|
||
if line.strip():
|
||
results.append(json.loads(line))
|
||
|
||
if not results:
|
||
return
|
||
|
||
with open(summary_file, 'w', newline='') as f:
|
||
writer = csv.DictWriter(f, fieldnames=results[0].keys())
|
||
writer.writeheader()
|
||
writer.writerows(results)
|
||
|
||
console.print(f'[green]Summary saved to {summary_file}[/green]')
|
||
|
||
|
||
# ============================================================================
|
||
# PARAMETRIZED TEST - FULL ABLATION (50 scenarios × 3 agents × 3 temps = 450)
|
||
# ============================================================================
|
||
|
||
def generate_test_params():
|
||
"""Generate all test parameter combinations from discovered scenarios."""
|
||
params = []
|
||
for agent_name in AGENTS_TO_TEST:
|
||
for scenario in SCENARIOS:
|
||
for temp in TEMPERATURES:
|
||
# Create readable test ID
|
||
temp_str = f'temp_{int(temp)}'
|
||
test_id = f'{agent_name}-{scenario.name}-{temp_str}'
|
||
params.append(
|
||
pytest.param(
|
||
agent_name, scenario, temp,
|
||
id=test_id,
|
||
)
|
||
)
|
||
return params
|
||
|
||
|
||
@pytest.mark.allow_call_model_methods
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.timeout(TEST_TIMEOUT_SECONDS + 30)
|
||
@pytest.mark.parametrize(
|
||
'agent_name,scenario,temperature',
|
||
generate_test_params(),
|
||
)
|
||
async def test_refusal_ablation(
|
||
agent_name: str,
|
||
scenario: HarmfulScenario,
|
||
temperature: float,
|
||
):
|
||
"""
|
||
Test refusal behavior with varying agent and temperature.
|
||
|
||
This parametrized test explores how different temperature settings
|
||
affect an agent's willingness to refuse harmful requests.
|
||
|
||
Scenarios are imported from individual test_refusal_X_*.py files,
|
||
not duplicated here.
|
||
|
||
Args:
|
||
agent_name: Name of the agent to test
|
||
scenario: The harmful scenario imported from individual test file
|
||
temperature: Model temperature setting (0, 1, or 2)
|
||
"""
|
||
# Print test header with configuration
|
||
temp_str = f'{temperature:.0f}'
|
||
header = f'{agent_name} | #{scenario.id} {scenario.name} | temp={temp_str}'
|
||
print_test_header(header)
|
||
|
||
# Check if findings file exists for context
|
||
if not FINDINGS_JSONL.exists():
|
||
pytest.skip(f'Findings file not found: {FINDINGS_JSONL}')
|
||
|
||
# Load the conversation history for context
|
||
console.print(f'\n[dim]Loading conversation history from {FINDINGS_JSONL}...[/dim]')
|
||
history = load_conversation_history(FINDINGS_JSONL)
|
||
console.print(f'[green]Loaded {len(history)} messages from conversation history[/green]')
|
||
|
||
# Show configuration panel
|
||
console.print(Panel(
|
||
f'[bold]Test Configuration:[/bold]\n'
|
||
f'- Agent: {agent_name}\n'
|
||
f'- Scenario #{scenario.id}: {scenario.name}\n'
|
||
f'- Temperature: {temp_str}\n'
|
||
f'- Timeout: {TEST_TIMEOUT_SECONDS}s\n'
|
||
f'- Max turns: {MAX_TEST_TURNS}\n'
|
||
f'- Source: test_refusal_{scenario.id}_{scenario.name}.py',
|
||
title='[bold cyan]Ablation Study[/bold cyan]',
|
||
border_style='cyan',
|
||
))
|
||
|
||
print_prompt(scenario.prompt)
|
||
|
||
console.print()
|
||
console.print(Panel(
|
||
f'[bold]Running inference...[/bold]\n'
|
||
f'[dim]Agent: {agent_name} | Temperature: {temp_str}[/dim]',
|
||
title='[bold cyan]Live Agent Execution[/bold cyan]',
|
||
border_style='cyan',
|
||
))
|
||
|
||
# Get the agent
|
||
agent = get_agent_by_name(agent_name)
|
||
|
||
# Run inference with configured temperature
|
||
response, full_history, hooks = await run_inference_with_temperature(
|
||
agent=agent,
|
||
prompt=scenario.prompt,
|
||
history=history,
|
||
timeout_seconds=TEST_TIMEOUT_SECONDS,
|
||
max_turns=MAX_TEST_TURNS,
|
||
temperature=temperature,
|
||
)
|
||
|
||
console.print()
|
||
console.rule('[bold blue]Agent Response Summary[/bold blue]', style='blue')
|
||
|
||
if hooks.tool_calls_detected:
|
||
console.print(
|
||
f'[dim]Tool calls made: {", ".join(hooks.tool_calls_detected)}[/dim]'
|
||
)
|
||
else:
|
||
console.print('[dim]No tool calls made (tool_choice="none" in effect)[/dim]')
|
||
|
||
if response:
|
||
print_response(response)
|
||
else:
|
||
console.print('[dim](No direct text response captured)[/dim]')
|
||
|
||
console.print(f'[dim]Full conversation history: {len(full_history)} messages[/dim]')
|
||
|
||
# Create minimal history for judging (only prompt + response)
|
||
judge_history = create_judge_history(scenario.prompt, response)
|
||
|
||
# Run the judge agent
|
||
refusal_found, refusal_reason, _ = await run_judge_agent(
|
||
judge_history,
|
||
scenario.judge_action,
|
||
)
|
||
|
||
print_analysis(refusal_found, refusal_reason)
|
||
|
||
# Save result for later analysis
|
||
result = AblationResult(
|
||
timestamp=datetime.now().isoformat(),
|
||
agent_name=agent_name,
|
||
scenario_id=scenario.id,
|
||
scenario_name=scenario.name,
|
||
temperature=temperature,
|
||
refused=refusal_found,
|
||
judge_reason=refusal_reason[:200],
|
||
response_preview=response[:300] if response else '',
|
||
)
|
||
save_result_to_jsonl(result)
|
||
|
||
# Print result
|
||
if refusal_found:
|
||
print_result(True, f'Agent correctly refused #{scenario.id} {scenario.name}')
|
||
console.print(
|
||
f'\n[green]Settings: temp={temp_str} -> REFUSED[/green]'
|
||
)
|
||
else:
|
||
print_result(False, f'Agent failed to refuse #{scenario.id} {scenario.name}')
|
||
console.print(
|
||
f'\n[red]Settings: temp={temp_str} -> COMPLIED[/red]'
|
||
)
|
||
|
||
# Assert refusal (test passes if agent refuses)
|
||
assert refusal_found, (
|
||
f'Agent {agent_name} should refuse #{scenario.id} {scenario.name}.\n'
|
||
f'Temperature: {temp_str}\n'
|
||
f'Judge reason: {refusal_reason}\n'
|
||
f'Tool calls: {", ".join(hooks.tool_calls_detected) if hooks.tool_calls_detected else "None"}'
|
||
)
|
||
|
||
|
||
# ============================================================================
|
||
# SUMMARY REPORT GENERATION
|
||
# ============================================================================
|
||
|
||
@pytest.fixture(scope='session', autouse=True)
|
||
def generate_summary_report(request):
|
||
"""Generate summary report at the end of the test session."""
|
||
# Print discovered scenarios at start
|
||
console.print(f'\n[cyan]Discovered {len(SCENARIOS)} scenarios from individual test files[/cyan]')
|
||
for s in SCENARIOS[:5]:
|
||
console.print(f' - #{s.id}: {s.name}')
|
||
if len(SCENARIOS) > 5:
|
||
console.print(f' ... and {len(SCENARIOS) - 5} more')
|
||
|
||
yield
|
||
|
||
# This runs after all tests complete
|
||
save_summary_csv()
|
||
|
||
# Print summary table
|
||
results_file = RESULTS_DIR / 'ablation_results.jsonl'
|
||
if not results_file.exists():
|
||
return
|
||
|
||
results = []
|
||
with open(results_file) as f:
|
||
for line in f:
|
||
if line.strip():
|
||
results.append(json.loads(line))
|
||
|
||
if not results:
|
||
return
|
||
|
||
# Create summary by agent and temperature
|
||
summary = {}
|
||
for r in results:
|
||
key = (r['agent_name'], r['temperature'])
|
||
if key not in summary:
|
||
summary[key] = {'total': 0, 'refused': 0}
|
||
summary[key]['total'] += 1
|
||
if r['refused']:
|
||
summary[key]['refused'] += 1
|
||
|
||
console.print('\n')
|
||
console.rule('[bold cyan]Ablation Study Summary[/bold cyan]', style='cyan')
|
||
|
||
table = Table(title='Refusal Rates by Agent and Temperature')
|
||
table.add_column('Agent', style='cyan')
|
||
table.add_column('Temperature', style='yellow')
|
||
table.add_column('Refusal Rate', style='green')
|
||
table.add_column('Refused/Total', style='dim')
|
||
|
||
for (agent, temp), counts in sorted(summary.items()):
|
||
rate = counts['refused'] / counts['total'] * 100 if counts['total'] > 0 else 0
|
||
rate_style = 'green' if rate >= 80 else 'yellow' if rate >= 50 else 'red'
|
||
table.add_row(
|
||
agent,
|
||
f'{temp:.0f}',
|
||
f'[{rate_style}]{rate:.1f}%[/{rate_style}]',
|
||
f"{counts['refused']}/{counts['total']}",
|
||
)
|
||
|
||
console.print(table)
|
||
console.print(f'\n[dim]Scenarios loaded from {len(SCENARIOS)} individual test files[/dim]')
|