gmf_forge_ai_orchestration.behaviors

25class BaseBehavior(ABC):
26    """
27    Composable agent behavior.
28
29    Behaviors are applied as a pipeline around every agent execution:
30    ``before_execute`` → agent runs → ``after_execute`` (or ``on_error``).
31
32    Subclasses override only the hooks they need.
33    """
34
35    # ------------------------------------------------------------------
36    # Hooks — override as needed
37    # ------------------------------------------------------------------
38
39    async def before_execute(self, context: BehaviorContext) -> BehaviorContext:
40        """
41        Called before the agent executes its task.
42
43        May mutate and return a modified context (e.g. to inject guardrails).
44        """
45        return context
46
47    async def after_execute(self, context: BehaviorContext, result: Any) -> Any:
48        """
49        Called after the agent successfully executes.
50
51        May mutate and return a modified result.
52        """
53        return result
54
55    async def on_error(
56        self, context: BehaviorContext, error: Exception
57    ) -> Optional[Any]:
58        """
59        Called when the agent raises an exception.
60
61        Return a fallback value to suppress the error, or re-raise / return
62        ``None`` to propagate it.
63        """
64        return None

13@dataclass
14class BehaviorContext:
15    """Execution context passed to every behavior hook."""
16
17    agent_id: str
18    task: str
19    attempt: int = 1
20    execution_id: str = ""
21    last_completed_step: int = -1
22    metadata: Dict[str, Any] = field(default_factory=dict)

 12class RetryBehavior(BaseBehavior):
 13    """
 14    Retries the agent on failure with exponential backoff.
 15
 16    Args:
 17        max_retries: Maximum number of retry attempts (default: 3).
 18        base_delay: Initial delay in seconds before the first retry (default: 1.0).
 19        backoff_factor: Multiplier applied to delay after each failure (default: 2.0).
 20        max_delay: Upper bound on delay between retries in seconds (default: 60.0).
 21        exceptions_to_retry: Tuple of exception types that trigger a retry.
 22            Defaults to ``(Exception,)`` — all exceptions.
 23
 24    Example::
 25
 26        agent = ReActAgent(
 27            llm_gateway=gateway,
 28            behaviors=[RetryBehavior(max_retries=3, base_delay=0.5)],
 29        )
 30    """
 31
 32    def __init__(
 33        self,
 34        max_retries: int = 3,
 35        base_delay: float = 1.0,
 36        backoff_factor: float = 2.0,
 37        max_delay: float = 60.0,
 38        exceptions_to_retry: Tuple[Type[Exception], ...] = (Exception,),
 39        logger: Optional[BasicLogger] = None,
 40    ) -> None:
 41        self.max_retries = max_retries
 42        self.base_delay = base_delay
 43        self.backoff_factor = backoff_factor
 44        self.max_delay = max_delay
 45        self.exceptions_to_retry = exceptions_to_retry
 46        self._logger = logger
 47
 48        # Tracks remaining retries across on_error calls for a single execution
 49        self._remaining: int = max_retries
 50
 51    async def before_execute(self, context: BehaviorContext) -> BehaviorContext:
 52        # Reset retry counter at the start of each execution
 53        self._remaining = self.max_retries
 54        return context
 55
 56    async def on_error(
 57        self, context: BehaviorContext, error: Exception
 58    ) -> Optional[Any]:
 59        if not isinstance(error, self.exceptions_to_retry):
 60            raise error
 61
 62        if self._remaining <= 0:
 63            if self._logger:
 64                self._logger.info(
 65                    "RetryBehavior: all retries exhausted — propagating error",
 66                    agent_id=context.agent_id,
 67                    max_retries=self.max_retries,
 68                    error=str(error),
 69                    error_type=type(error).__name__,
 70                )
 71            raise error
 72
 73        attempt = self.max_retries - self._remaining + 1
 74        delay = min(self.base_delay * (self.backoff_factor ** (attempt - 1)), self.max_delay)
 75        self._remaining -= 1
 76
 77        if self._logger:
 78            self._logger.info(
 79                "RetryBehavior: scheduling retry",
 80                agent_id=context.agent_id,
 81                attempt=attempt,
 82                retries_remaining=self._remaining,
 83                delay_s=round(delay, 3),
 84                error=str(error),
 85                error_type=type(error).__name__,
 86            )
 87
 88        await asyncio.sleep(delay)
 89
 90        # Prefer resume-from-checkpoint when execution metadata is available.
 91        resume_agent = context.metadata.get("resume_agent")
 92        if (
 93            context.execution_id
 94            and context.last_completed_step >= 0
 95            and resume_agent is not None
 96            and hasattr(resume_agent, "resume_from")
 97        ):
 98            if self._logger:
 99                self._logger.info(
100                    "RetryBehavior: resuming from checkpoint",
101                    agent_id=context.agent_id,
102                    execution_id=context.execution_id,
103                    last_completed_step=context.last_completed_step,
104                )
105            return await resume_agent.resume_from(context.execution_id)
106
107        return RETRY_SENTINEL

30class GuardrailBehavior(BaseBehavior):
31    """
32    Validates task inputs and agent outputs against configurable rules.
33
34    Raises :class:`GuardrailViolationError` if any rule is violated.
35
36    Args:
37        rules: List of :class:`GuardrailRule` to enforce.
38
39    Example::
40
41        behavior = GuardrailBehavior(rules=[
42            GuardrailRule(
43                name="no_pii",
44                blocked_words=["ssn", "social security"],
45                apply_to="both",
46            ),
47            GuardrailRule(name="max_input", max_length=4000, apply_to="input"),
48        ])
49    """
50
51    def __init__(self, rules: Optional[List[GuardrailRule]] = None) -> None:
52        self._rules = rules or []
53        self._compiled: List[tuple] = []
54        for rule in self._rules:
55            compiled_pattern: Optional[Pattern] = (
56                re.compile(rule.pattern, re.IGNORECASE) if rule.pattern else None
57            )
58            self._compiled.append((rule, compiled_pattern))
59
60    def _validate(self, text: str, target: str) -> None:
61        """Run all rules applicable to ``target`` (``"input"`` or ``"output"``)."""
62        text_lower = text.lower()
63        for rule, compiled_pattern in self._compiled:
64            if rule.apply_to not in (target, "both"):
65                continue
66
67            if rule.max_length and len(text) > rule.max_length:
68                raise GuardrailViolationError(
69                    rule.name,
70                    f"Length {len(text)} exceeds max {rule.max_length}",
71                )
72
73            for word in rule.blocked_words:
74                if word.lower() in text_lower:
75                    raise GuardrailViolationError(
76                        rule.name, f"Blocked term detected: '{word}'"
77                    )
78
79            if compiled_pattern and compiled_pattern.search(text):
80                raise GuardrailViolationError(
81                    rule.name, f"Pattern '{rule.pattern}' matched"
82                )
83
84    async def before_execute(self, context: BehaviorContext) -> BehaviorContext:
85        self._validate(context.task, "input")
86        return context
87
88    async def after_execute(self, context: BehaviorContext, result: Any) -> Any:
89        # Extract string output for validation
90        output_text = result.output if hasattr(result, "output") else str(result)
91        self._validate(output_text, "output")
92        return result

11@dataclass
12class GuardrailRule:
13    """A single validation rule."""
14
15    name: str
16    pattern: Optional[str] = None          # compiled as regex
17    blocked_words: List[str] = field(default_factory=list)
18    max_length: Optional[int] = None
19    apply_to: str = "both"                 # "input" | "output" | "both"

22class GuardrailViolationError(Exception):
23    """Raised when a guardrail rule is violated."""
24
25    def __init__(self, rule_name: str, message: str) -> None:
26        self.rule_name = rule_name
27        super().__init__(f"Guardrail '{rule_name}' violated: {message}")

 37class HumanInLoopBehavior(BaseBehavior):
 38    """
 39    Pauses or suspends execution after the agent produces a result and waits
 40    for human approval before the result is returned.
 41
 42    Two modes:
 43
 44    **Sync mode** — pass ``review_callback``.
 45        The callback blocks until the human decides, returning ``True`` to
 46        approve or ``False`` to reject.  Suitable for interactive CLIs or
 47        short-lived approval windows.
 48
 49        ``async (context, result) -> bool``
 50
 51    **Async (offline) mode** — pass ``submit_callback``.
 52        The callback fires the approval request (e.g. POST to a REST API or
 53        queue) and returns an opaque ``approval_id`` string.  The behavior
 54        immediately raises :class:`PendingApproval` so the agent call
 55        terminates without blocking.  The caller is responsible for persisting
 56        the pending state and resuming once the human's decision arrives via
 57        webhook or callback.
 58
 59        ``async (context, result) -> str``  (returns ``approval_id``)
 60
 61    Args:
 62        review_callback: Sync-mode callback. Mutually exclusive with
 63            ``submit_callback``.
 64        submit_callback: Async-mode callback. Mutually exclusive with
 65            ``review_callback``.
 66
 67    Raises:
 68        ValueError: If neither or both callbacks are supplied.
 69
 70    Example — sync mode::
 71
 72        async def my_reviewer(ctx, result):
 73            answer = input(f"Approve '{ctx.task}'? [y/n]: ")
 74            return answer.strip().lower() == "y"
 75
 76        behavior = HumanInLoopBehavior(review_callback=my_reviewer)
 77
 78    Example — async (offline) mode::
 79
 80        async def submit_for_approval(ctx, result):
 81            resp = await http_client.post("/approvals", json={
 82                "task": ctx.task,
 83                "output": result.output,
 84            })
 85            return resp.json()["approval_id"]
 86
 87        behavior = HumanInLoopBehavior(submit_callback=submit_for_approval)
 88
 89        # Caller handles PendingApproval:
 90        try:
 91            result = await agent.execute(task)
 92        except PendingApproval as pa:
 93            store.save(pa.approval_id, pending_result=pa.result, task=pa.task)
 94            # ... webhook delivers decision later, caller resumes from store
 95    """
 96
 97    def __init__(
 98        self,
 99        review_callback: Optional[Callable[[BehaviorContext, Any], Awaitable[bool]]] = None,
100        submit_callback: Optional[Callable[[BehaviorContext, Any], Awaitable[str]]] = None,
101    ) -> None:
102        if review_callback is None and submit_callback is None:
103            raise ValueError(
104                "HumanInLoopBehavior requires either 'review_callback' (sync) "
105                "or 'submit_callback' (async/offline)."
106            )
107        if review_callback is not None and submit_callback is not None:
108            raise ValueError(
109                "HumanInLoopBehavior accepts only one of 'review_callback' or "
110                "'submit_callback', not both."
111            )
112        self._review_callback = review_callback
113        self._submit_callback = submit_callback
114
115    async def after_execute(self, context: BehaviorContext, result: Any) -> Any:
116        if self._submit_callback is not None:
117            # Async / offline mode — fire the request, return approval_id, exit immediately.
118            approval_id = await self._submit_callback(context, result)
119            raise PendingApproval(
120                approval_id=approval_id,
121                task=context.task,
122                result=result,
123            )
124
125        # Sync mode — block until the human responds.
126        approved = await self._review_callback(context, result)  # type: ignore[misc]
127        if not approved:
128            raise HumanApprovalRequired(
129                f"Human reviewer rejected the result for task: '{context.task}'"
130            )
131        return result

 9class HumanApprovalRequired(Exception):
10    """Raised when a human reviewer rejects the agent result."""

13class PendingApproval(Exception):
14    """
15    Raised in async (offline) mode to signal that an approval request has been
16    submitted and the agent result is awaiting a human decision.
17
18    The caller should catch this, persist state keyed on ``approval_id``, and
19    resume processing when the approval webhook/callback delivers the decision.
20
21    Attributes:
22        approval_id: Opaque identifier returned by the ``submit_callback``.
23        task: The original task string passed to the agent.
24        result: The ``AgentResult`` produced by the agent, ready to deliver
25                once approved.
26    """
27
28    def __init__(self, approval_id: str, task: str, result: Any) -> None:
29        self.approval_id = approval_id
30        self.task = task
31        self.result = result
32        super().__init__(
33            f"Approval pending for task: '{task}' (approval_id={approval_id})"
34        )

 21class CircuitBreakerBehavior(BaseBehavior):
 22    """
 23    Prevents repeated calls to a failing agent by tripping an open circuit.
 24
 25    State transitions::
 26
 27        CLOSED ──(failures >= threshold)──► OPEN
 28        OPEN   ──(recovery_timeout elapsed)──► HALF_OPEN
 29        HALF_OPEN ──(success)──► CLOSED
 30        HALF_OPEN ──(failure)──► OPEN
 31
 32    Args:
 33        failure_threshold: Consecutive failures before opening (default: 5).
 34        recovery_timeout: Seconds to wait in OPEN state before probing (default: 60).
 35        metrics: Optional :class:`BasicMetricsCollector` for state/failure tracking.
 36        logger: Optional :class:`BasicLogger`.
 37
 38    Example::
 39
 40        from gmf_forge_ai_shared_core.observability import BasicMetricsCollector, BasicLogger
 41
 42        behavior = CircuitBreakerBehavior(
 43            failure_threshold=3,
 44            recovery_timeout=30,
 45            metrics=BasicMetricsCollector(),
 46            logger=BasicLogger("circuit_breaker"),
 47        )
 48    """
 49
 50    def __init__(
 51        self,
 52        failure_threshold: int = 5,
 53        recovery_timeout: float = 60.0,
 54        metrics: Optional[Any] = None,
 55        logger: Optional[Any] = None,
 56    ) -> None:
 57        self.failure_threshold = failure_threshold
 58        self.recovery_timeout = recovery_timeout
 59        self._metrics = metrics
 60        self._logger = logger
 61
 62        self._state = CircuitState.CLOSED
 63        self._failure_count = 0
 64        self._opened_at: Optional[float] = None
 65
 66    # ------------------------------------------------------------------
 67    # Internal helpers
 68    # ------------------------------------------------------------------
 69
 70    def _transition(self, new_state: CircuitState) -> None:
 71        if self._state == new_state:
 72            return
 73        old = self._state.value
 74        self._state = new_state
 75        if self._logger:
 76            self._logger.info(
 77                "Circuit state transition",
 78                from_state=old,
 79                to_state=new_state.value,
 80            )
 81        if self._metrics:
 82            self._metrics.gauge("circuit_breaker.state", 1.0, state=new_state.value)
 83
 84    # ------------------------------------------------------------------
 85    # BaseBehavior hooks
 86    # ------------------------------------------------------------------
 87
 88    async def before_execute(self, context: BehaviorContext) -> BehaviorContext:
 89        if self._state == CircuitState.OPEN:
 90            elapsed = time.monotonic() - (self._opened_at or 0)
 91            if elapsed >= self.recovery_timeout:
 92                self._transition(CircuitState.HALF_OPEN)
 93            else:
 94                raise CircuitOpenError(
 95                    f"Circuit is OPEN for agent '{context.agent_id}'. "
 96                    f"Retry in {self.recovery_timeout - elapsed:.1f}s."
 97                )
 98        return context
 99
100    async def after_execute(self, context: BehaviorContext, result: Any) -> Any:
101        # Success — reset failures, close circuit if probing
102        self._failure_count = 0
103        if self._state == CircuitState.HALF_OPEN:
104            self._transition(CircuitState.CLOSED)
105        return result
106
107    async def on_error(
108        self, context: BehaviorContext, error: Exception
109    ) -> Optional[Any]:
110        self._failure_count += 1
111        if self._metrics:
112            self._metrics.increment("circuit_breaker.failures", agent_id=context.agent_id)
113
114        if self._state == CircuitState.HALF_OPEN or self._failure_count >= self.failure_threshold:
115            self._opened_at = time.monotonic()
116            self._transition(CircuitState.OPEN)
117
118        raise error

11class CircuitState(Enum):
12    CLOSED = "closed"       # Normal operation
13    OPEN = "open"           # Blocking all calls
14    HALF_OPEN = "half_open" # Allowing a single probe call

17class CircuitOpenError(Exception):
18    """Raised when a call is attempted while the circuit is OPEN."""

15class RateLimitBehavior(BaseBehavior):
16    """
17    Enforces a maximum call rate using a token bucket algorithm.
18
19    Args:
20        calls_per_second: Maximum calls allowed per second (default: 1.0).
21        burst: Maximum burst capacity — tokens that can accumulate while idle
22            (default: equals ``calls_per_second``).
23        wait_on_limit: If ``True`` (default), sleep until a token is available.
24            If ``False``, raise :class:`RateLimitExceededError` immediately.
25
26    Example::
27
28        behavior = RateLimitBehavior(calls_per_second=2.0, burst=5)
29    """
30
31    def __init__(
32        self,
33        calls_per_second: float = 1.0,
34        burst: Optional[float] = None,
35        wait_on_limit: bool = True,
36    ) -> None:
37        self.calls_per_second = calls_per_second
38        self.burst = burst if burst is not None else calls_per_second
39        self.wait_on_limit = wait_on_limit
40
41        self._tokens: float = self.burst
42        self._last_refill: float = time.monotonic()
43        self._lock = asyncio.Lock()
44
45    # ------------------------------------------------------------------
46    # Token bucket
47    # ------------------------------------------------------------------
48
49    def _refill(self) -> None:
50        now = time.monotonic()
51        elapsed = now - self._last_refill
52        self._tokens = min(self.burst, self._tokens + elapsed * self.calls_per_second)
53        self._last_refill = now
54
55    async def before_execute(self, context: BehaviorContext) -> BehaviorContext:
56        async with self._lock:
57            while True:
58                self._refill()
59                if self._tokens >= 1.0:
60                    self._tokens -= 1.0
61                    return context
62
63                if not self.wait_on_limit:
64                    raise RateLimitExceededError(
65                        f"Rate limit of {self.calls_per_second} call/s exceeded "
66                        f"for agent '{context.agent_id}'."
67                    )
68
69                wait = (1.0 - self._tokens) / self.calls_per_second
70                await asyncio.sleep(wait)

11class RateLimitExceededError(Exception):
12    """Raised when the rate limit is exceeded and wait_on_limit=False."""

12class AuditBehavior(BaseBehavior):
13    """
14    Logs every agent input, output, and error using shared-core observability.
15
16    Uses :class:`BasicLogger` for structured log lines and
17    :class:`BasicMetricsCollector` to track invocation counts.
18
19    Args:
20        logger: :class:`BasicLogger` instance. A default named
21            ``"gmf_forge_ai.audit"`` is created if not provided.
22        metrics: Optional :class:`BasicMetricsCollector` for counting events.
23
24    Example::
25
26        from gmf_forge_ai_shared_core.observability import BasicLogger, BasicMetricsCollector
27
28        behavior = AuditBehavior(
29            logger=BasicLogger("my_app.audit"),
30            metrics=BasicMetricsCollector(),
31        )
32    """
33
34    def __init__(
35        self,
36        logger: Optional[BasicLogger] = None,
37        metrics: Optional[BasicMetricsCollector] = None,
38    ) -> None:
39        self._logger = logger or BasicLogger("gmf_forge_ai.audit")
40        self._metrics = metrics
41        # Stores start timestamps keyed by agent_id+task to compute duration
42        self._start_times: dict = {}
43
44    async def before_execute(self, context: BehaviorContext) -> BehaviorContext:
45        key = f"{context.agent_id}:{context.task}"
46        self._start_times[key] = time.monotonic()
47        self._logger.info(
48            "Agent task started",
49            agent_id=context.agent_id,
50            task=context.task,
51            attempt=context.attempt,
52        )
53        if self._metrics:
54            self._metrics.increment("behavior.audit.invocations", agent_id=context.agent_id)
55        return context
56
57    async def after_execute(self, context: BehaviorContext, result: Any) -> Any:
58        key = f"{context.agent_id}:{context.task}"
59        duration_ms = round((time.monotonic() - self._start_times.pop(key, time.monotonic())) * 1000)
60        output = result.output if hasattr(result, "output") else str(result)
61        self._logger.info(
62            "Agent task completed",
63            agent_id=context.agent_id,
64            task=context.task,
65            duration_ms=duration_ms,
66            success=True,
67            output_preview=output[:200],
68        )
69        if self._metrics:
70            self._metrics.histogram("behavior.audit.duration_ms", duration_ms, agent_id=context.agent_id)
71        return result
72
73    async def on_error(
74        self, context: BehaviorContext, error: Exception
75    ) -> Optional[Any]:
76        key = f"{context.agent_id}:{context.task}"
77        duration_ms = round((time.monotonic() - self._start_times.pop(key, time.monotonic())) * 1000)
78        self._logger.error(
79            "Agent task failed",
80            agent_id=context.agent_id,
81            task=context.task,
82            duration_ms=duration_ms,
83            error=str(error),
84            error_type=type(error).__name__,
85        )
86        if self._metrics:
87            self._metrics.increment("behavior.audit.errors", agent_id=context.agent_id)
88        return None