211 lines
7.4 KiB
Python
211 lines
7.4 KiB
Python
"""AlignmentGuard - 对齐守卫：约束注入 + 级联故障检测"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class AlignmentConfig:
|
||
"""对齐守卫配置"""
|
||
|
||
constraints: list[str] = field(default_factory=list)
|
||
cascade_max_interactions: int = 10
|
||
cascade_max_depth: int = 3
|
||
audit_enabled: bool = False
|
||
audit_model: str = "default"
|
||
|
||
|
||
@dataclass
|
||
class AlignmentCheckResult:
|
||
"""对齐检查结果"""
|
||
|
||
passed: bool
|
||
violations: list[str] = field(default_factory=list)
|
||
checked_by: str = "" # "rule" or "llm"
|
||
|
||
|
||
@dataclass
|
||
class CascadeAlert:
|
||
"""级联故障告警"""
|
||
|
||
session_id: str
|
||
alert_type: str # "interaction_limit" or "loop_depth"
|
||
current_value: int
|
||
threshold: int
|
||
message: str
|
||
|
||
|
||
class ConstraintInjector:
|
||
"""将全局约束注入到任务 input_data 中"""
|
||
|
||
def __init__(self, config: AlignmentConfig):
|
||
self._config = config
|
||
|
||
def inject(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||
"""注入约束指令到 input_data
|
||
|
||
在 input_data 中添加 'alignment_constraints' 键，值为约束列表。
|
||
不修改原始 dict，返回新 dict。
|
||
"""
|
||
result = {**input_data, "alignment_constraints": list(self._config.constraints)}
|
||
return result
|
||
|
||
|
||
class AlignmentGuard:
|
||
"""对齐守卫 — 扩展 QualityGate，增加约束注入和级联检测"""
|
||
|
||
def __init__(self, config: AlignmentConfig, llm_gateway=None):
|
||
self._config = config
|
||
self._injector = ConstraintInjector(config)
|
||
self._llm_gateway = llm_gateway
|
||
self._interaction_counts: dict[str, int] = {}
|
||
self._loop_depths: dict[str, int] = {}
|
||
|
||
def inject_constraints(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||
"""委托给 ConstraintInjector"""
|
||
return self._injector.inject(input_data)
|
||
|
||
async def check_output(
|
||
self,
|
||
output: dict[str, Any],
|
||
constraints: list[str] | None = None,
|
||
) -> AlignmentCheckResult:
|
||
"""检查输出是否符合约束
|
||
|
||
- 系统级约束：基于规则的检查（关键词 + 正则匹配）
|
||
- 组织级约束：LLM 语义检查（仅当 audit_enabled=True）
|
||
"""
|
||
effective_constraints = constraints if constraints is not None else self._config.constraints
|
||
if not effective_constraints:
|
||
return AlignmentCheckResult(passed=True, checked_by="rule")
|
||
|
||
# 1. 基于规则的检查：关键词/子串匹配
|
||
violations = self._rule_check(output, effective_constraints)
|
||
if violations:
|
||
return AlignmentCheckResult(
|
||
passed=False,
|
||
violations=violations,
|
||
checked_by="rule",
|
||
)
|
||
|
||
# 2. LLM 语义检查（仅当 audit_enabled=True 且有 llm_gateway）
|
||
if self._config.audit_enabled and self._llm_gateway is not None:
|
||
return await self._llm_check(output, effective_constraints)
|
||
|
||
return AlignmentCheckResult(passed=True, checked_by="rule")
|
||
|
||
def _rule_check(
|
||
self, output: dict[str, Any], constraints: list[str]
|
||
) -> list[str]:
|
||
"""基于规则的约束检查：将 output 内容拼接后做关键词/子串匹配"""
|
||
content = self._extract_text(output)
|
||
violations: list[str] = []
|
||
for constraint in constraints:
|
||
# 简单子串匹配：约束关键词出现在输出中即视为违规
|
||
if constraint.lower() in content.lower():
|
||
violations.append(constraint)
|
||
return violations
|
||
|
||
@staticmethod
|
||
def _extract_text(output: dict[str, Any]) -> str:
|
||
"""从 output dict 中提取所有文本内容"""
|
||
parts: list[str] = []
|
||
for value in output.values():
|
||
if isinstance(value, str):
|
||
parts.append(value)
|
||
else:
|
||
parts.append(str(value))
|
||
return " ".join(parts)
|
||
|
||
async def _llm_check(
|
||
self, output: dict[str, Any], constraints: list[str]
|
||
) -> AlignmentCheckResult:
|
||
"""LLM 语义检查"""
|
||
content = self._extract_text(output)
|
||
constraint_text = "\n".join(f"- {c}" for c in constraints)
|
||
messages = [
|
||
{
|
||
"role": "system",
|
||
"content": (
|
||
"You are an alignment auditor. Check if the following output "
|
||
"violates any of the listed constraints. "
|
||
"Reply with 'PASS' if no violations, or list the violated constraints."
|
||
),
|
||
},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
f"Constraints:\n{constraint_text}\n\nOutput:\n{content}"
|
||
),
|
||
},
|
||
]
|
||
try:
|
||
response = await self._llm_gateway.chat(
|
||
messages=messages, model=self._config.audit_model
|
||
)
|
||
reply = response.content.strip()
|
||
if reply.upper().startswith("PASS"):
|
||
return AlignmentCheckResult(passed=True, checked_by="llm")
|
||
else:
|
||
return AlignmentCheckResult(
|
||
passed=False,
|
||
violations=[reply],
|
||
checked_by="llm",
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"LLM audit failed: {e}")
|
||
return AlignmentCheckResult(
|
||
passed=False,
|
||
violations=[f"LLM audit unavailable: {e}"],
|
||
checked_by="rule",
|
||
)
|
||
|
||
def record_interaction(self, session_id: str) -> CascadeAlert | None:
|
||
"""记录一次 agent 间交互，超过阈值返回 CascadeAlert"""
|
||
self._interaction_counts[session_id] = (
|
||
self._interaction_counts.get(session_id, 0) + 1
|
||
)
|
||
count = self._interaction_counts[session_id]
|
||
if count > self._config.cascade_max_interactions:
|
||
return CascadeAlert(
|
||
session_id=session_id,
|
||
alert_type="interaction_limit",
|
||
current_value=count,
|
||
threshold=self._config.cascade_max_interactions,
|
||
message=(
|
||
f"Session {session_id} exceeded max interactions: "
|
||
f"{count} > {self._config.cascade_max_interactions}"
|
||
),
|
||
)
|
||
return None
|
||
|
||
def record_loop_depth(self, session_id: str, depth: int) -> CascadeAlert | None:
|
||
"""记录循环深度，超过阈值返回 CascadeAlert"""
|
||
self._loop_depths[session_id] = depth
|
||
if depth > self._config.cascade_max_depth:
|
||
return CascadeAlert(
|
||
session_id=session_id,
|
||
alert_type="loop_depth",
|
||
current_value=depth,
|
||
threshold=self._config.cascade_max_depth,
|
||
message=(
|
||
f"Session {session_id} exceeded max loop depth: "
|
||
f"{depth} > {self._config.cascade_max_depth}"
|
||
),
|
||
)
|
||
return None
|
||
|
||
def reset_session(self, session_id: str) -> None:
|
||
"""重置某个 session 的交互计数"""
|
||
self._interaction_counts.pop(session_id, None)
|
||
self._loop_depths.pop(session_id, None)
|
||
|
||
def get_interaction_count(self, session_id: str) -> int:
|
||
"""获取某个 session 的当前交互计数"""
|
||
return self._interaction_counts.get(session_id, 0)
|