mirror of
https://github.com/hwchase17/langchain.git
synced 2026-06-09 18:50:33 +00:00
d08245f70d
`PIIMiddleware` previously scrubbed detected PII only at the state level via its `after_model` / `before_model` hooks. Consumers reading the live stream — `astream_events(version="v3")` or `run.messages` / `run.tool_calls` / `run.values` — saw the raw model text, the raw tool-call args, the raw tool outputs, and the raw state snapshots until the run finished and the canonical conversation history was written. This change registers a stream transformer ahead of `MessagesTransformer` that redacts every wire surface of an agent run. The transformer holds a sliding lookback buffer (default 128 characters) per `(run_id, content-block index)` so PII patterns that straddle delta boundaries are caught before the safe prefix is released downstream. Anything older than the lookback is run through the configured detector and emitted; the trailing tail stays buffered until a later delta extends it past the cap or the block finishes. `_finalize_block` always re-runs detection over the full block snapshot so the finalized content lands fully redacted even when the in-flight buffer never released a tail (short responses, or PII arriving in the final delta). The `block` strategy is now supported on the streaming path via a buffering mode that withholds every delta until the block resolves — clean blocks release the full text at finalize, PII-bearing blocks zero the wire and let `after_model` / `apply_to_tool_results` raise `PIIDetectionError` on the original state message. Activation is gated on `apply_to_output=True`, matching the existing post-hoc semantics. The middleware's transformer factory is cloned by `StreamMux._make_child` into every subgraph scope, so attaching `PIIMiddleware` at the outer agent also redacts streamed deltas from sub-agents invoked inside tools. ## Tool-call and tools-channel coverage The transformer covers every wire surface of an agent run, not just AI message text: - **Streamed AI text deltas** (`content-block-delta` of type `text-delta`) — lookback machinery, redacted in place. - **Streamed tool-call args** (`content-block-delta` with `tool_call_chunk` / `server_tool_call_chunk` fields) — each delta carries the full cumulative args string; detection runs on the field directly and redacts in place. Verified empirically against `_compat_bridge.py` and the consumer-side `_merge_block_delta_into_store` snapshot-replace semantics. - **Finalized tool-call blocks** (`content-block-finish` with `tool_call` / `server_tool_call` / `invalid_tool_call`) — `args` dict walked recursively and each string leaf redacted. - **Tool execution events on the `tools` channel** — `tool-started.input`, `tool-output-delta`, `tool-finished.output`, `tool-error.message` all run through detection. String deltas use the same lookback machinery as text-deltas keyed by `tool_call_id`; structured payloads walk recursively. - **State snapshots on the `values` channel** — message lists are walked and each message's `.content` is redacted on a fresh copy. Graph state itself stays intact for the state-level enforcer (`apply_to_tool_results` via `before_model`) to act on independently. - **Legacy `(BaseMessage, metadata)` payloads** on the `messages` channel (Python 3.10 path, where `langgraph`'s `ASYNCIO_ACCEPTS_CONTEXT = sys.version_info >= (3, 11)` falls back to a code path that doesn't propagate the streaming callback into the chat model) — `.content` and `AIMessage.tool_calls[*].args` are scrubbed. For `block`, the event's `data` tuple is replaced with an empty-content copy so the original message stays in state for `after_model` to raise on. ## Worth a careful look - `_PIIStreamTransformer._mutate_text_delta` — lookback partition. Anything older than `lookback` characters is released after redaction; the tail stays buffered. Bulletproof against whitespace-permissive detectors (notably `credit_card`, whose regex matches across spaces). - `_PIIStreamTransformer._mutate_tool_call_chunk_delta` — direct in-place redaction of the cumulative args string. No buffer; the wire shape is cumulative-snapshot, the consumer-side merge is replace-not-append. - `_PIIStreamTransformer._mutate_legacy_payload` — the dual path: mutate-in-place for non-`block` (idempotent with `after_model`), replace-with-empty-copy for `block` (keeps original in graph state for `after_model` to raise on). - `_PIIStreamTransformer._redact_value` — the recursive walker. `BaseMessage` branch returns a fresh `.content`-redacted copy via `model_copy(update=...)` — never mutates in place — so tool-output payloads that wrap a `ToolMessage` and message lists in state snapshots flow through cleanly. - The new `transformers` attribute on `PIIMiddleware`: this is what makes `create_agent` pick the factory up. Multiple `PIIMiddleware` instances each register one transformer; ordering is preserved within the `before_builtins` lane. ## Compatibility Bumps `langgraph` to `>=1.2.1` for the `before_builtins` opt-in on `StreamTransformer`.