kaxil commented on code in PR #65485: URL: https://github.com/apache/airflow/pull/65485#discussion_r3114169828
########## providers/common/ai/src/airflow/providers/common/ai/privacy/history.py: ########## @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Pydantic AI history processors backed by Presidio.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Any + +from pydantic_ai.messages import BinaryContent + +from airflow.providers.common.ai.privacy.exceptions import ATTACHMENT_REJECTION_MSG, InputGuardAttachmentError +from airflow.providers.common.ai.privacy.presidio_guard import PresidioInputGuard + +if TYPE_CHECKING: + import logging + + from airflow.providers.common.ai.privacy.config import InputGuardConfig + + +def sanitize_user_content( + content: Any, + *, + guard: PresidioInputGuard, + config: InputGuardConfig, + source: str = "content", +) -> Any: + """Sanitize supported user-content payloads while respecting attachment policy.""" + if isinstance(content, BinaryContent): + if config.attachment_policy == "reject": Review Comment: This rejects `BinaryContent` regardless of origin. When `agent.run_sync` is used with a model that returns a `FilePart` (model generated an image or document, `FilePart.content` is `BinaryContent`), the next turn's history processor calls into this function and raises `InputGuardAttachmentError` even though the user never attached anything. The feature is documented as guarding "user input", so maybe only `UserPromptPart` content should go through the reject check. A separate knob for model-output attachments would be safer. ########## providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py: ########## @@ -180,15 +183,35 @@ def create_agent( def create_agent(self, *, instructions: str, **agent_kwargs) -> Agent[None, str]: ... def create_agent( - self, output_type: type[Any] = str, *, instructions: str, **agent_kwargs + self, + output_type: type[Any] = str, + *, + instructions: str, + input_guard: InputGuardConfig | dict[str, Any] | None = None, + **agent_kwargs, ) -> Agent[None, Any]: """ Create a pydantic-ai Agent configured with this hook's model. :param output_type: The expected output type from the agent (default: ``str``). :param instructions: System-level instructions for the agent. + :param input_guard: Optional outbound input-sanitization configuration. :param agent_kwargs: Additional keyword arguments passed to the Agent constructor. """ + config = InputGuardConfig.from_value(input_guard) + if config.enabled: + instructions = PresidioInputGuard(config, logger=self.log).sanitize_text( Review Comment: This only sanitizes the static `instructions` string passed to `Agent(...)`. Pydantic-ai merges three sources into the final `ModelRequest.instructions`: (a) static instructions, (b) `@agent.instructions` callables, and (c) `AbstractToolset.get_instructions()`. Only (a) is sanitized here. Because `ModelRequest.instructions` is a top-level attribute (not in `message.parts`), the history processor in `history.py` never inspects it either. So anyone adding dynamic instructions, or using a toolset that injects instructions (SQL toolset, MCP, etc.), gets unsanitized text in every request. Consider sanitizing at a `before_model_request` hook on `request_context.request.instructions` so all three sources are covered. ########## providers/common/ai/src/airflow/providers/common/ai/privacy/presidio_guard.py: ########## @@ -0,0 +1,162 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Presidio-backed sanitization helpers for model-bound text.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Mapping, Sequence +from functools import lru_cache +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +if TYPE_CHECKING: + import logging + + from airflow.providers.common.ai.privacy.config import InputGuardConfig + +# Masks up to this many characters per detected entity, effectively masking all practical PII. +_MAX_MASK_CHARS = 10_000 + + +@lru_cache(maxsize=1) +def _create_presidio_engines() -> tuple[Any, Any]: + """Create and cache shared Presidio engines for the provider process.""" + try: + from presidio_analyzer import AnalyzerEngine + from presidio_anonymizer import AnonymizerEngine + except ImportError as exc: + raise AirflowOptionalProviderFeatureException( + "input_guard requires the presidio extra: " + "pip install apache-airflow-providers-common-ai[presidio]" + ) from exc + + return AnalyzerEngine(), AnonymizerEngine() + + +class PresidioInputGuard: + """Sanitize outbound strings and JSON-like payloads with Presidio.""" + + def __init__(self, config: InputGuardConfig, logger: logging.Logger | None = None) -> None: + self.config = config + self.logger = logger + self._analyzer, self._anonymizer = _create_presidio_engines() + + def sanitize_text(self, text: str, *, source: str = "text") -> str: + """Analyze and anonymize one text blob.""" + if not text: + return text + + analyze_kwargs: dict[str, Any] = { + "text": text, + "language": self.config.language, + } + if self.config.entities: Review Comment: `if self.config.entities:` treats `entities=()` as falsy and drops it, which tells Presidio to analyze everything. Users who write `entities=()` expecting "detect nothing" get the opposite. Distinguishing `None` (no filter) from `()` (empty filter, skip analysis) would match least-surprise semantics. ########## providers/common/ai/src/airflow/providers/common/ai/operators/llm.py: ########## @@ -114,10 +117,24 @@ def llm_hook(self) -> PydanticAIHook: } return PydanticAIHook.get_hook(self.llm_conn_id, hook_params=hook_params) - def execute(self, context: Context) -> Any: - agent: Agent[None, Any] = self.llm_hook.create_agent( - output_type=self.output_type, instructions=self.system_prompt, **self.agent_params + def _create_agent( + self, + *, + output_type: type[Any] | None = None, + instructions: str | None = None, + **agent_kwargs: Any, + ) -> Agent[None, Any]: + """Create the pydantic-ai Agent with optional outbound input guarding.""" + if self.input_guard is not None: + agent_kwargs["input_guard"] = self.input_guard + return self.llm_hook.create_agent( + output_type=output_type or self.output_type, Review Comment: `output_type or self.output_type` is a subtle trap. It works today because classes are truthy, but if a caller ever passes a falsy value, the instance attribute silently takes over. `output_type if output_type is not None else self.output_type` reads clearer and matches how `instructions` is handled on the next line. ########## providers/common/ai/src/airflow/providers/common/ai/privacy/presidio_guard.py: ########## @@ -0,0 +1,162 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Presidio-backed sanitization helpers for model-bound text.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Mapping, Sequence +from functools import lru_cache +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +if TYPE_CHECKING: + import logging + + from airflow.providers.common.ai.privacy.config import InputGuardConfig + +# Masks up to this many characters per detected entity, effectively masking all practical PII. +_MAX_MASK_CHARS = 10_000 + + +@lru_cache(maxsize=1) Review Comment: `@lru_cache(maxsize=1)` keyed on `()` means the first call wins for the whole process. If the first task uses `language="en"` and a second task configures `language="fr"`, the second task gets the English engine (initialized with `supported_languages=None`, which is en only) and `analyzer.analyze(language="fr")` either errors or returns nothing. The cache should key on `language` (or the full relevant config), not on `()`. ########## providers/common/ai/src/airflow/providers/common/ai/privacy/history.py: ########## @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Pydantic AI history processors backed by Presidio.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Any + +from pydantic_ai.messages import BinaryContent + +from airflow.providers.common.ai.privacy.exceptions import ATTACHMENT_REJECTION_MSG, InputGuardAttachmentError +from airflow.providers.common.ai.privacy.presidio_guard import PresidioInputGuard + +if TYPE_CHECKING: + import logging + + from airflow.providers.common.ai.privacy.config import InputGuardConfig + + +def sanitize_user_content( + content: Any, + *, + guard: PresidioInputGuard, + config: InputGuardConfig, + source: str = "content", +) -> Any: + """Sanitize supported user-content payloads while respecting attachment policy.""" + if isinstance(content, BinaryContent): + if config.attachment_policy == "reject": + raise InputGuardAttachmentError(ATTACHMENT_REJECTION_MSG) + return content + if isinstance(content, str): + return guard.sanitize_text(content, source=source) + if isinstance(content, list): + return [ + sanitize_user_content(item, guard=guard, config=config, source=f"{source}[{i}]") + for i, item in enumerate(content) + ] + if isinstance(content, tuple): + return tuple( + sanitize_user_content(item, guard=guard, config=config, source=f"{source}[{i}]") + for i, item in enumerate(content) + ) + if isinstance(content, dict): + return { + key: sanitize_user_content(value, guard=guard, config=config, source=f"{source}.{key}") + for key, value in content.items() + } + # Non-text, non-container types (int, float, …) pass through unchanged. + return content Review Comment: Pydantic `BaseModel` tool returns bypass this function. `ToolReturnPart.content` is typed `Any`, and tools often return `BaseModel` instances like `Customer(email="[email protected]")`. It isn't `str`/`list`/`tuple`/`dict`, so it falls through here unchanged. Pydantic-ai then calls `part.model_response_str()` (which JSON-dumps it) when building the model request, so the PII reaches the model. Handling `BaseModel` via `.model_dump()` and dataclasses via `asdict` would close this. Verified in a minimal repro: a tool returning `Customer(email="[email protected]")` arrives at the model as `{"email":"[email protected]"}`. ########## providers/common/ai/tests/unit/common/ai/privacy/test_presidio_guard.py: ########## @@ -0,0 +1,494 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import sys +from dataclasses import dataclass +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from airflow.providers.common.ai.privacy.config import InputGuardConfig +from airflow.providers.common.ai.privacy.exceptions import ( + ATTACHMENT_REJECTION_MSG, + InputGuardAttachmentError, + InputGuardError, +) +from airflow.providers.common.ai.privacy.history import ( + build_history_processor, + contains_binary_content, + sanitize_file_content, + sanitize_user_content, +) +from airflow.providers.common.ai.privacy.presidio_guard import PresidioInputGuard, _create_presidio_engines + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_engines(analyzer_results=None, anonymized_text="<REDACTED>"): + """Return (analyzer_mock, anonymizer_mock) pre-configured for a single analyze call.""" + analyzer = MagicMock() Review Comment: All tests mock `AnalyzerEngine` and `AnonymizerEngine`, so no test exercises actual Presidio behaviour. The mode-specific outputs (what `replace` / `mask` / `redact` / `hash` actually produce), the `KeyError` from a bad template, and the missing-language-model case are invisible to CI. One end-to-end test guarded by `pytest.importorskip("presidio_analyzer")` would catch real regressions without inflating CI cost much. ########## providers/common/ai/src/airflow/providers/common/ai/privacy/presidio_guard.py: ########## @@ -0,0 +1,162 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Presidio-backed sanitization helpers for model-bound text.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Mapping, Sequence +from functools import lru_cache +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +if TYPE_CHECKING: + import logging + + from airflow.providers.common.ai.privacy.config import InputGuardConfig + +# Masks up to this many characters per detected entity, effectively masking all practical PII. +_MAX_MASK_CHARS = 10_000 + + +@lru_cache(maxsize=1) +def _create_presidio_engines() -> tuple[Any, Any]: + """Create and cache shared Presidio engines for the provider process.""" + try: + from presidio_analyzer import AnalyzerEngine + from presidio_anonymizer import AnonymizerEngine + except ImportError as exc: + raise AirflowOptionalProviderFeatureException( + "input_guard requires the presidio extra: " + "pip install apache-airflow-providers-common-ai[presidio]" + ) from exc + + return AnalyzerEngine(), AnonymizerEngine() + + +class PresidioInputGuard: + """Sanitize outbound strings and JSON-like payloads with Presidio.""" + + def __init__(self, config: InputGuardConfig, logger: logging.Logger | None = None) -> None: + self.config = config + self.logger = logger + self._analyzer, self._anonymizer = _create_presidio_engines() + + def sanitize_text(self, text: str, *, source: str = "text") -> str: + """Analyze and anonymize one text blob.""" + if not text: + return text + + analyze_kwargs: dict[str, Any] = { + "text": text, + "language": self.config.language, + } + if self.config.entities: + analyze_kwargs["entities"] = list(self.config.entities) + if self.config.score_threshold is not None: + analyze_kwargs["score_threshold"] = self.config.score_threshold + + results = self._analyzer.analyze(**analyze_kwargs) + if not results: + return text + + sanitized_text = self._anonymizer.anonymize( + text=text, + analyzer_results=results, + operators=self._build_operators(results), + ).text + self._log_sanitization(results=results, sanitized_text=sanitized_text, source=source) + return sanitized_text + + def sanitize_value(self, value: Any, *, source: str = "value") -> Any: + """Recursively sanitize text in JSON-like tool payloads.""" + if isinstance(value, str): + return self.sanitize_text(value, source=source) + if isinstance(value, Mapping): + return { + key: self.sanitize_value(subvalue, source=f"{source}.{key}") + for key, subvalue in value.items() + } + if isinstance(value, list): + return [ + self.sanitize_value(item, source=f"{source}[{index}]") for index, item in enumerate(value) + ] + if isinstance(value, tuple): + return tuple( + self.sanitize_value(item, source=f"{source}[{index}]") for index, item in enumerate(value) + ) + return value + + def _build_operators(self, results: Sequence[Any]) -> dict[str, Any]: + from presidio_anonymizer.entities import OperatorConfig + + if self.config.mode == "redact": + return {"DEFAULT": OperatorConfig("redact", {})} + if self.config.mode == "hash": + return {"DEFAULT": OperatorConfig("hash", {})} + if self.config.mode == "mask": + return { + "DEFAULT": OperatorConfig( + "mask", + { + "masking_char": self.config.mask_char, + "chars_to_mask": _MAX_MASK_CHARS, + "from_end": self.config.mask_from_end, + }, + ) + } + + operators: dict[str, Any] = {} + for result in results: + entity_type = getattr(result, "entity_type", "PII") + operators[entity_type] = OperatorConfig( + "replace", + {"new_value": self.config.replace_value_template.format(entity_type=entity_type)}, Review Comment: `replace_value_template` is user-supplied and goes through `str.format(entity_type=...)` with no validation. A typo like `"<{entity}>"` raises `KeyError: 'entity'` at sanitize time, after the expensive `analyze` call has already run. Templates like `"{0}"` or `"{__class__}"` work but read unintended attributes. Validating the template once in `InputGuardConfig.from_value` (render with a dummy `entity_type`, reject on `KeyError`) would catch this at DAG-parse time. ########## providers/common/ai/src/airflow/providers/common/ai/privacy/history.py: ########## @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Pydantic AI history processors backed by Presidio.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Any + +from pydantic_ai.messages import BinaryContent + +from airflow.providers.common.ai.privacy.exceptions import ATTACHMENT_REJECTION_MSG, InputGuardAttachmentError +from airflow.providers.common.ai.privacy.presidio_guard import PresidioInputGuard + +if TYPE_CHECKING: + import logging + + from airflow.providers.common.ai.privacy.config import InputGuardConfig + + +def sanitize_user_content( + content: Any, + *, + guard: PresidioInputGuard, + config: InputGuardConfig, + source: str = "content", +) -> Any: + """Sanitize supported user-content payloads while respecting attachment policy.""" + if isinstance(content, BinaryContent): + if config.attachment_policy == "reject": + raise InputGuardAttachmentError(ATTACHMENT_REJECTION_MSG) + return content + if isinstance(content, str): + return guard.sanitize_text(content, source=source) + if isinstance(content, list): + return [ + sanitize_user_content(item, guard=guard, config=config, source=f"{source}[{i}]") + for i, item in enumerate(content) + ] + if isinstance(content, tuple): + return tuple( + sanitize_user_content(item, guard=guard, config=config, source=f"{source}[{i}]") + for i, item in enumerate(content) + ) + if isinstance(content, dict): + return { + key: sanitize_user_content(value, guard=guard, config=config, source=f"{source}.{key}") + for key, value in content.items() + } + # Non-text, non-container types (int, float, …) pass through unchanged. + return content + + +def contains_binary_content(content: Any) -> bool: + """Return True when user content includes at least one binary attachment.""" + if isinstance(content, BinaryContent): + return True + if isinstance(content, list | tuple): + return any(contains_binary_content(item) for item in content) + if isinstance(content, dict): + return any(contains_binary_content(value) for value in content.values()) + return False + + +def sanitize_file_content( + user_content: Any, + input_guard: InputGuardConfig | dict[str, Any] | None, + logger: logging.Logger | None = None, +) -> Any: + """ + Sanitize file-analysis content using the given guard configuration. + + Performs an early binary-attachment check before invoking the guard so that + multimodal attachments are rejected before any LLM call is made. + + :param user_content: The prompt/content payload built from the analysed files. + :param input_guard: Guard configuration as an ``InputGuardConfig`` or raw dict. + :param logger: Logger to use for sanitization summaries. + :return: Sanitized content, or the original content when guarding is disabled. + """ + from airflow.providers.common.ai.privacy.config import InputGuardConfig as _Config + + config = _Config.from_value(input_guard) + if not config.enabled: + return user_content + if config.attachment_policy == "reject" and contains_binary_content(user_content): + raise InputGuardAttachmentError(ATTACHMENT_REJECTION_MSG) + guard = PresidioInputGuard(config, logger=logger) + return sanitize_user_content( + user_content, guard=guard, config=config, source="file_analysis.user_content" + ) + + +def build_history_processor(config: InputGuardConfig, logger: logging.Logger | None = None): + """Build a Pydantic AI history processor that sanitizes outbound content.""" + guard = PresidioInputGuard(config, logger=logger) + + def _processor(messages: list[Any]) -> list[Any]: + sanitized_messages: list[Any] = [] + for message_index, message in enumerate(messages): + parts = getattr(message, "parts", None) + if parts is None: + sanitized_messages.append(message) + continue + + new_parts = [] + changed = False + for part_index, part in enumerate(parts): + content = getattr(part, "content", None) Review Comment: Tool call args bypass the guard entirely. Real `pydantic_ai.messages.ToolCallPart` has no `content` field, just `args: str | dict[str, Any] | None`, so this `getattr` returns None and the part falls into the `continue` branch on line 126. On the next turn the LLM sees the raw tool args it emitted earlier, which routinely include PII (e.g. `send_email(to="[email protected]", ...)`). I verified this against pydantic-ai 1.80 with a `TestModel` + real tool. The `test_part_without_content_attr_is_kept` test uses a fake `ToolCallPart` dataclass that doesn't have `args`, so it confirms the bypass is intentional but hides the production impact. Same gap for `BuiltinToolCallPart.args`. Consider extending the processor to also sanitize `part.args` for these parts. ########## providers/common/ai/src/airflow/providers/common/ai/privacy/presidio_guard.py: ########## @@ -0,0 +1,162 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Presidio-backed sanitization helpers for model-bound text.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Mapping, Sequence +from functools import lru_cache +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +if TYPE_CHECKING: + import logging + + from airflow.providers.common.ai.privacy.config import InputGuardConfig + +# Masks up to this many characters per detected entity, effectively masking all practical PII. +_MAX_MASK_CHARS = 10_000 + + +@lru_cache(maxsize=1) +def _create_presidio_engines() -> tuple[Any, Any]: + """Create and cache shared Presidio engines for the provider process.""" + try: + from presidio_analyzer import AnalyzerEngine + from presidio_anonymizer import AnonymizerEngine + except ImportError as exc: + raise AirflowOptionalProviderFeatureException( + "input_guard requires the presidio extra: " + "pip install apache-airflow-providers-common-ai[presidio]" + ) from exc + + return AnalyzerEngine(), AnonymizerEngine() + + +class PresidioInputGuard: + """Sanitize outbound strings and JSON-like payloads with Presidio.""" + + def __init__(self, config: InputGuardConfig, logger: logging.Logger | None = None) -> None: + self.config = config + self.logger = logger + self._analyzer, self._anonymizer = _create_presidio_engines() + + def sanitize_text(self, text: str, *, source: str = "text") -> str: + """Analyze and anonymize one text blob.""" + if not text: + return text + + analyze_kwargs: dict[str, Any] = { + "text": text, + "language": self.config.language, + } + if self.config.entities: + analyze_kwargs["entities"] = list(self.config.entities) + if self.config.score_threshold is not None: + analyze_kwargs["score_threshold"] = self.config.score_threshold + + results = self._analyzer.analyze(**analyze_kwargs) + if not results: + return text + + sanitized_text = self._anonymizer.anonymize( + text=text, + analyzer_results=results, + operators=self._build_operators(results), + ).text + self._log_sanitization(results=results, sanitized_text=sanitized_text, source=source) + return sanitized_text + + def sanitize_value(self, value: Any, *, source: str = "value") -> Any: Review Comment: `sanitize_value` isn't called anywhere in the PR. `history.py::sanitize_user_content` covers the same ground (with slightly different semantics, `Mapping` vs `dict`). Either wire this up or drop it, two parallel implementations are a maintenance hazard. ########## providers/common/ai/src/airflow/providers/common/ai/privacy/history.py: ########## @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Pydantic AI history processors backed by Presidio.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Any + +from pydantic_ai.messages import BinaryContent + +from airflow.providers.common.ai.privacy.exceptions import ATTACHMENT_REJECTION_MSG, InputGuardAttachmentError +from airflow.providers.common.ai.privacy.presidio_guard import PresidioInputGuard + +if TYPE_CHECKING: + import logging + + from airflow.providers.common.ai.privacy.config import InputGuardConfig + + +def sanitize_user_content( + content: Any, + *, + guard: PresidioInputGuard, + config: InputGuardConfig, + source: str = "content", +) -> Any: + """Sanitize supported user-content payloads while respecting attachment policy.""" + if isinstance(content, BinaryContent): + if config.attachment_policy == "reject": + raise InputGuardAttachmentError(ATTACHMENT_REJECTION_MSG) + return content + if isinstance(content, str): + return guard.sanitize_text(content, source=source) + if isinstance(content, list): + return [ + sanitize_user_content(item, guard=guard, config=config, source=f"{source}[{i}]") + for i, item in enumerate(content) + ] + if isinstance(content, tuple): + return tuple( + sanitize_user_content(item, guard=guard, config=config, source=f"{source}[{i}]") + for i, item in enumerate(content) + ) + if isinstance(content, dict): + return { + key: sanitize_user_content(value, guard=guard, config=config, source=f"{source}.{key}") + for key, value in content.items() + } + # Non-text, non-container types (int, float, …) pass through unchanged. + return content + + +def contains_binary_content(content: Any) -> bool: + """Return True when user content includes at least one binary attachment.""" + if isinstance(content, BinaryContent): + return True + if isinstance(content, list | tuple): + return any(contains_binary_content(item) for item in content) + if isinstance(content, dict): + return any(contains_binary_content(value) for value in content.values()) + return False + + +def sanitize_file_content( + user_content: Any, + input_guard: InputGuardConfig | dict[str, Any] | None, + logger: logging.Logger | None = None, +) -> Any: + """ + Sanitize file-analysis content using the given guard configuration. + + Performs an early binary-attachment check before invoking the guard so that + multimodal attachments are rejected before any LLM call is made. + + :param user_content: The prompt/content payload built from the analysed files. + :param input_guard: Guard configuration as an ``InputGuardConfig`` or raw dict. + :param logger: Logger to use for sanitization summaries. + :return: Sanitized content, or the original content when guarding is disabled. + """ + from airflow.providers.common.ai.privacy.config import InputGuardConfig as _Config Review Comment: Inline import. There's no circular dependency here, `config.py` is already referenced at the top via `TYPE_CHECKING`, and importing it eagerly is fine. Moving this to the module-level imports matches the provider convention and the global "imports at top" rule. ########## providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py: ########## @@ -116,6 +115,10 @@ def execute(self, context: Context) -> Any: max_text_chars=self.max_text_chars, sample_rows=self.sample_rows, ) + if self.input_guard is not None: + request.user_content = sanitize_file_content( Review Comment: This calls the guard once on `request.user_content`, then `_create_agent` wires up the history processor which runs Presidio a second time on the same text once `agent.run_sync(request.user_content)` fires. For `replace` mode it's idempotent but wasteful. For `hash` mode it isn't idempotent: a hex hash can re-match `CREDIT_CARD` / `US_BANK_NUMBER` regexes and get re-hashed, producing a different final value than documented. And `_log_sanitization` emits two INFO lines for the same content. The early binary-reject check is the only piece that has to run before `run_sync`; the rest could be left to the history processor. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
