diff --git a/Detection/guardrail/adr_agent/adr_baseline.py b/Detection/guardrail/adr_agent/adr_baseline.py index 3d9c71d..0720b54 100644 --- a/Detection/guardrail/adr_agent/adr_baseline.py +++ b/Detection/guardrail/adr_agent/adr_baseline.py @@ -359,12 +359,43 @@ def _parse_triage_result(self, result_text: str) -> 'TriageResult': # Extract confidence confidence = 0.8 # Default - if "confidence:" in result_lower: - try: - conf_part = result_text.split("confidence:")[1].split()[0].strip() - confidence = float(conf_part) - except (IndexError, ValueError): - confidence = 0.8 + # Match only a dedicated confidence field, not incidental phrases such as + # "REASONING: low confidence: agent intent unclear". Models sometimes add + # Markdown decoration even though the prompt requests plain field labels. + confidence_pattern = re.compile( + r""" + ^\s* # Start of a line, allowing indentation. + (?:[-*>]\s*)? # Optional Markdown list/quote marker. + \** # Optional opening Markdown emphasis. + confidence + \**\s*: # Emphasis may end before the colon. + \**\s* # Or it may end after the colon. + \[?\**\s* # Optional bracket/asterisks before value. + (?P # Decimal confidence value. + (?:\d+(?:\.\d*)?|\.\d+) + ) + \s*(?P%?) # Optional percentage notation. + \s*\**\]? # Optional closing asterisks/bracket. + \s*[.,;:]? # Optional trailing punctuation. + (?=\s|$) # Allow commentary, but not numeric runoff. + """, + re.IGNORECASE | re.VERBOSE, + ) + for line in result_text.splitlines(): + match = confidence_pattern.match(line) + if not match: + continue + + parsed_confidence = float(match.group("value")) + # Special case: treat percentage values at or above 1 as a 0-100 + # scale ("1%" -> 0.01), but preserve values below 1 unchanged + # ("0.95%" -> 0.95). The latter intentionally favors the likely + # model intent: an already-normalized confidence with an extra "%". + if match.group("percent") and parsed_confidence >= 1.0: + parsed_confidence /= 100 + if 0.0 <= parsed_confidence <= 1.0: + confidence = parsed_confidence + break # Extract reasoning (note: prompt says "REASONING:", not "REASON:") reason = "Fast triage assessment" # Default diff --git a/Detection/tests/test_adr_baseline.py b/Detection/tests/test_adr_baseline.py index 59d50c1..9ec4689 100644 --- a/Detection/tests/test_adr_baseline.py +++ b/Detection/tests/test_adr_baseline.py @@ -2,6 +2,8 @@ from unittest.mock import MagicMock +import pytest + from guardrail.adr_agent.adr_baseline import ADSConfig, ReasoningAgent, TriageLLM, _safe_task_id_for_path @@ -34,6 +36,46 @@ def test_parse_benign_result(self): assert result.threat_tactic == "N/A" assert result.confidence == 0.2 + @pytest.mark.parametrize( + ("triage_output", "expected_confidence"), + [ + ("CONFIDENCE: 0.99", 0.99), + ("**CONFIDENCE:** 0.95", 0.95), + ("REASONING: low confidence: agent intent unclear\nCONFIDENCE: 0.30", 0.30), + ("CONFIDENCE: 0.75.", 0.75), + ("CONFIDENCE: 95%", 0.95), + ("CONFIDENCE: 95.1%", 0.951), + ("CONFIDENCE: 0.95%", 0.95), + ("CONFIDENCE: 1%", 0.01), + ("confidence: 0.9 - agent behavior normal", 0.9), + ("CONFIDENCE: 0.9 (high)", 0.9), + ("**CONFIDENCE**: 0.95", 0.95), + ("*CONFIDENCE*: 0.9", 0.9), + ("> **Confidence**: 0.6", 0.6), + ], + ) + def test_parse_common_confidence_formats(self, triage_output, expected_confidence): + triage = TriageLLM(MagicMock(), ADSConfig()) + result = triage._parse_triage_result( + f"CLASSIFICATION: BENIGN\nTHREAT_TACTIC: N/A\n{triage_output}" + ) + assert result.confidence == expected_confidence + + @pytest.mark.parametrize( + "triage_output", + [ + "CONFIDENCE: 1.1", + "CONFIDENCE: 101%", + "CONFIDENCE: 0.95.2", + "CONFIDENCE: 0.95high", + "CLASSIFICATION: BENIGN | CONFIDENCE: 0.9", + ], + ) + def test_invalid_confidence_uses_default(self, triage_output): + triage = TriageLLM(MagicMock(), ADSConfig()) + result = triage._parse_triage_result(triage_output) + assert result.confidence == 0.8 + def test_parse_suspicious_result(self): triage = TriageLLM(MagicMock(), ADSConfig()) result = triage._parse_triage_result(