From 60a58a0fdb9e45381cfecce6ff1946fb68b8b8e3 Mon Sep 17 00:00:00 2001 From: Chiguyong Date: Mon, 6 Jul 2026 00:32:35 +0800 Subject: [PATCH] =?UTF-8?q?feat(llm):=20U2=20=E2=80=94=20Prompt=20Cache=20?= =?UTF-8?q?=E5=85=A8=E9=93=BE=E8=B7=AF=E7=9B=91=E6=8E=A7=20+=20PII=20?= =?UTF-8?q?=E8=84=B1=E6=95=8F=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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：脱敏失败拒绝发送（不降级到原文），等保合规要求 --- agentkit.yaml | 18 ++ src/agentkit/llm/gateway.py | 61 +++++-- src/agentkit/llm/pii_filter.py | 227 ++++++++++++++++++++++++ src/agentkit/llm/protocol.py | 16 +- src/agentkit/llm/providers/anthropic.py | 8 + src/agentkit/telemetry/metrics.py | 47 +++++ tests/unit/test_pii_filter.py | 200 +++++++++++++++++++++ tests/unit/test_prompt_cache_layers.py | 141 ++++++++++++++- 8 files changed, 697 insertions(+), 21 deletions(-) create mode 100644 src/agentkit/llm/pii_filter.py create mode 100644 tests/unit/test_pii_filter.py diff --git a/agentkit.yaml b/agentkit.yaml index 62c1411..7c4d53d 100644 --- a/agentkit.yaml +++ b/agentkit.yaml @@ -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 diff --git a/src/agentkit/llm/gateway.py b/src/agentkit/llm/gateway.py index 387b578..7496f43 100644 --- a/src/agentkit/llm/gateway.py +++ b/src/agentkit/llm/gateway.py @@ -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. diff --git a/src/agentkit/llm/pii_filter.py b/src/agentkit/llm/pii_filter.py new file mode 100644 index 0000000..70f1aea --- /dev/null +++ b/src/agentkit/llm/pii_filter.py @@ -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 diff --git a/src/agentkit/llm/protocol.py b/src/agentkit/llm/protocol.py index 21369c3..d08997c 100644 --- a/src/agentkit/llm/protocol.py +++ b/src/agentkit/llm/protocol.py @@ -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: diff --git a/src/agentkit/llm/providers/anthropic.py b/src/agentkit/llm/providers/anthropic.py index 3a1c197..a57ad4a 100644 --- a/src/agentkit/llm/providers/anthropic.py +++ b/src/agentkit/llm/providers/anthropic.py @@ -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 diff --git a/src/agentkit/telemetry/metrics.py b/src/agentkit/telemetry/metrics.py index 5c4303c..fde9b1f 100644 --- a/src/agentkit/telemetry/metrics.py +++ b/src/agentkit/telemetry/metrics.py @@ -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 diff --git a/tests/unit/test_pii_filter.py b/tests/unit/test_pii_filter.py new file mode 100644 index 0000000..ee4af34 --- /dev/null +++ b/tests/unit/test_pii_filter.py @@ -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 diff --git a/tests/unit/test_prompt_cache_layers.py b/tests/unit/test_prompt_cache_layers.py index d3e61f0..cf128a6 100644 --- a/tests/unit/test_prompt_cache_layers.py +++ b/tests/unit/test_prompt_cache_layers.py @@ -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()