Copilot commented on code in PR #65485: URL: https://github.com/apache/airflow/pull/65485#discussion_r3106673485
########## providers/common/ai/src/airflow/providers/common/ai/privacy/config.py: ########## @@ -0,0 +1,61 @@ +# 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. +"""Configuration for LLM input guarding.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal + +from airflow.providers.common.ai.privacy.exceptions import InputGuardError + +InputGuardMode = Literal["redact", "mask", "replace", "hash"] +AttachmentPolicy = Literal["reject", "allow_unmodified"] + + +@dataclass(slots=True) +class InputGuardConfig: + """Runtime configuration for outbound LLM input sanitization.""" + + enabled: bool = False + language: str = "en" + entities: tuple[str, ...] | None = None + mode: InputGuardMode = "replace" + replace_value_template: str = "<{entity_type}>" + mask_char: str = "*" + mask_from_end: bool = False + score_threshold: float | None = None + attachment_policy: AttachmentPolicy = "reject" + log_sanitized_text: bool = False + log_preview_chars: int = 200 + + @classmethod + def from_value(cls, value: InputGuardConfig | dict[str, Any] | None) -> InputGuardConfig: + """Normalize supported config inputs into an InputGuardConfig.""" + if value is None: + return cls() + if isinstance(value, cls): + return value + if not isinstance(value, dict): + raise InputGuardError("input_guard must be a dict or InputGuardConfig instance.") + try: + entities = value.get("entities") + if entities is not None: + value = {**value, "entities": tuple(entities)} Review Comment: `InputGuardConfig.from_value()` converts `entities` via `tuple(entities)` without guarding against a single string input. If a user passes `{"entities": "EMAIL_ADDRESS"}`, this becomes a tuple of characters and will silently misconfigure Presidio. Consider explicitly handling `str` (either wrap as a 1-tuple or raise `InputGuardError`) and validating that the iterable contains strings. ########## 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: `sanitize_file_content()` does a local import of `InputGuardConfig` even though there’s no apparent circular dependency here. Moving this import to module scope would simplify the function and aligns with the project guideline to avoid imports inside function bodies unless necessary. ########## 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]: """ Review Comment: `PydanticAIHook.create_agent()` now accepts an `input_guard` kwarg, but the overload signatures above don’t include it. This makes type checkers/IDEs think `input_guard` is invalid. Update the overloads to include the new keyword-only parameter (and corresponding return types) so the public API remains accurately typed. ########## 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( + request.user_content, self.input_guard, logger=self.log + ) Review Comment: With `input_guard` enabled, `LLMFileAnalysisOperator` sanitizes `request.user_content` via `sanitize_file_content()`, and then `_create_agent()` forwards the same `input_guard` to `PydanticAIHook.create_agent()`, which adds a history processor that will sanitize the prompt content again before the model call. This double-sanitization adds extra Presidio passes (and potentially duplicate INFO logs). Consider doing only the early binary-attachment rejection here and relying on the hook’s history processor for text sanitization, or avoid passing `input_guard` to the agent when content has already been sanitized. -- 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]
