Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions Detection/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -280,6 +280,34 @@ uv run python main_detector.py --tasks 1-10 --results-dir "$BENCH"

> **Note:** Default detector is `adr` (ADR dual-agent) and requires API keys + Claude CLI. For keyless smoke tests, use `--detector llamafirewall`.

#### Optional structured Tier-1 routing

ADR uses its existing text triage contract by default. To opt into the strict
structured routing contract, add this field under `adr_framework.triage_llm`
in `config_detector.yaml`:

```yaml
decision_contract: "structured_risk_route_v1"
```

The structured contract is supported for ADR-Bench and AgentDojo. It reports
an explicit `BENIGN_EXIT` or `ESCALATE` route together with risk, injection,
unsafe-plan/action/effect, authorization, evidence, and provenance-needed
fields. Validated fields are persisted in the optional
`structured_triage` result key and a `<task_id>_structured_triage.json` debug
artifact, including for escalated tasks. Free-form fields in this audit object
are untrusted model-generated data, not instructions. Stock results omit the key.
Missing, malformed, or internally inconsistent structured output
fails toward `ESCALATE` without a validated audit object. On escalation, the
handoff may include bounded,
validated MCP component names from framework system metadata so existing
source review targets the declared components; free-form Tier-1 text is not
copied into that trusted prompt slot. The ADR-Bench profile preserves the
stock high-recall escalation conditions: structured fields annotate those
routes rather than narrowing them, and Tier 2 evaluates scenario/component
threat separately from whether an unsafe action or effect occurred. Omitting
the option preserves the stock prompt, request, parser, and Tier-2 handoff.

### Output

Detection results saved in the benchmark directory:
Expand Down
3 changes: 3 additions & 0 deletions Detection/config_detector.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,9 @@ adr_framework:
# Triage LLM Configuration (High Recall)
triage_llm:
model: "gpt-4o"
# Optional: "structured_risk_route_v1" enables schema-constrained routing.
# Omit this setting (or use "stock_text") to preserve the stock text contract.
# decision_contract: "structured_risk_route_v1"
cost_per_1m_input: 2.50 # USD - update if model changes
cost_per_1m_output: 10.00
temperature: 0
Expand Down
60 changes: 57 additions & 3 deletions Detection/guardrail/adr_agent/adr_baseline.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import json
import logging
import subprocess
from copy import deepcopy
from pathlib import Path
from typing import Dict, List, Any, Optional
import sys
Expand Down Expand Up @@ -222,6 +223,18 @@ def __init__(self, config_data=None):
def get_triage_model(self) -> str:
return self.triage_config.get('model', 'gpt-4o')

def get_triage_decision_contract(self) -> str:
"""Return the selected Tier-1 output contract.

Omitting the setting preserves the original text prompt and parser.
Contract names are validated when the triage implementation is built.
"""

value = self.triage_config.get('decision_contract', 'stock_text')
if not isinstance(value, str) or not value:
raise ValueError("ADR triage decision_contract must be a nonempty string")
return value

def get_triage_rates(self) -> tuple[float, float]:
return (self.triage_config.get('cost_per_1m_input', 2.50),
self.triage_config.get('cost_per_1m_output', 10.00))
Expand Down Expand Up @@ -249,7 +262,9 @@ def __init__(self, config_data=None, benchmark_type: str = "adr_bench", **kwargs

# Initialize dual-agent system using API keys
self.openai_client = get_openai_client()
self.triage_llm = TriageLLM(self.openai_client, self.config, benchmark_type=benchmark_type)
self.triage_llm = _build_triage_llm(
self.openai_client, self.config, benchmark_type=benchmark_type
)
self.reasoning_agent = ReasoningAgent(self.config, benchmark_type=benchmark_type)

logger.info(f"ADR initialized ({benchmark_type}): {self.config.get_triage_model()} + {self.config.get_reasoning_model()}")
Expand Down Expand Up @@ -333,13 +348,25 @@ def _analyze_messages(self, messages: List[Dict[str, Any]], task_id: str) -> Det
# downstream can fail or redact it further.
logger.warning(f"🔍 {deterministic_result.reason}")

structured_triage = None
# Check if triage is enabled
if self.config.enable_triage:
# Stage 1: Triage LLM (first line of defense per proposal),
# unless the deterministic check already found something - no
# need to pay for an LLM call when we already have a
# definitive signal.
triage_result = deterministic_result or self.triage_llm.analyze(messages)
audit = getattr(triage_result, 'structured_output', None)
if isinstance(audit, dict):
structured_triage = deepcopy(audit)
# Save before Tier 2 so its failure cannot erase the routing audit.
# Free-form evidence stays in artifacts, never in the trusted handoff.
audit_file = self.reasoning_agent.debug_log_dir / f"{_safe_task_id_for_path(task_id)}_structured_triage.json"
try:
with open(audit_file, 'w') as f:
json.dump({'task_id': task_id, 'structured_triage': structured_triage}, f, indent=2)
except OSError:
logger.warning("Could not write structured triage audit artifact")

# Fast path for clearly benign (saves Claude resources)
if not triage_result.is_suspicious:
Expand All @@ -362,7 +389,8 @@ def _analyze_messages(self, messages: List[Dict[str, Any]], task_id: str) -> Det
'analysis_time': analysis_time,
'input_tokens': triage_result.input_tokens,
'output_tokens': triage_result.output_tokens,
'cost_usd': triage_cost
'cost_usd': triage_cost,
**({'structured_triage': structured_triage} if structured_triage is not None else {}),
}, f, indent=2)
logger.info(f"📝 Triage-only log saved: {triage_log_file}")

Expand All @@ -383,7 +411,8 @@ def _analyze_messages(self, messages: List[Dict[str, Any]], task_id: str) -> Det
analysis_time=analysis_time,
input_tokens=triage_result.input_tokens,
output_tokens=triage_result.output_tokens,
cost_usd=triage_cost
cost_usd=triage_cost,
structured_triage=structured_triage,
)

triage_reasoning = f"Triage escalation: {triage_result.prompt_reason or triage_result.reason}"
Expand Down Expand Up @@ -420,6 +449,8 @@ def _analyze_messages(self, messages: List[Dict[str, Any]], task_id: str) -> Det
reasoning_result.output_tokens = triage_tokens_out + (reasoning_result.output_tokens or 0)
reasoning_result.cost_usd = total_cost

if structured_triage is not None:
reasoning_result.structured_triage = structured_triage
return reasoning_result


Expand Down Expand Up @@ -663,6 +694,29 @@ def _format_conversation(self, messages: List[Dict[str, Any]]) -> str:
return "\n".join(formatted)


def _build_triage_llm(
openai_client: Any,
config: ADSConfig,
benchmark_type: str,
) -> TriageLLM:
"""Build the configured Tier-1 implementation.

The structured implementation is imported lazily so the stock path keeps
the same class, prompt, parser, and OpenAI request body as before.
"""

decision_contract = config.get_triage_decision_contract()
if decision_contract == "stock_text":
return TriageLLM(openai_client, config, benchmark_type=benchmark_type)
if decision_contract == "structured_risk_route_v1":
from .structured_risk_route_triage import StructuredRiskRouteTriageLLM

return StructuredRiskRouteTriageLLM(
openai_client, config, benchmark_type=benchmark_type
)
raise ValueError(f"Unknown ADR triage decision contract: {decision_contract!r}")


class ReasoningAgent:
"""High-precision reasoning with persistent workspace and enterprise MCP intelligence"""

Expand Down
Loading
Loading