This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 787e323ca46 Support custom redaction and message length cap in
LLMRetryPolicy (#70830)
787e323ca46 is described below
commit 787e323ca462841a0c098b2be494c20c4b585642
Author: Wei Lee <[email protected]>
AuthorDate: Fri Aug 7 22:57:50 2026 +0800
Support custom redaction and message length cap in LLMRetryPolicy (#70830)
---
providers/common/ai/docs/retry_policies.rst | 74 +++++++++++--
.../airflow/providers/common/ai/policies/retry.py | 79 +++++++++----
.../ai/tests/unit/common/ai/policies/test_retry.py | 122 +++++++++++++++++++++
3 files changed, 244 insertions(+), 31 deletions(-)
diff --git a/providers/common/ai/docs/retry_policies.rst
b/providers/common/ai/docs/retry_policies.rst
index b498050b685..c1962e41ce8 100644
--- a/providers/common/ai/docs/retry_policies.rst
+++ b/providers/common/ai/docs/retry_policies.rst
@@ -71,8 +71,9 @@ How it works
When a task fails, ``LLMRetryPolicy``:
1. Sends the exception message to the configured LLM. By default, the message
- is first masked through Airflow's secrets masker (see ``redact_exception``
- below) before it is added to the prompt.
+ is first masked through Airflow's secrets masker (see ``redactor`` below)
+ and truncated to ``max_exception_length`` characters before it is added
+ to the prompt.
2. The LLM classifies the error into a category (``rate_limit``, ``auth``,
``network``, ``data``, ``transient``, ``permanent``)
3. Based on the classification, returns RETRY (with a suggested delay) or FAIL
@@ -158,22 +159,71 @@ Parameters
* - ``timeout``
- 30.0
- Max seconds to wait for the LLM response before falling back.
+ * - ``redactor``
+ - None (uses ``redact_registered_secrets``)
+ - Callable ``(str) -> str`` applied to the exception's string
+ representation before it is added to the classification prompt. The
+ default only masks values already registered via ``mask_secret()``
+ (e.g. connection passwords Airflow captured while resolving the
+ failing task's connections) -- it is not general-purpose PII
+ detection and will not catch arbitrary sensitive strings that were
+ never registered as secrets. Passing a custom callable **replaces**
+ the default masker entirely rather than stacking on top of it.
* - ``redact_exception``
- True
- - When ``True``, the exception's string representation is passed through
- Airflow's secrets masker before being added to the classification
- prompt. This only masks values already registered via
- ``mask_secret()`` (e.g. connection passwords Airflow captured while
- resolving the failing task's connections) -- it is not general-purpose
- PII detection and will not catch arbitrary sensitive strings that were
- never registered as secrets. Set to ``False`` only if you are certain
- your exception messages contain no sensitive data and you need the
- raw text for accurate classification.
+ - Whether to redact the exception's string representation before it is
+ added to the classification prompt. Set to ``False`` to disable
+ redaction entirely. Raises ``ValueError`` at construction time if
+ combined with an explicit ``redactor``.
+ * - ``max_exception_length``
+ - 4096
+ - Maximum number of characters of the (already redacted) exception
+ message included in the prompt. Longer messages are truncated with a
+ trailing ``"... (truncated)"`` marker. Must be a positive integer.
+
+Custom redactors
+----------------
+
+The default ``redactor`` only masks values already registered with Airflow's
+secrets masker via ``mask_secret()``. It does not detect free-text PII --
+email addresses, customer names, account numbers -- that were never
+registered as secrets. If your task's exception messages can contain that
+kind of data, supply your own ``redactor`` callable. It **replaces** the
+default masker rather than running in addition to it, so combine your own
+logic with
:func:`~airflow.providers.common.ai.policies.retry.redact_registered_secrets`
+yourself if you still want known-secret masking too:
+
+.. code-block:: python
+
+ import re
+
+ from airflow.providers.common.ai.policies.retry import
redact_registered_secrets
+
+ EMAIL_RE = re.compile(r"[\w.+-]+@[\w-]+\.[\w.-]+")
+
+
+ def redact_emails_and_secrets(message: str) -> str:
+ return redact_registered_secrets(EMAIL_RE.sub("<email>", message))
+
+
+ llm_policy = LLMRetryPolicy(
+ llm_conn_id="pydanticai_default",
+ redactor=redact_emails_and_secrets,
+ max_exception_length=2048, # keep long tracebacks from inflating
token cost
+ )
+
+To disable redaction entirely (for example, if you are certain your
+exception messages contain no sensitive data and need the raw text for
+accurate classification), pass ``redact_exception=False``:
+
+.. code-block:: python
+
+ LLMRetryPolicy(llm_conn_id="pydanticai_default", redact_exception=False)
Local LLM support
-----------------
-By default, ``redact_exception`` already masks known secrets before the
+By default, the built-in ``redactor`` already masks known secrets before the
exception data reaches the LLM provider. For environments where exception
data must not leave your own infrastructure at all -- even in masked form --
point to a local model via Ollama or vLLM instead, so the classification
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
index 768e16270e2..864e5b6dad4 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
@@ -24,7 +24,7 @@ from __future__ import annotations
import logging
from datetime import timedelta
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
from pydantic import BaseModel
@@ -43,12 +43,14 @@ except ImportError:
) from None
if TYPE_CHECKING:
+ from collections.abc import Callable
+
from airflow.sdk.definitions.context import Context
from airflow.sdk.definitions.retry_policy import RetryRule
log = logging.getLogger(__name__)
-__all__ = ["ErrorClassification", "LLMRetryPolicy"]
+__all__ = ["ErrorClassification", "LLMRetryPolicy",
"redact_registered_secrets"]
DEFAULT_INSTRUCTIONS = (
"You are an error classifier for a data pipeline system. "
@@ -79,6 +81,12 @@ class ErrorClassification(BaseModel):
"""Brief explanation of the classification decision."""
+def redact_registered_secrets(message: str) -> str:
+ """Mask values registered via ``mask_secret()``; the default ``redactor``
for :class:`LLMRetryPolicy`."""
+ # redact() is typed for arbitrary containers; a str in always yields a str
out.
+ return cast("str", redact(message))
+
+
class LLMRetryPolicy(RetryPolicy):
"""
Retry policy that uses an LLM to classify errors and decide retry
behaviour.
@@ -102,12 +110,24 @@ class LLMRetryPolicy(RetryPolicy):
falling back. Defaults to 30s. The LLM provider's own timeout
(e.g. 600s for Anthropic) is much longer; this keeps the retry
decision path fast even when the provider is degraded.
- :param redact_exception: When ``True`` (the default), the exception's
- string representation is passed through Airflow's secrets masker
- (:func:`~airflow.sdk.log.redact`) before being added to the prompt.
- Set to ``False`` only if you are certain your exception messages
- contain no sensitive data and you need the raw text for accurate
- classification.
+ :param redactor: Callable applied to the exception's string representation
+ before it is added to the classification prompt. Defaults to
+
:func:`~airflow.providers.common.ai.policies.retry.redact_registered_secrets`,
+ which only masks values already registered via ``mask_secret()``.
+ Pass a custom callable to replace the default masking entirely --
+ for example to redact free-text PII the secrets masker cannot see.
+ To disable masking altogether, use ``redact_exception=False`` --
+ not ``redactor=None``.
+ :param redact_exception: Whether to redact the exception's string
+ representation before it is added to the classification prompt.
+ Defaults to ``True``. Set to ``False`` to send the raw exception
+ text as-is. Passing ``redact_exception=False`` together with
+ an explicit ``redactor`` raises ``ValueError`` at construction time,
+ since the two settings would otherwise conflict silently.
+ :param max_exception_length: Maximum number of characters of the
+ (already redacted) exception message included in the prompt. Longer
+ messages are truncated with a trailing ``"... (truncated)"`` marker.
+ Must be a positive integer. Defaults to 4096.
.. warning::
The exception's string representation is sent to the configured
@@ -115,15 +135,18 @@ class LLMRetryPolicy(RetryPolicy):
etc.) as part of the classification prompt, so it may leak whatever
the failing task put in the exception message — connection strings,
credential fragments, PII, or other secrets. By default
- ``_classify()`` runs the message through Airflow's secrets masker
- (:func:`~airflow.sdk.log.redact`, controlled by ``redact_exception``),
- which masks values already registered via ``mask_secret()`` (for
- example, connection passwords Airflow captured while resolving the
- failing task's connections). This does **not** perform
- general-purpose PII detection and will not catch arbitrary sensitive
- strings that were never registered as secrets. You are still
- responsible for confirming that your task's exception messages are
- safe to send to a third-party LLM provider.
+ ``_classify()`` runs the message through
+
:func:`~airflow.providers.common.ai.policies.retry.redact_registered_secrets`
+ via ``redactor``, which masks values already registered via
+ ``mask_secret()`` (for example, connection passwords Airflow
+ captured while resolving the failing task's connections). This does
+ **not** perform general-purpose PII detection and will not catch
+ arbitrary sensitive strings that were never registered as secrets --
+ for free-text PII (emails, customer names, etc.) supply your own
+ ``redactor``, or pass ``redact_exception=False`` to disable
+ redaction altogether. You are still responsible for confirming that
+ your task's exception messages are safe to send to a third-party
+ LLM provider.
"""
def __init__(
@@ -134,14 +157,29 @@ class LLMRetryPolicy(RetryPolicy):
fallback_rules: list[RetryRule] | None = None,
timeout: float = 30.0,
*,
+ redactor: Callable[[str], str] | None = None,
redact_exception: bool = True,
+ max_exception_length: int = 4096,
) -> None:
+ if max_exception_length <= 0:
+ raise ValueError(f"max_exception_length must be a positive
integer, got {max_exception_length}")
+ if not redact_exception and redactor is not None:
+ raise ValueError(
+ "redactor must not be set when redact_exception=False --
passing an explicit "
+ "redactor while also disabling redaction is contradictory.
Either drop "
+ "redact_exception=False to keep using redactor, or drop
redactor to disable "
+ "redaction entirely."
+ )
self.llm_conn_id = llm_conn_id
self.model_id = model_id
self.instructions = instructions or DEFAULT_INSTRUCTIONS
self.fallback_rules = fallback_rules
self.timeout = timeout
+ self.redactor: Callable[[str], str] | None = (
+ None if not redact_exception else redactor if redactor is not None
else redact_registered_secrets
+ )
self.redact_exception = redact_exception
+ self.max_exception_length = max_exception_length
def evaluate(
self,
@@ -174,11 +212,14 @@ class LLMRetryPolicy(RetryPolicy):
instructions=self.instructions,
)
- exception_message = redact(str(exception)) if self.redact_exception
else str(exception)
+ # Redact before truncating -- truncating first could cut a registered
secret in half.
+ message = self.redactor(str(exception)) if self.redactor is not None
else str(exception)
+ if len(message) > self.max_exception_length:
+ message = f"{message[: self.max_exception_length]}... (truncated)"
prompt = (
f"Classify this error from a data pipeline task "
f"(attempt {try_number} of {max_tries}):\n\n"
- f"{type(exception).__name__}: {exception_message}"
+ f"{type(exception).__name__}: {message}"
)
from pydantic_ai.settings import ModelSettings
diff --git a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
index 6175938e07a..9b38eb243e1 100644
--- a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
+++ b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
@@ -28,6 +28,7 @@ pytest.importorskip("airflow.sdk.definitions.retry_policy",
reason="RetryPolicy
from airflow.providers.common.ai.policies.retry import (
ErrorClassification,
LLMRetryPolicy,
+ redact_registered_secrets,
)
from airflow.sdk._shared.secrets_masker import reset_secrets_masker
from airflow.sdk.definitions.retry_policy import RetryAction, RetryRule
@@ -48,6 +49,18 @@ def _make_mock_agent(category, should_retry, delay=0,
reasoning="test"):
return mock_agent
[email protected]_redact
+def test_redact_registered_secrets_masks_only_registered_values():
+ """Docs tell Dag authors to import and wrap this, so both the name and its
narrow scope are contracts."""
+ reset_secrets_masker()
+ mask_secret("super-secret-conn-password")
+
+ assert (
+ redact_registered_secrets("contact [email protected] with
super-secret-conn-password")
+ == "contact [email protected] with ***"
+ )
+
+
class TestLLMClassifyDecisions:
"""Test that _classify maps LLM classification to correct
RetryDecisions."""
@@ -154,6 +167,115 @@ class TestLLMClassifyDecisions:
prompt = mock_agent.run_sync.call_args[0][0]
assert secret_value in prompt
+ @pytest.mark.enable_redact
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ def test_explicit_redactor_none_still_applies_default_masking(self,
mock_hook_cls):
+ """redactor=None means "use the default masker" -- the same as
omitting it."""
+ reset_secrets_masker()
+ secret_value = "super-secret-conn-password"
+ mask_secret(secret_value)
+
+ mock_agent = _make_mock_agent("auth", should_retry=False)
+ mock_hook_cls.return_value.create_agent.return_value = mock_agent
+
+ policy = LLMRetryPolicy(llm_conn_id="test", redactor=None)
+ policy.evaluate(
+ ConnectionError(f"could not authenticate with password
{secret_value}"),
+ try_number=1,
+ max_tries=3,
+ )
+
+ prompt = mock_agent.run_sync.call_args[0][0]
+ assert secret_value not in prompt
+ assert "***" in prompt
+
+ def test_redact_exception_false_with_explicit_redactor_raises(self):
+ with pytest.raises(ValueError, match="redactor must not be set when
redact_exception=False"):
+ LLMRetryPolicy(llm_conn_id="test", redact_exception=False,
redactor=lambda message: message)
+
+ @pytest.mark.enable_redact
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ def test_custom_redactor_replaces_masker_instead_of_stacking(self,
mock_hook_cls):
+ """A custom redactor replaces the secrets masker entirely -- it is not
applied on top."""
+ reset_secrets_masker()
+ secret_value = "super-secret-conn-password"
+ mask_secret(secret_value)
+
+ mock_agent = _make_mock_agent("auth", should_retry=False)
+ mock_hook_cls.return_value.create_agent.return_value = mock_agent
+
+ policy = LLMRetryPolicy(llm_conn_id="test", redactor=lambda s:
s.replace("authenticate", "REDACTED"))
+ policy.evaluate(
+ ConnectionError(f"could not authenticate with password
{secret_value}"),
+ try_number=1,
+ max_tries=3,
+ )
+
+ prompt = mock_agent.run_sync.call_args[0][0]
+ # The registered secret is untouched by the masker...
+ assert secret_value in prompt
+ # ...but the custom redactor's own transformation did apply.
+ assert "REDACTED" in prompt
+
+ @pytest.mark.parametrize(
+ ("max_exception_length", "message_length", "expect_truncated"),
+ [
+ pytest.param(4096, 4096, False, id="default-limit-exact-fit"),
+ pytest.param(4096, 5000, True, id="default-limit-exceeded"),
+ pytest.param(10, 20, True, id="custom-limit-exceeded"),
+ pytest.param(10, 5, False, id="custom-limit-under"),
+ ],
+ )
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ def test_message_truncated_when_over_max_exception_length(
+ self, mock_hook_cls, max_exception_length, message_length,
expect_truncated
+ ):
+ mock_agent = _make_mock_agent("data", should_retry=False)
+ mock_hook_cls.return_value.create_agent.return_value = mock_agent
+
+ policy = LLMRetryPolicy(
+ llm_conn_id="test", redact_exception=False,
max_exception_length=max_exception_length
+ )
+ policy.evaluate(ValueError("x" * message_length), try_number=1,
max_tries=3)
+
+ prompt = mock_agent.run_sync.call_args[0][0]
+ assert ("... (truncated)" in prompt) is expect_truncated
+ if expect_truncated:
+ assert f"{'x' * max_exception_length}... (truncated)" in prompt
+ else:
+ assert "x" * message_length in prompt
+
+ @pytest.mark.enable_redact
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ def test_truncation_happens_after_redaction(self, mock_hook_cls):
+ """Redact-then-truncate must not equal truncate-then-redact for this
input.
+
+ The secret sits right at the truncation boundary: truncating first
would slice
+ it in half so the masker could no longer recognize and mask it.
+ """
+ reset_secrets_masker()
+ secret_value = "super-secret-conn-password"
+ mask_secret(secret_value)
+ max_exception_length = 20
+ # Padding places the secret so it straddles the truncation boundary.
+ padding = "a" * (max_exception_length - 5)
+ message = f"{padding}{secret_value}"
+
+ mock_agent = _make_mock_agent("auth", should_retry=False)
+ mock_hook_cls.return_value.create_agent.return_value = mock_agent
+
+ policy = LLMRetryPolicy(llm_conn_id="test",
max_exception_length=max_exception_length)
+ policy.evaluate(ConnectionError(message), try_number=1, max_tries=3)
+
+ prompt = mock_agent.run_sync.call_args[0][0]
+ assert secret_value not in prompt
+ assert "***" in prompt
+
+ @pytest.mark.parametrize("max_exception_length", [0, -1, -100])
+ def test_non_positive_max_exception_length_raises(self,
max_exception_length):
+ with pytest.raises(ValueError, match="max_exception_length must be a
positive integer"):
+ LLMRetryPolicy(llm_conn_id="test",
max_exception_length=max_exception_length)
+
@patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
def test_custom_instructions_forwarded_to_agent(self, mock_hook_cls):
mock_hook_cls.return_value.create_agent.return_value =
_make_mock_agent("x", False)
