01 · 数据模型与策略引擎

本章讲主线:一条消息进来,LlamaFirewall 怎么根据它的角色挑扫描器、怎么把多个扫描结果汇总成一个裁决,以及扫单条 vs. 扫整段对话的几个方法差在哪。

1. 先认识数据模型(全在一个文件里)

所有基础类型定义在 src/llamafirewall/llamafirewall_data_types.py,很短,值得先看。

消息与角色

Message 是一个 dataclass:role + content + 可选 tool_calls(llamafirewall_data_types.py:50)。为了好用,库给每个角色包了一个子类构造器:

子类	角色	语义
`UserMessage`	`Role.USER`	用户输入
`SystemMessage`	`Role.SYSTEM`	系统提示
`AssistantMessage`	`Role.ASSISTANT`	LLM 输出(可带 `tool_calls`)
`ToolMessage`	`Role.TOOL`	工具 / 网页返回的结果
`MemoryMessage`	`Role.MEMORY`	从记忆里读出的内容

Trace 只是 List[Message] 的别名(llamafirewall_data_types.py:95)——也就是「一段对话历史」。

裁决结果

ScanResult(llamafirewall_data_types.py:42)是所有扫描器的统一输出:

@dataclass
class ScanResult:
    decision: ScanDecision   # ALLOW / BLOCK / HUMAN_IN_THE_LOOP_REQUIRED
    reason: str              # 人读的解释
    score: float             # 0.0~1.0 可信度/风险分
    status: ScanStatus = ScanStatus.SUCCESS  # SUCCESS / ERROR / SKIPPED

三档裁决是这个框架的一个关键设计:除了 ALLOW/BLOCK,还有 HUMAN_IN_THE_LOOP_REQUIRED——「拿不准,交给人」。目前只有 AlignmentCheck 用它(见第 2 章)。

2. 策略 = 一张 `Role → [Scanner]` 的表

「用哪些扫描器查哪种消息」这件事,被抽象成一个映射类型(config.py:13):

Configuration: TypeAlias = Mapping[Role, Sequence[ScannerType | str]]

注意值可以是 ScannerType(内置)或 str(自定义扫描器的注册名)——这为扩展留了口子(第 4 章)。

构造 LlamaFirewall 时不传表,就用默认表(llamafirewall.py:90):

角色	默认扫描器	为什么
`TOOL`	`CODE_SHIELD`, `PROMPT_GUARD`	工具/网页结果不可信,既可能夹注入又可能夹危险代码
`USER`	`PROMPT_GUARD`	用户输入主要防提示注入
`ASSISTANT`	`CODE_SHIELD`	模型输出主要防生成危险代码
`SYSTEM`	(空)	系统提示通常可信
`MEMORY`	(空)	默认不扫

库还内置了两个用例预设(config.py:21),用 LlamaFirewall.from_usecase(...) 一键套用:

预设 `UseCase`	配置
`CHAT_BOT`	`USER` 和 `SYSTEM` 都用 PromptGuard
`CODING_ASSISTANT`	`ASSISTANT` 和 `TOOL` 都用 CodeShield

3. 扫描器工厂:懒加载 + 注册表

引擎不直接 import 所有扫描器(那样会把 torch、transformers、codeshield 全拖进来)。它用一个工厂函数 create_scanner(scanner_type)(llamafirewall.py:46)按需才 import 并实例化:

# 示意,非源码:工厂的核心分支
def create_scanner(scanner_type):
    if isinstance(scanner_type, str) and scanner_type in custom_scanner_registry:
        return custom_scanner_registry[scanner_type]()      # 自定义扫描器
    if scanner_type == ScannerType.PROMPT_GUARD:
        from .scanners.prompt_guard_scanner import PromptGuardScanner  # 用到才 import
        return PromptGuardScanner()
    # ... 其余类型同理

重点看两件事:

懒加载:每种扫描器的重依赖只在真正用到时才加载(scanners/__init__.py 里的 __getattr__ 也是同一思路)。
自定义优先:先查 custom_scanner_registry(第 4 章讲注册),命中就用你的类。

⚠ 一个性能坑(inferred): create_scanner 在每次 scan() 的循环里被调用 (llamafirewall.py:117),也就是每扫一条消息就新建一次扫描器实例。对 PromptGuard 这类构造函数里会 _load_model_and_tokenizer() 的扫描器,意味着可能反复从本地缓存重载模型。高吞吐场景要留意。

4. `scan()`:扫单条消息 + 汇总裁决

scan(input, trace=None)(llamafirewall.py:107)是最核心的方法。流程:

1. scanners = self.scanners.get(input.role, [])   # 按角色取扫描器,没配就空
2. for scanner_type in scanners:
     scanner = create_scanner(scanner_type)
     result  = asyncio.run(scanner.scan(input, trace))   # 同步壳里跑异步 scan
     记录 result.reason 和 result.score
     decisions[result.decision] = max(该裁决已有分, result.score)
3. 汇总 → ScanResult

汇总规则(关键)

只有一个扫描器时,直接返回它的结果(llamafirewall.py:133)。
多个扫描器时(llamafirewall.py:144):
- 只要 decisions 里出现 BLOCK,最终裁决就是 BLOCK,分数取 BLOCK 那一档的分。
- 否则,取分数最高的那个裁决(max(decisions, key=decisions.get))。
- reason 把各扫描器的原因拼起来去重("; ".join(set(reasons)))。

一个值得注意的不对称(inferred): scan() 的汇总只把 BLOCK 当特殊值,没有为 HUMAN_IN_THE_LOOP_REQUIRED 单列优先级——它只能靠「分数最高」这条路胜出。而下面的 scan_async 和 scan_replay 循环里却把 HITL 和 BLOCK 平等对待(都触发短路)。所以多扫描器 + 同步 scan() 的场景里,HITL 的处理不如在重放循环里直接。

5. 三个入口方法,差在「短路时机」

方法	输入	跑法	短路条件	用途
`scan`	单条消息	同步(内部 `asyncio.run` 逐个跑)	无(跑完所有扫描器再汇总)	最常用
`scan_async`	单条消息	异步	一旦某扫描器 BLOCK 或 HITL 就立刻返回	集成到异步框架
`scan_replay`	整段 Trace	逐条调 `scan`	某条被 BLOCK 或 HITL 就 `break`	扫历史对话

`scan_async`:遇到坏消息就早退

scan_async(llamafirewall.py:164)不做「跑完再汇总」,而是边跑边看:任一扫描器返回 BLOCK 或 HITL,立即返回该结果;全部放行才返回默认 ALLOW。省算力,适合在线路径。

`scan_replay`:重放整段对话

scan_replay(trace)(llamafirewall.py:184)用来回放一整段对话,逐条扫描。它的巧妙点在 给每条消息喂的是「它之前的历史」:

# 示意,非源码:重放时构造「过去的轨迹」
for current_ix, message in enumerate(trace):
    past_trace = [msg for ix, msg in enumerate(trace) if ix < current_ix]  # 只取当前之前的
    scan_result = self.scan(message, past_trace if past_trace else None)
    if scan_result.decision in (BLOCK, HITL):
        break   # 一旦出问题,后面不再扫

这正是 AlignmentCheck 需要的:判断「第 N 步动作是否跑偏」时,要看用户最初的目标 + 到目前为止的轨迹(见第 2 章)。真实实现见 llamafirewall.py:197-206。

`scan_replay_build_trace`:增量维护轨迹

还有一个增量版 scan_replay_build_trace(message, stored_trace)(llamafirewall.py:208):扫一条新消息, 只有裁决是 ALLOW 时才把它追加进轨迹,BLOCK/HITL 的消息不进历史。适合流式对话里边收边扫。

上述方法都有 _async 版本(scan_replay_async、scan_replay_build_trace_async),逻辑一致,只是用 await scan_async 替代 asyncio.run(scan),方便在 Jupyter / async 上下文里用。

6. 小结

数据模型极简:Message(role, content) + ScanResult(decision, reason, score)。
策略就是 Role → [Scanner] 一张表,可自定义、有默认、有预设。
引擎按角色路由、懒加载扫描器、汇总裁决;BLOCK 一票否决是核心规则。
单条用 scan/scan_async,整段用 scan_replay,后者给每条消息喂「它之前的历史」。

下一章进入具体扫描器,看每种风险到底怎么被检出来。

1. 先认识数据模型(全在一个文件里)​

消息与角色​

裁决结果​

2. 策略 = 一张 Role → [Scanner] 的表​

3. 扫描器工厂:懒加载 + 注册表​

4. scan():扫单条消息 + 汇总裁决​

汇总规则(关键)​

5. 三个入口方法,差在「短路时机」​

scan_async:遇到坏消息就早退​

scan_replay:重放整段对话​

scan_replay_build_trace:增量维护轨迹​

6. 小结​