feat: add RetryPolicy for automatic RETRY failure mode handling

Author: baranylcn

gateframe/__init__.py

@@ -2,6 +2,7 @@
 from gateframe.core.contract import ValidationContract, ValidationResult
 from gateframe.core.escalation import EscalationRoute, EscalationRouter
 from gateframe.core.failure import FailureMode, FailureResult
+from gateframe.core.retry import RetryPolicy, RetryResult
 from gateframe.rules.boundary import AllowedValues, BoundaryRule
 from gateframe.rules.confidence import ConfidenceRule
 from gateframe.rules.semantic import LlmJudge, SemanticRule
@@ -18,6 +19,8 @@
     "LlmJudge",
     "SemanticRule",
     "StructuralRule",
+    "RetryPolicy",
+    "RetryResult",
     "ValidationContract",
     "ValidationResult",
     "WorkflowContext",

gateframe/core/retry.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from gateframe.core.contract import ValidationContract, ValidationResult
+from gateframe.core.failure import FailureMode
+
+
+@dataclass
+class RetryResult:
+    """Outcome of a RetryPolicy.run() call.
+
+    Attributes:
+        final: The last ValidationResult produced (pass or exhausted retries).
+        attempts: Total number of validate calls made (1 = no retry needed).
+        succeeded: True if any attempt passed validation.
+    """
+
+    final: ValidationResult
+    attempts: int
+    succeeded: bool
+    history: list[ValidationResult] = field(default_factory=list)
+
+
+class RetryPolicy:
+    """Executes a ValidationContract with automatic retry on RETRY failure mode.
+
+    The caller supplies a ``prompt_fn`` callable that returns a fresh output
+    on each invocation (e.g. re-calls the LLM).  On each attempt, if the
+    result contains a RETRY failure, ``prompt_fn`` is called again with the
+    previous failures injected as ``retry_failures`` context so the caller
+    can craft a follow-up prompt.
+
+    Args:
+        max_retries: Maximum number of *additional* attempts after the first.
+            Total attempts = max_retries + 1.
+
+    Example::
+
+        policy = RetryPolicy(max_retries=2)
+        result = policy.run(
+            contract,
+            prompt_fn=lambda **ctx: call_llm(ctx.get("retry_failures")),
+        )
+        if result.succeeded:
+            use(result.final)
+    """
+
+    def __init__(self, max_retries: int = 3) -> None:
+        if max_retries < 0:
+            raise ValueError("max_retries must be >= 0")
+        self.max_retries = max_retries
+
+    def run(
+        self,
+        contract: ValidationContract,
+        prompt_fn: Callable[..., Any],
+        **context: Any,
+    ) -> RetryResult:
+        """Run *contract* against output from *prompt_fn*, retrying on RETRY failures.
+
+        Args:
+            contract: The ValidationContract to evaluate.
+            prompt_fn: Called each attempt to produce the output to validate.
+                Receives ``**context`` plus ``retry_failures`` (list of
+                FailureResult) on retry attempts so callers can adjust prompts.
+            **context: Extra keyword arguments forwarded to both ``prompt_fn``
+                and ``contract.validate()``.
+
+        Returns:
+            A :class:`RetryResult` with the final result and attempt metadata.
+        """
+        history: list[ValidationResult] = []
+        retry_failures: list = []
+
+        for attempt in range(self.max_retries + 1):
+            call_context = dict(context)
+            if retry_failures:
+                call_context["retry_failures"] = retry_failures
+
+            output = prompt_fn(**call_context)
+            result = contract.validate(output, **call_context)
+            history.append(result)
+
+            if result.passed:
+                return RetryResult(
+                    final=result,
+                    attempts=attempt + 1,
+                    succeeded=True,
+                    history=history,
+                )
+
+            has_retry = any(f.failure_mode is FailureMode.RETRY for f in result.failures)
+            if not has_retry or attempt == self.max_retries:
+                break
+
+            retry_failures = [f for f in result.failures if f.failure_mode is FailureMode.RETRY]
+
+        return RetryResult(
+            final=history[-1],
+            attempts=len(history),
+            succeeded=False,
+            history=history,
+        )

tests/core/test_context.py

@@ -194,7 +194,6 @@ def test_reset_restores_initial_confidence_and_clears_history(self) -> None:
         assert ctx.history == []
         assert ctx.threshold_breached is False
 
-
     def test_concurrent_updates_do_not_corrupt_confidence(self) -> None:
         ctx = WorkflowContext("wf_concurrent", soft_fail_penalty=0.01)
         result = _make_result(passed=False, failures=[_soft_failure()])

tests/core/test_retry.py

@@ -0,0 +1,129 @@
+from typing import Any
+
+import pytest
+
+from gateframe.core.contract import ValidationContract
+from gateframe.core.failure import FailureMode, FailureResult
+from gateframe.core.retry import RetryPolicy
+from gateframe.core.rule import Rule
+
+
+class _PassingRule(Rule):
+    def validate(self, output: Any, **context: Any) -> FailureResult | None:
+        return None
+
+
+class _FailingRule(Rule):
+    def __init__(self, name: str = "failing", mode: FailureMode = FailureMode.RETRY) -> None:
+        super().__init__(name)
+        self._mode = mode
+
+    def validate(self, output: Any, **context: Any) -> FailureResult | None:
+        return FailureResult(rule_name=self.name, failure_mode=self._mode, message="rejected")
+
+
+class _PassOnAttempt(Rule):
+    """Passes on the Nth call (1-indexed)."""
+
+    def __init__(self, pass_on: int) -> None:
+        super().__init__("pass_on_attempt")
+        self._pass_on = pass_on
+        self._calls = 0
+
+    def validate(self, output: Any, **context: Any) -> FailureResult | None:
+        self._calls += 1
+        if self._calls >= self._pass_on:
+            return None
+        return FailureResult(rule_name=self.name, failure_mode=FailureMode.RETRY, message="not yet")
+
+
+class TestRetryPolicy:
+    def test_passes_on_first_attempt(self) -> None:
+        contract = ValidationContract("c", [_PassingRule("r")])
+        policy = RetryPolicy(max_retries=2)
+        result = policy.run(contract, prompt_fn=lambda **ctx: {})
+        assert result.succeeded is True
+        assert result.attempts == 1
+
+    def test_retries_until_pass(self) -> None:
+        rule = _PassOnAttempt(pass_on=3)
+        contract = ValidationContract("c", [rule])
+        policy = RetryPolicy(max_retries=3)
+        result = policy.run(contract, prompt_fn=lambda **ctx: {})
+        assert result.succeeded is True
+        assert result.attempts == 3
+
+    def test_exhausts_retries_and_fails(self) -> None:
+        contract = ValidationContract("c", [_FailingRule(mode=FailureMode.RETRY)])
+        policy = RetryPolicy(max_retries=2)
+        result = policy.run(contract, prompt_fn=lambda **ctx: {})
+        assert result.succeeded is False
+        assert result.attempts == 3  # 1 initial + 2 retries
+
+    def test_does_not_retry_on_hard_fail(self) -> None:
+        contract = ValidationContract("c", [_FailingRule(mode=FailureMode.HARD_FAIL)])
+        policy = RetryPolicy(max_retries=3)
+        result = policy.run(contract, prompt_fn=lambda **ctx: {})
+        assert result.succeeded is False
+        assert result.attempts == 1
+
+    def test_does_not_retry_on_soft_fail(self) -> None:
+        contract = ValidationContract("c", [_FailingRule(mode=FailureMode.SOFT_FAIL)])
+        policy = RetryPolicy(max_retries=3)
+        result = policy.run(contract, prompt_fn=lambda **ctx: {})
+        assert result.succeeded is False
+        assert result.attempts == 1
+
+    def test_retry_failures_passed_to_prompt_fn(self) -> None:
+        received: list = []
+
+        def prompt_fn(**ctx: Any) -> dict:
+            received.append(ctx.get("retry_failures"))
+            return {}
+
+        contract = ValidationContract("c", [_FailingRule(mode=FailureMode.RETRY)])
+        policy = RetryPolicy(max_retries=1)
+        policy.run(contract, prompt_fn=prompt_fn)
+
+        assert received[0] is None  # first call has no retry context
+        assert received[1] is not None  # second call has failures injected
+        assert received[1][0].failure_mode is FailureMode.RETRY
+
+    def test_context_forwarded_to_contract(self) -> None:
+        class _ContextRule(Rule):
+            def validate(self, output: Any, **context: Any) -> FailureResult | None:
+                if context.get("role") != "admin":
+                    return FailureResult(
+                        rule_name=self.name,
+                        failure_mode=FailureMode.HARD_FAIL,
+                        message="forbidden",
+                    )
+                return None
+
+        contract = ValidationContract("c", [_ContextRule("role_check")])
+        policy = RetryPolicy()
+        result = policy.run(contract, prompt_fn=lambda **ctx: {}, role="admin")
+        assert result.succeeded is True
+
+    def test_history_contains_all_attempts(self) -> None:
+        contract = ValidationContract("c", [_FailingRule(mode=FailureMode.RETRY)])
+        policy = RetryPolicy(max_retries=2)
+        result = policy.run(contract, prompt_fn=lambda **ctx: {})
+        assert len(result.history) == 3
+
+    def test_max_retries_zero_means_single_attempt(self) -> None:
+        contract = ValidationContract("c", [_FailingRule(mode=FailureMode.RETRY)])
+        policy = RetryPolicy(max_retries=0)
+        result = policy.run(contract, prompt_fn=lambda **ctx: {})
+        assert result.attempts == 1
+        assert result.succeeded is False
+
+    def test_negative_max_retries_raises(self) -> None:
+        with pytest.raises(ValueError, match="max_retries"):
+            RetryPolicy(max_retries=-1)
+
+    def test_retry_result_final_is_last_result(self) -> None:
+        contract = ValidationContract("c", [_FailingRule(mode=FailureMode.RETRY)])
+        policy = RetryPolicy(max_retries=1)
+        result = policy.run(contract, prompt_fn=lambda **ctx: {})
+        assert result.final is result.history[-1]