"""PII detection and handling middleware for agents.""" from __future__ import annotations from functools import partial from typing import TYPE_CHECKING, Any, ClassVar, Literal from langchain_core.messages import AIMessage, AnyMessage, BaseMessage, HumanMessage, ToolMessage from langgraph.stream import StreamTransformer from typing_extensions import override from langchain.agents.middleware._redaction import ( PIIDetectionError, PIIMatch, RedactionRule, ResolvedRedactionRule, apply_strategy, detect_credit_card, detect_email, detect_ip, detect_mac_address, detect_url, ) from langchain.agents.middleware.types import ( AgentMiddleware, AgentState, ContextT, ResponseT, hook_config, ) if TYPE_CHECKING: from collections.abc import Callable from langgraph.runtime import Runtime from langgraph.stream._types import ProtocolEvent _DEFAULT_STREAM_LOOKBACK = 128 """Default trailing-buffer size for cross-delta PII detection. The transformer always holds the last `lookback` characters in a per-content block buffer so that PII patterns straddling delta boundaries are detected before any text is released downstream. 128 comfortably covers the built-in detectors (the credit-card regex tops out at 19 characters; URLs and emails are typically well under 100) while bounding first-token latency. """ class _PIIStreamTransformer(StreamTransformer): """Mutates `content-block-delta` text on `messages` events in flight. Runs before built-in stream transformers so the redacted text is what every downstream consumer sees — both the main protocol event log and the `run.messages` projection that `MessagesTransformer` snapshots into. Holds a sliding buffer of the most recent text per (run_id, content block index) so PII patterns that straddle delta boundaries are caught. Anything older than `lookback` characters is redacted with the resolved rule's strategy and emitted as the new delta text; the trailing tail stays in the buffer until a later delta extends it past the cap or the block's finish event flushes the snapshot. """ before_builtins: ClassVar[bool] = True required_stream_modes: ClassVar[tuple[str, ...]] = ("messages", "tools", "values") def __init__( self, scope: tuple[str, ...] = (), *, rule: ResolvedRedactionRule, lookback: int = _DEFAULT_STREAM_LOOKBACK, ) -> None: super().__init__(scope) self._rule = rule self._lookback = lookback # Text/reasoning deltas keyed by `(run_id, content_block_index)`. self._buffers: dict[tuple[str, int], str] = {} # Tool-output-delta buffers keyed by `tool_call_id`. Held in a # separate dict so `_drop_run` on the messages channel can't # sweep active tool-output state. self._tool_buffers: dict[str, str] = {} def init(self) -> dict[str, Any]: # No projection — this transformer mutates events in place rather # than building a derived view. return {} def process(self, event: ProtocolEvent) -> bool: method = event["method"] if method == "messages": return self._process_messages_event(event) if method == "tools": return self._process_tools_event(event) if method == "values": return self._process_values_event(event) return True def _process_values_event(self, event: ProtocolEvent) -> bool: """Redact the state snapshot on the `values` channel. State snapshots emitted between nodes carry the full state dict, which typically includes the messages list. Walking the snapshot with `_redact_value` returns a fresh structure where every message has a redacted copy of its content — the original objects in graph state remain intact for the state-level enforcer (`apply_to_tool_results` via `before_model`) to act on independently when the agent loops back. """ data = event["params"].get("data") if data is None: return True event["params"]["data"] = self._redact_value(data) return True def _process_messages_event(self, event: ProtocolEvent) -> bool: params = event["params"] data = params.get("data") if not isinstance(data, tuple) or len(data) != 2: # noqa: PLR2004 return True payload, metadata = data # Legacy `(BaseMessage, metadata)` shape: the langgraph→langchain # integration emits this when a model only implements `_generate` # (or when its `_astream` falls back), producing a single event # carrying the full message rather than streamed content-block # deltas. Swap in a redacted copy so the consumer sees scrubbed # text on the wire while the original stays intact in graph state # for `after_model` to act on independently. Under `block`, # `_redact_base_message` raises `PIIDetectionError` via # `apply_strategy` before we get here. if isinstance(payload, BaseMessage): redacted = self._redact_base_message(payload) if redacted is not payload: params["data"] = (redacted, metadata) return True if not isinstance(payload, dict): return True kind = payload.get("event") run_id = str(metadata.get("run_id") or "") if metadata else "" if kind == "content-block-delta": self._mutate_delta(payload, run_id) elif kind == "content-block-finish": self._finalize_block(payload, run_id) elif kind in {"message-finish", "error"}: self._drop_run(run_id) return True def _process_tools_event(self, event: ProtocolEvent) -> bool: data = event["params"].get("data") if not isinstance(data, dict): return True kind = data.get("event") tool_call_id = data.get("tool_call_id") if kind == "tool-started": # Tool inputs may be a dict (multi-arg tools), a string # (single-arg tools — `BaseTool._parse_input` passes the # raw string through), or a list (array-input tools). # `_redact_value` handles all three uniformly. if "input" in data: data["input"] = self._redact_value(data["input"]) elif kind == "tool-output-delta": # Use the tool_call_id as buffer key when present; fall back # to a None-keyed slot for the rare malformed/custom emitter # case (the buffer becomes shared but at least redaction runs). self._mutate_tool_output_delta( data, tool_call_id if isinstance(tool_call_id, str) else "" ) elif kind == "tool-finished": if "output" in data: data["output"] = self._redact_value(data["output"]) if isinstance(tool_call_id, str): self._tool_buffers.pop(tool_call_id, None) elif kind == "tool-error": msg = data.get("message") if isinstance(msg, str) and msg: matches = self._rule.detector(msg) if matches: data["message"] = apply_strategy(msg, matches, self._rule.strategy) if isinstance(tool_call_id, str): self._tool_buffers.pop(tool_call_id, None) return True def _mutate_tool_output_delta(self, data: dict[str, Any], tool_call_id: str) -> None: """Redact a `tool-output-delta` payload. String deltas go through the same lookback machinery as text-deltas, keyed by `tool_call_id` in the disjoint `_tool_buffers` dict so `_drop_run` on the messages channel can't sweep active tool-output state. Structured deltas (dict/list) walk recursively without buffering — they don't have a position-stable shape across deltas to buffer against. """ delta = data.get("delta") if isinstance(delta, str): held = self._tool_buffers.get(tool_call_id, "") combined = held + delta matches = self._rule.detector(combined) if matches: # `apply_strategy` raises `PIIDetectionError` under # `strategy="block"`, failing the run immediately — # cleaner than withholding deltas until `after_model` # raises later. combined = apply_strategy(combined, matches, self._rule.strategy) emit_end = max(0, len(combined) - self._lookback) self._tool_buffers[tool_call_id] = combined[emit_end:] data["delta"] = combined[:emit_end] elif isinstance(delta, (dict, list)): data["delta"] = self._redact_value(delta) def _redact_tool_call_list(self, calls: list[Any] | None) -> tuple[list[Any], bool]: """Walk a list of tool-call (or invalid-tool-call) dicts. Returns `(new_list, changed)`. Each element's `args` is run through `_redact_value` regardless of its type — `tool_call.args` is a dict, `invalid_tool_call.args` is a raw JSON string, and `_redact_value` handles both shapes uniformly. If nothing changed, returns the input list and `changed=False`. """ if not calls: return calls or [], False new_calls: list[Any] = [] changed = False for tc in calls: if isinstance(tc, dict) and "args" in tc and tc["args"] is not None: redacted = self._redact_value(tc["args"]) if redacted != tc["args"]: new_tc = dict(tc) new_tc["args"] = redacted new_calls.append(new_tc) changed = True continue new_calls.append(tc) return new_calls, changed def _redact_value(self, value: Any) -> Any: """Recursively redact PII in string leaves of a nested structure. Returns a new value where every `str` leaf that contains PII has been replaced (or emptied under `block`). Non-string leaves and the structure itself are preserved. `BaseMessage` payloads (typically `ToolMessage` from `tool-finished.output`, or any message reached via the `values` channel) return a fresh copy with `.content` redacted plus `AIMessage.tool_calls[*].args` / `invalid_tool_calls[*].args` walked. The original object stays intact for state-level enforcers (`after_model`, `before_model` with `apply_to_tool_results`) to act on independently. Scope mirrors the pre-streaming state-level surfaces: `.content` (string or list-of-content-blocks) and `tool_calls` args. Other message attributes (`additional_kwargs`, `response_metadata`, `ToolMessage.artifact`) are intentionally not walked here — they aren't scrubbed in graph state by the existing hooks, so scrubbing them on the wire would create a wire/state divergence. """ if isinstance(value, str): if not value: return value matches = self._rule.detector(value) if not matches: return value # `apply_strategy` raises `PIIDetectionError` under `block` # — the run fails immediately rather than buffering until a # state-level hook can raise. return apply_strategy(value, matches, self._rule.strategy) if isinstance(value, BaseMessage): return self._redact_base_message(value) if isinstance(value, dict): return {k: self._redact_value(v) for k, v in value.items()} if isinstance(value, list): return [self._redact_value(v) for v in value] if isinstance(value, tuple): return tuple(self._redact_value(v) for v in value) return value def _redact_base_message(self, value: BaseMessage) -> BaseMessage: """Return a fresh copy of `value` with PII-carrying surfaces redacted.""" update: dict[str, Any] = {} content = value.content if isinstance(content, str) and content: matches = self._rule.detector(content) if matches: update["content"] = apply_strategy(content, matches, self._rule.strategy) elif isinstance(content, list) and content: # Structured content-blocks shape: # `[{"type": "text", "text": "..."}, {"type": "tool_call", ...}, ...]`. redacted_content = self._redact_value(content) if redacted_content != content: update["content"] = redacted_content # `AIMessage.tool_calls` and `.invalid_tool_calls` carry PII in # `args` independently of `.content`. `tool_call.args` is a # dict; `invalid_tool_call.args` is a raw JSON string — # `_redact_value` handles both shapes via the recursion. if isinstance(value, AIMessage): new_tc_list, tc_changed = self._redact_tool_call_list(value.tool_calls) if tc_changed: update["tool_calls"] = new_tc_list new_inv_list, inv_changed = self._redact_tool_call_list(value.invalid_tool_calls) if inv_changed: update["invalid_tool_calls"] = new_inv_list if not update: return value return value.model_copy(update=update) def _mutate_delta(self, payload: dict[str, Any], run_id: str) -> None: delta = payload.get("delta") if not isinstance(delta, dict): return delta_type = delta.get("type") if delta_type == "text-delta": self._mutate_string_field_delta(delta, payload, run_id, "text") return if delta_type == "reasoning-delta": # Reasoning content (chain-of-thought from extended-thinking # models) is a real PII surface — models echo back # user-supplied data or synthesize it from context. Run the # same lookback machinery as text-delta against the # `reasoning` field. Block indices are unique within a # message regardless of block type, so the buffer key # `(run_id, index)` naturally disjoint from text-delta keys. self._mutate_string_field_delta(delta, payload, run_id, "reasoning") return if delta_type == "block-delta": fields = delta.get("fields") if isinstance(fields, dict) and fields.get("type") in { "tool_call_chunk", "server_tool_call_chunk", }: self._mutate_tool_call_chunk_delta(fields) # Other delta types (`data-delta`, vendor block types) pass # through. The pre-streaming middleware scrubbed `.content` text # on state messages only; binary payloads and provider-specific # block shapes are out of scope for parity with that surface. def _mutate_string_field_delta( self, delta: dict[str, Any], payload: dict[str, Any], run_id: str, field: str, ) -> None: """Apply the lookback-buffer redaction to a string field on a delta. Shared by `text-delta` (`field="text"`) and `reasoning-delta` (`field="reasoning"`). Buffer is keyed by `(run_id, block_index)`; block indices are unique within a message so different block types share the same key space without collision. """ text = delta.get(field) if not isinstance(text, str) or not text: return index = payload.get("index") if not isinstance(index, int): return key = (run_id, index) held = self._buffers.get(key, "") combined = held + text # Run detection on the full accumulated buffer before splitting. # Detecting only on the about-to-emit prefix would miss matches # that straddle the lookback boundary — the detector's regex # needs a complete, boundary-anchored hit, so a truncated prefix # would fail to match and the partial PII would leak on the # wire. Under `strategy="block"`, `apply_strategy` raises # `PIIDetectionError` here, failing the run as soon as PII # arrives rather than buffering until `after_model`. matches = self._rule.detector(combined) if matches: combined = apply_strategy(combined, matches, self._rule.strategy) emit_end = max(0, len(combined) - self._lookback) self._buffers[key] = combined[emit_end:] delta[field] = combined[:emit_end] def _mutate_tool_call_chunk_delta(self, fields: dict[str, Any]) -> None: """Redact cumulative tool-call args with lookback withholding. Each `tool_call_chunk` `block-delta` event carries the full accumulated args string (verified against `_compat_bridge.py` — `delta_source = current` for these block types — and against the consumer-side `_merge_block_delta_into_store`, which replaces wholesale rather than appends). Detection runs on the full cumulative args so any complete PII anywhere in the string is redacted before emission. Lookback withholding then trims the trailing the lookback window characters from what reaches the consumer — those characters might be the start of a partial PII match that completes in a future cumulative delta. The trimmed tail surfaces at `content-block- finish` where `_finalize_block` redacts the parsed args dict. For args that fit within the lookback window (the typical case), this withholds the entire args string during streaming — the redacted args dict appears only at finalize. For args that exceed the lookback window, the safe prefix streams incrementally as the cumulative state grows. PII that appears more than the lookback window characters from the cumulative tail in a delta where it hasn't yet completed can still surface in the emit prefix — same residual exposure as PII longer than the lookback window on the text path. The `content-block-finish` snapshot redaction is the backstop. """ args = fields.get("args") if not isinstance(args, str) or not args: return matches = self._rule.detector(args) if matches: # `apply_strategy` raises `PIIDetectionError` under # `strategy="block"` — the run fails the moment a complete # PII pattern surfaces in the cumulative args string. args = apply_strategy(args, matches, self._rule.strategy) emit_end = max(0, len(args) - self._lookback) fields["args"] = args[:emit_end] def _finalize_block(self, payload: dict[str, Any], run_id: str) -> None: index = payload.get("index") if not isinstance(index, int): return key = (run_id, index) # The finalized block carries the model's original concatenation # of deltas, not what we emitted on the wire. Re-run detection over # its full text so the snapshot matches the redacted stream. content = payload.get("content") if isinstance(content, dict): ctype = content.get("type") if ctype == "text": self._finalize_string_field(content, "text") elif ctype == "reasoning": self._finalize_string_field(content, "reasoning") elif ( ctype in {"tool_call", "server_tool_call", "invalid_tool_call"} and "args" in content and content["args"] is not None ): # `tool_call` / `server_tool_call` args are dicts; # `invalid_tool_call.args` is the raw unparsed JSON # string. `_redact_value` handles both shapes. content["args"] = self._redact_value(content["args"]) self._buffers.pop(key, None) def _finalize_string_field(self, content: dict[str, Any], field: str) -> None: """Re-redact a string content-block field on `content-block-finish`. Used for `text` and `reasoning` content blocks. Under `strategy="block"` `apply_strategy` raises `PIIDetectionError`, failing the run immediately. """ text = content.get(field) if not isinstance(text, str) or not text: return matches = self._rule.detector(text) if not matches: return content[field] = apply_strategy(text, matches, self._rule.strategy) def _drop_run(self, run_id: str) -> None: # Release any buffered tails for this run_id — content-block-finish # should have already done so for normal completion, but message-finish # / error paths need an explicit sweep so abandoned blocks don't # accumulate in long-lived processes. stale = [key for key in self._buffers if key[0] == run_id] for key in stale: del self._buffers[key] def finalize(self) -> None: self._buffers.clear() self._tool_buffers.clear() def fail(self, err: BaseException) -> None: # noqa: ARG002 self._buffers.clear() self._tool_buffers.clear() class PIIMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]): """Detect and handle Personally Identifiable Information (PII) in conversations. This middleware detects common PII types and applies configurable strategies to handle them. It can detect emails, credit cards, IP addresses, MAC addresses, and URLs in both user input and agent output. Built-in PII types: - `email`: Email addresses - `credit_card`: Credit card numbers (validated with Luhn algorithm) - `ip`: IP addresses (validated with stdlib) - `mac_address`: MAC addresses - `url`: URLs (both `http`/`https` and bare URLs) Strategies: - `block`: Raise an exception when PII is detected - `redact`: Replace PII with `[REDACTED_TYPE]` placeholders - `mask`: Partially mask PII (e.g., `****-****-****-1234` for credit card) - `hash`: Replace PII with deterministic hash (e.g., ``) Strategy Selection Guide: | Strategy | Preserves Identity? | Best For | | -------- | ------------------- | --------------------------------------- | | `block` | N/A | Avoid PII completely | | `redact` | No | General compliance, log sanitization | | `mask` | No | Human readability, customer service UIs | | `hash` | Yes (pseudonymous) | Analytics, debugging | Example: ```python from langchain.agents.middleware import PIIMiddleware from langchain.agents import create_agent # Redact all emails in user input agent = create_agent( "openai:gpt-5.5", middleware=[ PIIMiddleware("email", strategy="redact"), ], ) # Use different strategies for different PII types agent = create_agent( "openai:gpt-5.5", middleware=[ PIIMiddleware("credit_card", strategy="mask"), PIIMiddleware("url", strategy="redact"), PIIMiddleware("ip", strategy="hash"), ], ) # Custom PII type with regex agent = create_agent( "openai:gpt-5.5", middleware=[ PIIMiddleware("api_key", detector=r"sk-[a-zA-Z0-9]{32}", strategy="block"), ], ) ``` """ def __init__( self, # From a typing point of view, the literals are covered by 'str'. # Nonetheless, we escape PYI051 to keep hints and autocompletion for the caller. pii_type: Literal["email", "credit_card", "ip", "mac_address", "url"] | str, # noqa: PYI051 *, strategy: Literal["block", "redact", "mask", "hash"] = "redact", detector: Callable[[str], list[PIIMatch]] | str | None = None, apply_to_input: bool = True, apply_to_output: bool = False, apply_to_tool_results: bool = False, ) -> None: """Initialize the PII detection middleware. Args: pii_type: Type of PII to detect. Can be a built-in type (`email`, `credit_card`, `ip`, `mac_address`, `url`) or a custom type name. strategy: How to handle detected PII. Options: * `block`: Raise `PIIDetectionError` when PII is detected * `redact`: Replace with `[REDACTED_TYPE]` placeholders * `mask`: Partially mask PII (show last few characters) * `hash`: Replace with deterministic hash (format: ``) detector: Custom detector function or regex pattern. * If `Callable`: Function that takes content string and returns list of `PIIMatch` objects * If `str`: Regex pattern to match PII * If `None`: Uses built-in detector for the `pii_type` apply_to_input: Whether to check user messages before model call. apply_to_output: Whether to check AI messages after model call. When `True`, a stream transformer is also installed so that every wire surface of an agent run is redacted in flight: * Streamed AI text deltas (`content-block-delta` of type `text-delta`) * Streamed tool-call arguments (`content-block-delta` with `tool_call_chunk` / `server_tool_call_chunk` fields, plus the finalized `tool_call` content block on `content-block-finish`) * Tool execution events on the `tools` channel (`tool-started.input`, `tool-output-delta`, `tool-finished.output`, `tool-error.message`) * State snapshots on the `values` channel — message lists are walked and each message's `.content` is redacted on a fresh copy (state itself stays intact for `before_model` / `after_model` to act on independently) State-level redaction via `after_model` (and `before_model` with `apply_to_tool_results`) remains the canonical enforcer; the streaming transformer ensures consumers reading `astream_events(version="v3")` or `run.messages` / `run.tool_calls` / `run.values` never see PII on the wire. apply_to_tool_results: Whether to check tool result messages after tool execution. Raises: ValueError: If `pii_type` is not built-in and no detector is provided. """ super().__init__() self.apply_to_input = apply_to_input self.apply_to_output = apply_to_output self.apply_to_tool_results = apply_to_tool_results self._resolved_rule: ResolvedRedactionRule = RedactionRule( pii_type=pii_type, strategy=strategy, detector=detector, ).resolve() self.pii_type = self._resolved_rule.pii_type self.strategy = self._resolved_rule.strategy self.detector = self._resolved_rule.detector # Stream transformer scrubs the streamed surface of the same # messages that the state-level hooks scrub in graph state. # Installed whenever any output-side scrubbing is enabled — # `apply_to_output` covers AI messages (text, tool-call args, # reasoning), `apply_to_tool_results` covers tool execution # (the `tools` channel + ToolMessage content on `values` and # `messages`). For `block` the transformer raises # `PIIDetectionError` directly from its event handler the # moment a complete PII pattern is detected, failing the run # via langgraph's `StreamMux.afail` path. The state-level # `after_model` / `before_model` hooks remain a backstop for # non-streaming consumers. if self.apply_to_output or self.apply_to_tool_results: self.transformers = ( partial( _PIIStreamTransformer, rule=self._resolved_rule, ), ) @property def name(self) -> str: """Name of the middleware.""" return f"{self.__class__.__name__}[{self.pii_type}]" def _process_content(self, content: str) -> tuple[str, list[PIIMatch]]: """Apply the configured redaction rule to the provided content.""" matches = self.detector(content) if not matches: return content, [] sanitized = apply_strategy(content, matches, self.strategy) return sanitized, matches @hook_config(can_jump_to=["end"]) @override def before_model( self, state: AgentState[Any], runtime: Runtime[ContextT], ) -> dict[str, Any] | None: """Check user messages and tool results for PII before model invocation. Args: state: The current agent state. runtime: The langgraph runtime. Returns: Updated state with PII handled according to strategy, or `None` if no PII detected. Raises: PIIDetectionError: If PII is detected and strategy is `'block'`. """ if not self.apply_to_input and not self.apply_to_tool_results: return None messages = state["messages"] if not messages: return None new_messages = list(messages) any_modified = False # Check user input if enabled if self.apply_to_input: # Get last user message last_user_msg = None last_user_idx = None for i in range(len(messages) - 1, -1, -1): if isinstance(messages[i], HumanMessage): last_user_msg = messages[i] last_user_idx = i break if last_user_idx is not None and last_user_msg and last_user_msg.content: # Detect PII in message content content = str(last_user_msg.content) new_content, matches = self._process_content(content) if matches: updated_message: AnyMessage = HumanMessage( content=new_content, id=last_user_msg.id, name=last_user_msg.name, ) new_messages[last_user_idx] = updated_message any_modified = True # Check tool results if enabled if self.apply_to_tool_results: # Find the last AIMessage, then process all `ToolMessage` objects after it last_ai_idx = None for i in range(len(messages) - 1, -1, -1): if isinstance(messages[i], AIMessage): last_ai_idx = i break if last_ai_idx is not None: # Get all tool messages after the last AI message for i in range(last_ai_idx + 1, len(messages)): msg = messages[i] if isinstance(msg, ToolMessage): tool_msg = msg if not tool_msg.content: continue content = str(tool_msg.content) new_content, matches = self._process_content(content) if not matches: continue # Create updated tool message updated_message = ToolMessage( content=new_content, id=tool_msg.id, name=tool_msg.name, tool_call_id=tool_msg.tool_call_id, ) new_messages[i] = updated_message any_modified = True if any_modified: return {"messages": new_messages} return None @hook_config(can_jump_to=["end"]) async def abefore_model( self, state: AgentState[Any], runtime: Runtime[ContextT], ) -> dict[str, Any] | None: """Async check user messages and tool results for PII before model invocation. Args: state: The current agent state. runtime: The langgraph runtime. Returns: Updated state with PII handled according to strategy, or `None` if no PII detected. Raises: PIIDetectionError: If PII is detected and strategy is `'block'`. """ return self.before_model(state, runtime) @override def after_model( self, state: AgentState[Any], runtime: Runtime[ContextT], ) -> dict[str, Any] | None: """Check AI messages for PII after model invocation. Args: state: The current agent state. runtime: The langgraph runtime. Returns: Updated state with PII handled according to strategy, or None if no PII detected. Raises: PIIDetectionError: If PII is detected and strategy is `'block'`. """ if not self.apply_to_output: return None messages = state["messages"] if not messages: return None # Get last AI message last_ai_msg = None last_ai_idx = None for i in range(len(messages) - 1, -1, -1): msg = messages[i] if isinstance(msg, AIMessage): last_ai_msg = msg last_ai_idx = i break if last_ai_idx is None or not last_ai_msg or not last_ai_msg.content: return None # Detect PII in message content content = str(last_ai_msg.content) new_content, matches = self._process_content(content) if not matches: return None # Create updated message updated_message = AIMessage( content=new_content, id=last_ai_msg.id, name=last_ai_msg.name, tool_calls=last_ai_msg.tool_calls, ) # Return updated messages new_messages = list(messages) new_messages[last_ai_idx] = updated_message return {"messages": new_messages} async def aafter_model( self, state: AgentState[Any], runtime: Runtime[ContextT], ) -> dict[str, Any] | None: """Async check AI messages for PII after model invocation. Args: state: The current agent state. runtime: The langgraph runtime. Returns: Updated state with PII handled according to strategy, or None if no PII detected. Raises: PIIDetectionError: If PII is detected and strategy is `'block'`. """ return self.after_model(state, runtime) __all__ = [ "PIIDetectionError", "PIIMatch", "PIIMiddleware", "detect_credit_card", "detect_email", "detect_ip", "detect_mac_address", "detect_url", ]