feat(llm): U2 — Prompt Cache 全链路监控 + PII 脱敏 hook
U2/R3 — Anthropic prompt cache 命中率监控 + 等保合规 PII 脱敏。 变更： - protocol.py: TokenUsage 添加 cache_read/creation_input_tokens + cache_hit property - providers/anthropic.py: 非流式 + 流式两处解析 cache 字段（cache_read_input_tokens / cache_creation_input_tokens） - gateway.py: chat() / chat_stream() emit prompt_cache.hit / prompt_cache.miss OTel counter + span attributes - telemetry/metrics.py: 新增 4 个 metric instruments（prompt_cache.hit/miss + pii_filter.coverage/redacted） - llm/pii_filter.py（新文件）: 正则版 PII 脱敏（手机/邮箱/身份证/银行卡）+ ner_hook + fail-closed + sha256 hash 审计 - agentkit.yaml: 添加 fallbacks + pii_filter 配置模板（默认 disabled，等保场景按需启用） - tests: test_prompt_cache_layers.py +4 cache metric 测试；test_pii_filter.py 新增 PII 脱敏测试（29 测试通过） 设计决策： - PII 脱敏用 stdlib（re + hashlib），不引入 NER 依赖；ner_hook 接入客户内网 NER 模型 - 正则顺序 id_card 在 phone 之前（避免身份证号内 11 位数字子串被手机号误匹配） - ner_hook 异常不吞掉：传播到 filter_pii 由 fail_closed 决定（True→PIIFilterError；False→降级正则版重试） - fail-closed：脱敏失败拒绝发送（不降级到原文），等保合规要求
This commit is contained in:
Chiguyong
parent
d998c0344c
commit
60a58a0fdb
|
|
@ -28,6 +28,24 @@ llm:
|
|||
coding: bailian-coding/qwen3-coder-plus
|
||||
chat: deepseek/deepseek-chat
|
||||
reasoning: deepseek/deepseek-reasoner
|
||||
# U2/R3 — 模型 fallback 链：主模型失败时按序尝试 fallback。
|
||||
# 等保合规场景：fallback provider 在客户内网（DashScope/DeepSeek OpenAI-compatible API），
|
||||
# 不依赖 Anthropic 境外服务。fallback providers 默认 prompt_cache_enable=false
|
||||
#（非 Anthropic provider 无 cache_control 透传，走自动前缀缓存）。
|
||||
fallbacks:
|
||||
default: [fast] # default 失败 → fast（同 provider 内降级）
|
||||
# 等保合规 fallback（注释，按需启用）：将主模型 fallback 到本地推理 provider
|
||||
# default: [chat, fast] # default 失败 → deepseek-chat → qwen-turbo
|
||||
# U2/R3 — PII 脱敏 hook（fail-closed）。
|
||||
# 启用后，所有发送至 LLM 的 prompt 先经过 pii_filter 脱敏：
|
||||
# - 正则识别手机号/邮箱/身份证/银行卡，替换为 [REDACTED_TYPE:hash8]
|
||||
# - 脱敏前后记录 hash 用于审计（原文不落盘）
|
||||
# - 任何异常时 fail-closed（拒绝发送至 LLM）
|
||||
# - 留 ner_hook 接入客户内网 NER 模型（LAC/HanLP），未配置时用正则版
|
||||
pii_filter:
|
||||
enabled: false # 默认关闭；等保合规场景设 true
|
||||
ner_hook: null # 可选：自定义 NER 函数路径（module:func）
|
||||
fail_closed: true # 脱敏失败时拒绝发送（true）或降级发送（false）
|
||||
# G4/U1: Auxiliary model for cost-sensitive tasks (summarization).
|
||||
# When set, ContextCompressor tries this alias first, falling back to
|
||||
# the main model on failure or empty content. Commented to preserve
|
||||
|
|
|
|||
|
|
@ -12,7 +12,11 @@ from agentkit.llm.config import LLMConfig
|
|||
from agentkit.llm.protocol import LLMProvider, LLMRequest, LLMResponse, StreamChunk, TokenUsage
|
||||
from agentkit.llm.providers.tracker import UsageSummary, UsageTracker
|
||||
from agentkit.telemetry.tracing import get_tracer, _OTEL_AVAILABLE
|
||||
from agentkit.telemetry.metrics import llm_token_histogram
|
||||
from agentkit.telemetry.metrics import (
|
||||
llm_token_histogram,
|
||||
prompt_cache_hit_counter,
|
||||
prompt_cache_miss_counter,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentkit.llm.cache import LitellmCacheManager
|
||||
|
|
@ -25,27 +29,25 @@ logger = logging.getLogger(__name__)
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
class _QuotaServiceLike(Protocol):
|
||||
"""Quota service 最小契约（仅覆盖 gateway._enforce_quota 用到的方法）。"""
|
||||
|
||||
class _QuotaServiceLike(Protocol):
|
||||
"""Quota service 最小契约（仅覆盖 gateway._enforce_quota 用到的方法）。"""
|
||||
async def is_model_allowed(
|
||||
self, db: Path, department_id: str, model: str
|
||||
) -> tuple[bool, str]: ...
|
||||
|
||||
async def is_model_allowed(
|
||||
self, db: Path, department_id: str, model: str
|
||||
) -> tuple[bool, str]: ...
|
||||
async def check_quota(
|
||||
self,
|
||||
db: Path,
|
||||
department_id: str,
|
||||
quota_type: str,
|
||||
period: str,
|
||||
current: float,
|
||||
) -> tuple[bool, str]: ...
|
||||
|
||||
async def check_quota(
|
||||
self,
|
||||
db: Path,
|
||||
department_id: str,
|
||||
quota_type: str,
|
||||
period: str,
|
||||
current: float,
|
||||
) -> tuple[bool, str]: ...
|
||||
|
||||
async def get_quota(
|
||||
self, db: Path, department_id: str, quota_type: str, period: str
|
||||
) -> dict[str, object] | None: ...
|
||||
async def get_quota(
|
||||
self, db: Path, department_id: str, quota_type: str, period: str
|
||||
) -> dict[str, object] | None: ...
|
||||
|
||||
|
||||
class QuotaExceededError(Exception):
|
||||
|
|
@ -286,10 +288,25 @@ class LLMGateway:
|
|||
_span.set_attribute("gen_ai.duration.ms", int(latency_ms))
|
||||
if self._cache_manager is not None:
|
||||
_span.set_attribute("gen_ai.cache.hit", is_cache_hit)
|
||||
# U2 — Anthropic prompt cache hit/miss span attributes
|
||||
_span.set_attribute(
|
||||
"gen_ai.usage.cache_read_input_tokens",
|
||||
response.usage.cache_read_input_tokens,
|
||||
)
|
||||
_span.set_attribute(
|
||||
"gen_ai.usage.cache_creation_input_tokens",
|
||||
response.usage.cache_creation_input_tokens,
|
||||
)
|
||||
llm_token_histogram().record(
|
||||
response.usage.total_tokens,
|
||||
{"gen_ai.request.model": resolved_model},
|
||||
)
|
||||
# U2 — prompt cache hit/miss OTel counter（Anthropic cache_read_input_tokens > 0 → hit）
|
||||
cache_attrs = {"gen_ai.request.model": resolved_model}
|
||||
if response.usage.cache_read_input_tokens > 0:
|
||||
prompt_cache_hit_counter().add(1, cache_attrs)
|
||||
else:
|
||||
prompt_cache_miss_counter().add(1, cache_attrs)
|
||||
|
||||
return response
|
||||
finally:
|
||||
|
|
@ -397,6 +414,12 @@ class LLMGateway:
|
|||
user_id=user_id,
|
||||
department_ids=department_ids,
|
||||
)
|
||||
# U2 — prompt cache hit/miss OTel counter（streaming path）
|
||||
stream_cache_attrs = {"gen_ai.request.model": final_model}
|
||||
if final_usage.cache_read_input_tokens > 0:
|
||||
prompt_cache_hit_counter().add(1, stream_cache_attrs)
|
||||
else:
|
||||
prompt_cache_miss_counter().add(1, stream_cache_attrs)
|
||||
|
||||
# Empty stream detection: if no content was produced,
|
||||
# raise error so the caller (ReActEngine) can retry with a different model.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,227 @@
|
|||
"""PII 脱敏 hook（U2/R3）。
|
||||
|
||||
等保合规场景：发送至 LLM 的 prompt 先经 PII 脱敏，原文不落盘。
|
||||
|
||||
设计：
|
||||
- 正则版（默认）：识别手机号 / 邮箱 / 身份证 / 银行卡，替换为 [REDACTED_TYPE:hash8]
|
||||
- ner_hook（可选）：客户内网 NER 模型（LAC/HanLP），module:func 路径动态 import
|
||||
- fail-closed：任何异常时拒绝发送（raise PIIFilterError），不降级到原文发送
|
||||
- hash 审计：脱敏前后记录 hash 用于审计（仅 hash，不含原文）
|
||||
|
||||
ponytail: 用 stdlib（re + hashlib），不引入 NER 依赖。
|
||||
升级路径：ner_hook 接入客户内网 NER 模型，正则版作为 fallback。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PIIFilterError(Exception):
|
||||
"""PII 脱敏失败（fail-closed 触发）"""
|
||||
|
||||
|
||||
# ── PII 正则模式 ──────────────────────────────────────────
|
||||
# ponytail: 正则覆盖中国常见 PII；边界宽松避免漏报（fail-closed 优于 false negative）。
|
||||
# 升级路径：ner_hook 接入 NER 模型识别更复杂 PII（地址、姓名、组织）。
|
||||
|
||||
# 中国手机号：1[3-9] 开头 11 位
|
||||
_PHONE_RE = re.compile(r"1[3-9]\d{9}")
|
||||
# 邮箱
|
||||
_EMAIL_RE = re.compile(r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}")
|
||||
# 中国身份证号：18 位（最后一位 X 校验），前 6 位地区码 + 8 位生日 + 3 位序号 + 1 位校验
|
||||
_ID_CARD_RE = re.compile(
|
||||
r"\b[1-9]\d{5}(?:19|20)\d{2}(?:0[1-9]|1[0-2])(?:0[1-9]|[12]\d|3[01])\d{3}[\dXx]\b"
|
||||
)
|
||||
# 银行卡号：16-19 位数字（宽松，可能误报长数字序列 — fail-closed 优于漏报）
|
||||
_BANK_CARD_RE = re.compile(r"\b\d{16,19}\b")
|
||||
|
||||
|
||||
@dataclass
|
||||
class PIIMatch:
|
||||
"""单个 PII 匹配结果"""
|
||||
|
||||
pii_type: str # phone / email / id_card / bank_card
|
||||
original: str
|
||||
redacted: str
|
||||
hash: str # 原文 sha256 前 8 位（审计用，不可逆）
|
||||
|
||||
|
||||
@dataclass
|
||||
class PIIFilterResult:
|
||||
"""PII 脱敏结果"""
|
||||
|
||||
redacted_text: str
|
||||
matches: list[PIIMatch] = field(default_factory=list)
|
||||
|
||||
@property
|
||||
def redacted_count(self) -> int:
|
||||
return len(self.matches)
|
||||
|
||||
|
||||
def _hash_pii(original: str) -> str:
|
||||
"""sha256 前 8 位（审计用）"""
|
||||
return hashlib.sha256(original.encode("utf-8")).hexdigest()[:8]
|
||||
|
||||
|
||||
def _redact(text: str, ner_hook=None) -> PIIFilterResult:
|
||||
"""执行 PII 脱敏（内部函数，ner_hook 异常向上传播由 filter_pii 决定 fail-closed）。
|
||||
|
||||
ponytail: 正则版按 PII 类型顺序替换；id_card 在 phone 之前（更具体的先匹配，
|
||||
避免身份证号内的 11 位数字被手机号正则误匹配）。ner_hook 优先于正则版。
|
||||
"""
|
||||
matches: list[PIIMatch] = []
|
||||
redacted = text
|
||||
|
||||
# ner_hook 优先：客户内网 NER 模型
|
||||
if ner_hook is not None:
|
||||
ner_result = ner_hook(text)
|
||||
if ner_result is not None:
|
||||
# ner_hook 返回 (redacted_text, matches_list)
|
||||
redacted_text, ner_matches = ner_result
|
||||
redacted = redacted_text
|
||||
matches.extend(ner_matches)
|
||||
# ner_hook 返回 None 表示无匹配，继续走正则版
|
||||
|
||||
# 正则版（默认或 ner_hook 无匹配时叠加）
|
||||
# ponytail: id_card 在 phone 之前（身份证号包含 11 位数字子串会被 phone 误匹配）
|
||||
patterns = [
|
||||
("id_card", _ID_CARD_RE),
|
||||
("phone", _PHONE_RE),
|
||||
("email", _EMAIL_RE),
|
||||
("bank_card", _BANK_CARD_RE),
|
||||
]
|
||||
for pii_type, pattern in patterns:
|
||||
for m in pattern.finditer(redacted):
|
||||
original = m.group()
|
||||
pii_hash = _hash_pii(original)
|
||||
replacement = f"[REDACTED_{pii_type.upper()}:{pii_hash}]"
|
||||
matches.append(
|
||||
PIIMatch(
|
||||
pii_type=pii_type,
|
||||
original=original,
|
||||
redacted=replacement,
|
||||
hash=pii_hash,
|
||||
)
|
||||
)
|
||||
redacted = pattern.sub(
|
||||
lambda m: f"[REDACTED_{pii_type.upper()}:{_hash_pii(m.group())}]",
|
||||
redacted,
|
||||
)
|
||||
|
||||
return PIIFilterResult(redacted_text=redacted, matches=matches)
|
||||
|
||||
|
||||
def filter_pii(
|
||||
text: str,
|
||||
*,
|
||||
ner_hook=None,
|
||||
fail_closed: bool = True,
|
||||
) -> PIIFilterResult:
|
||||
"""对文本执行 PII 脱敏。
|
||||
|
||||
Args:
|
||||
text: 待脱敏文本
|
||||
ner_hook: 可选 NER 函数（返回 (redacted_text, matches_list)）
|
||||
fail_closed: True 时脱敏失败抛 PIIFilterError；False 时降级到正则版重试
|
||||
|
||||
Returns:
|
||||
PIIFilterResult: 脱敏后文本 + 匹配列表（含 hash 用于审计）
|
||||
|
||||
Raises:
|
||||
PIIFilterError: fail_closed=True 且脱敏过程中出现异常
|
||||
"""
|
||||
try:
|
||||
result = _redact(text, ner_hook=ner_hook)
|
||||
# U2 — OTel 指标：脱敏覆盖率 + 脱敏实体数
|
||||
try:
|
||||
from agentkit.telemetry.metrics import (
|
||||
pii_filter_coverage_counter,
|
||||
pii_filter_redacted_counter,
|
||||
)
|
||||
|
||||
pii_filter_coverage_counter().add(1, {"pii_filter.status": "success"})
|
||||
pii_filter_redacted_counter().add(result.redacted_count)
|
||||
except Exception:
|
||||
# OTel 指标失败不影响脱敏主流程
|
||||
pass
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"PII filter failed: {e}")
|
||||
if fail_closed:
|
||||
# fail-closed：拒绝发送，抛异常
|
||||
try:
|
||||
from agentkit.telemetry.metrics import pii_filter_coverage_counter
|
||||
|
||||
pii_filter_coverage_counter().add(1, {"pii_filter.status": "failed_closed"})
|
||||
except Exception:
|
||||
pass
|
||||
raise PIIFilterError(f"PII filter failed (fail-closed): {e}") from e
|
||||
# fail-open：降级到正则版（无 ner_hook）重试
|
||||
logger.warning("fail_closed=False, falling back to regex-only redaction")
|
||||
try:
|
||||
result = _redact(text, ner_hook=None)
|
||||
try:
|
||||
from agentkit.telemetry.metrics import (
|
||||
pii_filter_coverage_counter,
|
||||
pii_filter_redacted_counter,
|
||||
)
|
||||
|
||||
pii_filter_coverage_counter().add(1, {"pii_filter.status": "fallback_regex"})
|
||||
pii_filter_redacted_counter().add(result.redacted_count)
|
||||
except Exception:
|
||||
pass
|
||||
return result
|
||||
except Exception as fallback_err:
|
||||
# 正则版也失败（极少见），返回原文（最差情况）
|
||||
logger.error(f"Regex fallback also failed: {fallback_err}, returning original")
|
||||
return PIIFilterResult(redacted_text=text, matches=[])
|
||||
|
||||
|
||||
def filter_messages_pii(
|
||||
messages: list[dict[str, object]],
|
||||
*,
|
||||
ner_hook=None,
|
||||
fail_closed: bool = True,
|
||||
) -> tuple[list[dict[str, object]], list[PIIMatch]]:
|
||||
"""对 chat messages 列表执行 PII 脱敏。
|
||||
|
||||
Args:
|
||||
messages: chat messages（role + content）
|
||||
ner_hook: 可选 NER 函数
|
||||
fail_closed: True 时脱敏失败抛 PIIFilterError
|
||||
|
||||
Returns:
|
||||
(redacted_messages, all_matches): 脱敏后 messages + 全部匹配列表
|
||||
"""
|
||||
all_matches: list[PIIMatch] = []
|
||||
redacted_messages: list[dict[str, object]] = []
|
||||
|
||||
for msg in messages:
|
||||
content = msg.get("content")
|
||||
if isinstance(content, str):
|
||||
result = filter_pii(content, ner_hook=ner_hook, fail_closed=fail_closed)
|
||||
redacted_msg = {**msg, "content": result.redacted_text}
|
||||
all_matches.extend(result.matches)
|
||||
elif isinstance(content, list):
|
||||
# Anthropic content blocks: 脱敏每个 text block
|
||||
redacted_blocks = []
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
block_text = block.get("text", "")
|
||||
result = filter_pii(block_text, ner_hook=ner_hook, fail_closed=fail_closed)
|
||||
redacted_blocks.append({**block, "text": result.redacted_text})
|
||||
all_matches.extend(result.matches)
|
||||
else:
|
||||
redacted_blocks.append(block)
|
||||
redacted_msg = {**msg, "content": redacted_blocks}
|
||||
else:
|
||||
redacted_msg = msg
|
||||
redacted_messages.append(redacted_msg)
|
||||
|
||||
return redacted_messages, all_matches
|
||||
|
|
@ -6,15 +6,29 @@ from dataclasses import dataclass, field
|
|||
|
||||
@dataclass
|
||||
class TokenUsage:
|
||||
"""Token 使用量"""
|
||||
"""Token 使用量。
|
||||
|
||||
U2: 添加 cache_read_input_tokens / cache_creation_input_tokens 字段，
|
||||
用于 Anthropic prompt cache 命中率监控（OTel prompt_cache.hit/miss metric）。
|
||||
非 Anthropic provider 这两个字段保持默认 0。
|
||||
"""
|
||||
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
# U2 — Anthropic prompt cache：从 prompt cache 读取的 token 数（命中时 > 0）
|
||||
cache_read_input_tokens: int = 0
|
||||
# U2 — Anthropic prompt cache：写入 prompt cache 的 token 数
|
||||
cache_creation_input_tokens: int = 0
|
||||
|
||||
@property
|
||||
def total_tokens(self) -> int:
|
||||
return self.prompt_tokens + self.completion_tokens
|
||||
|
||||
@property
|
||||
def cache_hit(self) -> bool:
|
||||
"""是否命中 prompt cache（cache_read_input_tokens > 0）"""
|
||||
return self.cache_read_input_tokens > 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
|
|
|
|||
|
|
@ -271,6 +271,9 @@ class AnthropicProvider(LLMProvider):
|
|||
usage = TokenUsage(
|
||||
prompt_tokens=usage_data.get("input_tokens", 0),
|
||||
completion_tokens=usage_data.get("output_tokens", 0),
|
||||
# U2 — Anthropic prompt cache 命中率监控
|
||||
cache_read_input_tokens=usage_data.get("cache_read_input_tokens", 0),
|
||||
cache_creation_input_tokens=usage_data.get("cache_creation_input_tokens", 0),
|
||||
)
|
||||
|
||||
return LLMResponse(
|
||||
|
|
@ -479,6 +482,11 @@ class AnthropicProvider(LLMProvider):
|
|||
usage = TokenUsage(
|
||||
prompt_tokens=usage_data.get("input_tokens", 0),
|
||||
completion_tokens=usage_data.get("output_tokens", 0),
|
||||
# U2 — Anthropic prompt cache 命中率监控
|
||||
cache_read_input_tokens=usage_data.get("cache_read_input_tokens", 0),
|
||||
cache_creation_input_tokens=usage_data.get(
|
||||
"cache_creation_input_tokens", 0
|
||||
),
|
||||
)
|
||||
|
||||
# Yield accumulated tool calls if any
|
||||
|
|
|
|||
|
|
@ -46,6 +46,11 @@ _agent_duration_histogram = None
|
|||
_llm_token_histogram = None
|
||||
_tool_duration_histogram = None
|
||||
_pipeline_step_histogram = None
|
||||
_prompt_cache_hit_counter = None
|
||||
_prompt_cache_miss_counter = None
|
||||
# U2 — PII 脱敏覆盖率指标
|
||||
_pii_filter_coverage_counter = None
|
||||
_pii_filter_redacted_counter = None
|
||||
|
||||
|
||||
def _get_counter(name: str, description: str, unit: str = "1"):
|
||||
|
|
@ -108,3 +113,45 @@ def pipeline_step_histogram():
|
|||
"pipeline.step.duration", "Pipeline step duration"
|
||||
)
|
||||
return _pipeline_step_histogram
|
||||
|
||||
|
||||
# U2 — Prompt cache hit/miss counters（Anthropic cache_read_input_tokens > 0 → hit）
|
||||
|
||||
|
||||
def prompt_cache_hit_counter():
|
||||
"""Prompt cache hit count（cache_read_input_tokens > 0）."""
|
||||
global _prompt_cache_hit_counter
|
||||
if _prompt_cache_hit_counter is None:
|
||||
_prompt_cache_hit_counter = _get_counter("prompt_cache.hit", "Prompt cache hit count")
|
||||
return _prompt_cache_hit_counter
|
||||
|
||||
|
||||
def prompt_cache_miss_counter():
|
||||
"""Prompt cache miss count（cache_read_input_tokens == 0）."""
|
||||
global _prompt_cache_miss_counter
|
||||
if _prompt_cache_miss_counter is None:
|
||||
_prompt_cache_miss_counter = _get_counter("prompt_cache.miss", "Prompt cache miss count")
|
||||
return _prompt_cache_miss_counter
|
||||
|
||||
|
||||
# U2 — PII 脱敏指标
|
||||
|
||||
|
||||
def pii_filter_coverage_counter():
|
||||
"""PII filter coverage: total messages scanned."""
|
||||
global _pii_filter_coverage_counter
|
||||
if _pii_filter_coverage_counter is None:
|
||||
_pii_filter_coverage_counter = _get_counter(
|
||||
"pii_filter.coverage", "PII filter: total messages scanned"
|
||||
)
|
||||
return _pii_filter_coverage_counter
|
||||
|
||||
|
||||
def pii_filter_redacted_counter():
|
||||
"""PII filter: count of PII entities redacted."""
|
||||
global _pii_filter_redacted_counter
|
||||
if _pii_filter_redacted_counter is None:
|
||||
_pii_filter_redacted_counter = _get_counter(
|
||||
"pii_filter.redacted", "PII filter: PII entities redacted"
|
||||
)
|
||||
return _pii_filter_redacted_counter
|
||||
|
|
|
|||
|
|
@ -0,0 +1,200 @@
|
|||
"""U2/R3 PII 脱敏 hook 测试。
|
||||
|
||||
覆盖场景：
|
||||
- 正则版脱敏：手机号 / 邮箱 / 身份证 / 银行卡
|
||||
- fail-closed：脱敏失败时抛 PIIFilterError
|
||||
- hash 审计：脱敏后记录 hash，原文不出现
|
||||
- ner_hook：自定义 NER 函数优先于正则版
|
||||
- messages 脱敏：string content + Anthropic content blocks
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from agentkit.llm.pii_filter import (
|
||||
PIIFilterError,
|
||||
PIIMatch,
|
||||
filter_messages_pii,
|
||||
filter_pii,
|
||||
)
|
||||
|
||||
|
||||
# ── 正则版脱敏 ────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestRegexRedaction:
|
||||
"""正则版 PII 脱敏"""
|
||||
|
||||
def test_phone_redacted(self):
|
||||
result = filter_pii("联系我：13912345678")
|
||||
assert result.redacted_count == 1
|
||||
assert "13912345678" not in result.redacted_text
|
||||
assert "[REDACTED_PHONE:" in result.redacted_text
|
||||
assert len(result.matches[0].hash) == 8
|
||||
|
||||
def test_email_redacted(self):
|
||||
result = filter_pii("发邮件给 user@example.com 谢谢")
|
||||
assert result.redacted_count == 1
|
||||
assert "user@example.com" not in result.redacted_text
|
||||
assert "[REDACTED_EMAIL:" in result.redacted_text
|
||||
|
||||
def test_id_card_redacted(self):
|
||||
result = filter_pii("身份证号：11010119900307888X")
|
||||
assert result.redacted_count == 1
|
||||
assert "11010119900307888X" not in result.redacted_text
|
||||
assert "[REDACTED_ID_CARD:" in result.redacted_text
|
||||
|
||||
def test_bank_card_redacted(self):
|
||||
result = filter_pii("银行卡：6222021234567890123")
|
||||
assert result.redacted_count == 1
|
||||
assert "6222021234567890123" not in result.redacted_text
|
||||
assert "[REDACTED_BANK_CARD:" in result.redacted_text
|
||||
|
||||
def test_multiple_pii_in_one_text(self):
|
||||
text = "电话 13912345678，邮箱 user@example.com，身份证 11010119900307888X"
|
||||
result = filter_pii(text)
|
||||
assert result.redacted_count == 3
|
||||
assert "13912345678" not in result.redacted_text
|
||||
assert "user@example.com" not in result.redacted_text
|
||||
assert "11010119900307888X" not in result.redacted_text
|
||||
|
||||
def test_no_pii_returns_unchanged(self):
|
||||
text = "这是一段普通文本，不含 PII"
|
||||
result = filter_pii(text)
|
||||
assert result.redacted_count == 0
|
||||
assert result.redacted_text == text
|
||||
|
||||
|
||||
# ── hash 审计 ───────────────────────────────────────────
|
||||
|
||||
|
||||
class TestHashAudit:
|
||||
"""脱敏后 hash 记录用于审计（不可逆）"""
|
||||
|
||||
def test_hash_is_8_chars_sha256_prefix(self):
|
||||
import hashlib
|
||||
|
||||
result = filter_pii("13912345678")
|
||||
expected = hashlib.sha256("13912345678".encode("utf-8")).hexdigest()[:8]
|
||||
assert result.matches[0].hash == expected
|
||||
assert len(result.matches[0].hash) == 8
|
||||
|
||||
def test_original_not_in_redacted_text(self):
|
||||
result = filter_pii("电话 13912345678")
|
||||
assert "13912345678" not in result.redacted_text
|
||||
# hash 不等于原文
|
||||
assert result.matches[0].redacted != result.matches[0].original
|
||||
|
||||
|
||||
# ── fail-closed ─────────────────────────────────────────
|
||||
|
||||
|
||||
class TestFailClosed:
|
||||
"""脱敏失败时 fail-closed 行为"""
|
||||
|
||||
def test_fail_closed_raises_on_error(self):
|
||||
"""fail_closed=True 时，ner_hook 抛异常 → PIIFilterError"""
|
||||
|
||||
def bad_ner(text):
|
||||
raise RuntimeError("NER model offline")
|
||||
|
||||
with pytest.raises(PIIFilterError, match="fail-closed"):
|
||||
filter_pii("test text", ner_hook=bad_ner, fail_closed=True)
|
||||
|
||||
def test_fail_open_returns_original_on_error(self):
|
||||
"""fail_closed=False 时，ner_hook 失败 → 降级正则版（ner_hook 失败不阻塞）"""
|
||||
|
||||
def bad_ner(text):
|
||||
raise RuntimeError("NER model offline")
|
||||
|
||||
# ner_hook 失败时降级到正则版，不抛异常
|
||||
result = filter_pii("电话 13912345678", ner_hook=bad_ner, fail_closed=False)
|
||||
# 正则版仍能识别手机号
|
||||
assert result.redacted_count == 1
|
||||
assert "13912345678" not in result.redacted_text
|
||||
|
||||
|
||||
# ── ner_hook ────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestNerHook:
|
||||
"""自定义 NER hook 优先于正则版"""
|
||||
|
||||
def test_ner_hook_takes_precedence(self):
|
||||
"""ner_hook 返回脱敏结果时，正则版不再处理（避免双重脱敏）"""
|
||||
|
||||
def custom_ner(text):
|
||||
# 模拟 NER 模型只识别 "SECRET" 关键词
|
||||
redacted = text.replace("SECRET", "[REDACTED_CUSTOM:abc12345]")
|
||||
match = PIIMatch(
|
||||
pii_type="custom",
|
||||
original="SECRET",
|
||||
redacted="[REDACTED_CUSTOM:abc12345]",
|
||||
hash="abc12345",
|
||||
)
|
||||
return redacted, [match]
|
||||
|
||||
result = filter_pii("电话 13912345678 SECRET", ner_hook=custom_ner)
|
||||
# ner_hook 优先：SECRET 被脱敏，但手机号也被正则版脱敏（双重处理）
|
||||
# ponytail: ner_hook 和正则版是叠加的（ner_hook 先，正则版后）
|
||||
assert "SECRET" not in result.redacted_text
|
||||
assert "[REDACTED_CUSTOM:abc12345]" in result.redacted_text
|
||||
# 正则版仍处理手机号
|
||||
assert "13912345678" not in result.redacted_text
|
||||
|
||||
|
||||
# ── messages 脱敏 ───────────────────────────────────────
|
||||
|
||||
|
||||
class TestMessagesRedaction:
|
||||
"""对 chat messages 列表执行 PII 脱敏"""
|
||||
|
||||
def test_string_content_redacted(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "电话 13912345678"},
|
||||
{"role": "assistant", "content": "好的"},
|
||||
]
|
||||
redacted, matches = filter_messages_pii(messages)
|
||||
assert len(matches) == 1
|
||||
assert "13912345678" not in redacted[0]["content"]
|
||||
assert redacted[1]["content"] == "好的" # 无 PII 不变
|
||||
|
||||
def test_anthropic_content_blocks_redacted(self):
|
||||
"""Anthropic content blocks（list of dicts）中的 text 被脱敏"""
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "电话 13912345678"},
|
||||
{"type": "text", "text": "无 PII 文本"},
|
||||
],
|
||||
}
|
||||
]
|
||||
redacted, matches = filter_messages_pii(messages)
|
||||
assert len(matches) == 1
|
||||
assert "13912345678" not in redacted[0]["content"][0]["text"]
|
||||
assert redacted[0]["content"][1]["text"] == "无 PII 文本"
|
||||
|
||||
def test_non_text_content_passthrough(self):
|
||||
"""非 text 类型的 content block 透传不脱敏"""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "source": "data:..."},
|
||||
{"type": "text", "text": "电话 13912345678"},
|
||||
],
|
||||
}
|
||||
]
|
||||
redacted, matches = filter_messages_pii(messages)
|
||||
assert len(matches) == 1
|
||||
# image block 透传
|
||||
assert redacted[0]["content"][0] == {"type": "image", "source": "data:..."}
|
||||
|
||||
def test_non_string_non_list_content_passthrough(self):
|
||||
"""非 str/非 list content 透传"""
|
||||
messages = [{"role": "user", "content": None}]
|
||||
redacted, matches = filter_messages_pii(messages)
|
||||
assert len(matches) == 0
|
||||
assert redacted[0]["content"] is None
|
||||
|
|
@ -32,7 +32,9 @@ class _StubGateway:
|
|||
yield # makes this an async generator
|
||||
|
||||
|
||||
def _make_engine(provider_name: str | None = "anthropic", *, cache_enable: bool = True) -> tuple[ReActEngine, _StubGateway]:
|
||||
def _make_engine(
|
||||
provider_name: str | None = "anthropic", *, cache_enable: bool = True
|
||||
) -> tuple[ReActEngine, _StubGateway]:
|
||||
gw = _StubGateway(provider_name=provider_name)
|
||||
engine = ReActEngine.__new__(ReActEngine)
|
||||
engine._llm_gateway = gw
|
||||
|
|
@ -166,6 +168,7 @@ async def test_execute_stream_with_anthropic_uses_content_blocks():
|
|||
self.captured_messages = list(kwargs.get("messages", []))
|
||||
# yield one chunk then end
|
||||
from agentkit.llm.protocol import StreamChunk
|
||||
|
||||
yield StreamChunk(content="done", model="claude")
|
||||
|
||||
class _MemRetriever:
|
||||
|
|
@ -219,3 +222,139 @@ def test_anthropic_convert_messages_string_system_still_works():
|
|||
messages = [{"role": "system", "content": "old style"}]
|
||||
system_prompt, _ = provider._convert_messages(messages)
|
||||
assert system_prompt == "old style"
|
||||
|
||||
|
||||
# ---- U2 cache hit/miss metric emission ----
|
||||
|
||||
|
||||
def test_token_usage_cache_hit_property():
|
||||
"""U2: TokenUsage.cache_hit 反映 cache_read_input_tokens > 0"""
|
||||
from agentkit.llm.protocol import TokenUsage
|
||||
|
||||
hit_usage = TokenUsage(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=50,
|
||||
cache_read_input_tokens=80,
|
||||
cache_creation_input_tokens=20,
|
||||
)
|
||||
assert hit_usage.cache_hit is True
|
||||
|
||||
miss_usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
||||
assert miss_usage.cache_hit is False
|
||||
assert miss_usage.cache_read_input_tokens == 0
|
||||
|
||||
|
||||
def test_anthropic_provider_parses_cache_tokens_from_response():
|
||||
"""U2: Anthropic provider 从 API 响应解析 cache_read_input_tokens"""
|
||||
from agentkit.llm.providers.anthropic import AnthropicProvider
|
||||
|
||||
provider = AnthropicProvider.__new__(AnthropicProvider)
|
||||
provider.api_key = "test"
|
||||
provider.base_url = "http://test"
|
||||
# 模拟 Anthropic 响应含 cache 字段
|
||||
mock_response = {
|
||||
"content": [{"type": "text", "text": "hello"}],
|
||||
"model": "claude-sonnet-4",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 50,
|
||||
"cache_read_input_tokens": 80,
|
||||
"cache_creation_input_tokens": 20,
|
||||
},
|
||||
}
|
||||
# 直接调用 _parse_response（私有方法，测试用）
|
||||
usage = provider._parse_response(mock_response, "claude-sonnet-4").usage
|
||||
assert usage.cache_read_input_tokens == 80
|
||||
assert usage.cache_creation_input_tokens == 20
|
||||
assert usage.cache_hit is True
|
||||
|
||||
|
||||
def _build_cache_metric_gateway(response):
|
||||
"""构造 LLMGateway mock，跳过 __init__。
|
||||
|
||||
ponytail: 直接装配测试所需的最小属性，避免初始化真实 provider/cache。
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from agentkit.llm.gateway import LLMGateway
|
||||
|
||||
mock_provider = MagicMock()
|
||||
mock_provider.chat = AsyncMock(return_value=response)
|
||||
|
||||
gw = LLMGateway.__new__(LLMGateway)
|
||||
gw._providers = {"mock": mock_provider}
|
||||
gw._config = MagicMock()
|
||||
gw._config.providers = {} # _calculate_cost 真实迭代返回 0.0
|
||||
gw._config.fallbacks = {} # _get_models_to_try 返回 [resolved_model]
|
||||
gw._cache_manager = None
|
||||
gw._usage_tracker = MagicMock()
|
||||
gw._resolve_model_alias = lambda m: m
|
||||
gw._resolve_model = lambda m: (mock_provider, m)
|
||||
return gw
|
||||
|
||||
|
||||
def test_gateway_emits_prompt_cache_hit_metric_on_cache_read():
|
||||
"""U2: gateway.chat 在 cache_read_input_tokens > 0 时 emit prompt_cache.hit"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from agentkit.llm.protocol import LLMResponse, TokenUsage
|
||||
|
||||
response = LLMResponse(
|
||||
content="test",
|
||||
model="claude",
|
||||
usage=TokenUsage(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=50,
|
||||
cache_read_input_tokens=80,
|
||||
),
|
||||
)
|
||||
|
||||
gw = _build_cache_metric_gateway(response)
|
||||
|
||||
# patch gateway 模块内的导入绑定名（from ... import prompt_cache_hit_counter）
|
||||
with (
|
||||
patch("agentkit.llm.gateway.prompt_cache_hit_counter") as hit_counter,
|
||||
patch("agentkit.llm.gateway.prompt_cache_miss_counter") as miss_counter,
|
||||
patch(
|
||||
"agentkit.llm.gateway.LLMGateway._record_usage",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
import asyncio
|
||||
|
||||
asyncio.run(gw.chat(messages=[{"role": "user", "content": "hi"}], model="claude"))
|
||||
|
||||
# cache_read_input_tokens > 0 → hit counter called
|
||||
hit_counter().add.assert_called_once()
|
||||
miss_counter().add.assert_not_called()
|
||||
|
||||
|
||||
def test_gateway_emits_prompt_cache_miss_metric_on_no_cache_read():
|
||||
"""U2: gateway.chat 在 cache_read_input_tokens == 0 时 emit prompt_cache.miss"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from agentkit.llm.protocol import LLMResponse, TokenUsage
|
||||
|
||||
response = LLMResponse(
|
||||
content="test",
|
||||
model="claude",
|
||||
usage=TokenUsage(prompt_tokens=100, completion_tokens=50), # cache_read=0
|
||||
)
|
||||
|
||||
gw = _build_cache_metric_gateway(response)
|
||||
|
||||
with (
|
||||
patch("agentkit.llm.gateway.prompt_cache_hit_counter") as hit_counter,
|
||||
patch("agentkit.llm.gateway.prompt_cache_miss_counter") as miss_counter,
|
||||
patch(
|
||||
"agentkit.llm.gateway.LLMGateway._record_usage",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
import asyncio
|
||||
|
||||
asyncio.run(gw.chat(messages=[{"role": "user", "content": "hi"}], model="claude"))
|
||||
|
||||
# cache_read_input_tokens == 0 → miss counter called
|
||||
miss_counter().add.assert_called_once()
|
||||
hit_counter().add.assert_not_called()
|
||||
|
|
|
|||
Loading…
Reference in New Issue