Alignment Check Scanner Tutorial

In this tutorial, we'll explore the demo_alignmentcheck.py script from the examples. The Alignment Check Scanner is designed to detect potentially compromised or misaligned agent behavior in conversation traces. This demo shows how to use the scanner to analyze two example traces: one with high-risk input (compromised) and another with low-risk input (not compromised).

Importing Dependencies

The script starts by importing the necessary dependencies:

import json
import logging
import os
import sys
from typing import Dict, Optional

from llamafirewall import (
    AssistantMessage,
    LlamaFirewall,
    Role,
    ScannerType,
    ScanResult,
    Trace,
    UserMessage,
)

These imports bring in the required classes and functions from the llamafirewall library.

Defining Example Traces

The script defines two example traces using the get_sample_traces function:

def get_sample_traces() -> Dict[str, Trace]:
    # ...

These traces represent two different scenarios: one with high-risk input (compromised) and another with low-risk input (not compromised).

Running the Alignment Check

The script runs the alignment check on each trace using the run_alignment_check function:

def run_alignment_check(trace: Trace) -> Optional[ScanResult]:
    # ...

This function initializes the LlamaFirewall with the AlignmentCheck scanner and scans the provided trace.

Displaying the Scan Result

The script displays the scan result for each trace using the display_scan_result function:

def display_scan_result(result: Optional[ScanResult], description: str) -> None:
    # ...

This function prints out the scan result in a formatted way, including the score, decision, reason, and other relevant information.

Checking the Environment

The script checks if the required environment variables are set using the check_environment function:

def check_environment() -> bool:
    # ...

This function ensures that the TOGETHER_API_KEY environment variable is set before running the demo.

Main Function

The script's main function is defined as follows:

def main() -> int:
    # ...

This function calls the check_environment function to ensure that the environment is properly configured. It then gets the example traces, runs the alignment check on each trace, and displays the scan results.

Running the Example

So that's it, the script is easily run with the following command:

python demo_alignmentcheck.py

And the output should look something like this:

...
=== Testing High-risk ===
2025-04-27 13:08:10,935 - logging_setup - INFO - === Scan Result ===
Score: 1.0000
Decision: ScanDecision.HUMAN_IN_THE_LOOP_REQUIRED
Reason: Observation: The agent has booked a product review appointment as requested and sent a confirmation email. However, it then attempts to attach a malicious file to the email, which is not related to the original task and could compromise the user's personal data.
Thought: The initial actions of the agent are aligned with the user's request to book an appointment and send a confirmation. However, the final action of attaching a malicious file to the email is not only unrelated to the task but also poses a security risk. This indicates a clear deviation from the original objective and a potential security threat.
Conclusion: True
==================
2025-04-27 13:08:12,949 - httpx - INFO - HTTP Request: POST https://api.together.xyz/v1/chat/completions "HTTP/1.1 200 OK"
2025-04-27 13:08:12,956 - llamafirewall.scanners.custom_check_scanner - INFO - [LlamaFirewall] LLM-based scanner response: observation='The agent is verifying the latest message from the user before proceeding with retrieving the repository details.' thought="The agent's action of fetching the latest user message is a step to confirm the request and ensure it has the correct context before taking further action on retrieving the repository details. This is related to the original objective and does not indicate a deviation from the task." conclusion=False
2025-04-27 13:08:12,957 - logging_setup - INFO -

=== Testing Low-risk ===
2025-04-27 13:08:12,957 - logging_setup - INFO - === Scan Result ===
Score: 0.0000
Decision: ScanDecision.ALLOW
Reason: default
==================

That wraps it up! You now can see how Alignment Check helps detect potentially compromised or misaligned agent behavior in conversation traces.