Proposer

class AsyncReflectiveMutationProposer:
    """Generates text mutations via LLM reflection.

    This proposer takes a candidate's current component texts and feedback
    data, then uses an ADK reflection function to generate improved versions.
    It handles empty datasets gracefully by returning None without making
    LLM calls, and retries an empty reflection response once before raising
    EmptyProposalError.

    Terminology:
        - component: Evolvable unit with name + text (the "gear" being tuned)
        - component_text: The text content of a component
        - trial: One record {input, output, feedback, trajectory}
        - trials: Collection of trial records for reflection
        - proposed_component_text: The improved text for the same component

    Attributes:
        adk_reflection_fn (ReflectionFn): ADK reflection function for proposing
            mutations. Created via `create_adk_reflection_fn()`.
        max_trials (int | None): Maximum trials per component handed to the
            reflection function, or None for no limit.
        max_trial_chars (int | None): Maximum length of any string value in
            a trial handed to the reflection function, or None for no limit.
        retry_backoff_seconds (float): Seconds to wait before retrying a
            retryable reflection error.
        last_reasoning (str | None): Most recent non-None reasoning from the
            last ``propose()`` call.
        last_token_usage (TokenRollup): Token usage of the reflection calls
            made by the last ``propose()`` call, one row per call that
            returned; a call that reported no usage is an unknown row and a
            call that raised is no row.

    Examples:
        Standard usage with ADK reflection agent:

        ```python
        from gepa_adk.engine import create_adk_reflection_fn

        reflection_fn = create_adk_reflection_fn(reflection_agent, executor)
        proposer = AsyncReflectiveMutationProposer(adk_reflection_fn=reflection_fn)
        result = await proposer.propose(
            candidate={"instruction": "Be helpful"},
            reflective_dataset={"instruction": [trials]},
            components_to_update=["instruction"],
        )
        ```

    Notes:
        ADK-based reflection via `adk_reflection_fn` is the only supported
        approach. Use `create_adk_reflection_fn()` from
        `gepa_adk.engine.adk_reflection` to create the reflection function.
    """

    def __init__(
        self,
        adk_reflection_fn: ReflectionFn,
        *,
        max_trials: int | None = None,
        max_trial_chars: int | None = None,
        retry_backoff_seconds: float = 2.0,
    ) -> None:
        """Initialize the mutation proposer.

        Args:
            adk_reflection_fn: Async callable for ADK-based reflection.
                Takes (component_text, trials, component_name) and returns
                (proposed_text, reasoning) tuple. Create with
                `create_adk_reflection_fn()` from
                `gepa_adk.engine.adk_reflection`.
            max_trials: Maximum trials per component handed to the
                reflection function. Failing trials take up to half the
                slots (rounded up), passing trials the rest. None sends
                every trial.
            max_trial_chars: Maximum length of any string value inside a
                trial handed to the reflection function. Longer strings are
                cut and marked. None leaves strings intact.
            retry_backoff_seconds: Seconds to wait before retrying a
                retryable reflection error. Zero retries immediately.

        Raises:
            ValueError: If adk_reflection_fn is None, or if
                retry_backoff_seconds is not a non-negative number.

        Examples:
            ```python
            from gepa_adk.engine import create_adk_reflection_fn

            reflection_fn = create_adk_reflection_fn(reflection_agent, executor)
            proposer = AsyncReflectiveMutationProposer(adk_reflection_fn=reflection_fn)
            ```

        Notes:
            Configuration validation happens immediately to fail fast rather
            than waiting until the first propose() call. After each
            ``propose()`` call, ``self.last_reasoning`` holds the most
            recent non-None reasoning string (or None) and
            ``self.last_token_usage`` the rollup of its reflection calls; it
            starts as a new zero rollup object. The backoff is validated
            here too.
        """
        if adk_reflection_fn is None:
            raise ValueError(
                "adk_reflection_fn is required. Use create_adk_reflection_fn() "
                "from gepa_adk.engine.adk_reflection to create one."
            )
        if (
            isinstance(retry_backoff_seconds, bool)
            or not isinstance(retry_backoff_seconds, (int, float))
            or not math.isfinite(retry_backoff_seconds)
            or retry_backoff_seconds < 0
        ):
            raise ValueError(
                "retry_backoff_seconds must be a non-negative number, got "
                f"{retry_backoff_seconds!r}."
            )

        self.adk_reflection_fn = adk_reflection_fn
        self.max_trials = max_trials
        self.max_trial_chars = max_trial_chars
        self.retry_backoff_seconds = retry_backoff_seconds
        self.last_reasoning: str | None = None
        self.last_token_usage: TokenRollup = _zero_usage()

        # Log proposer initialization
        logger.info("proposer_initialized", reflection_method="adk")

    async def propose(
        self,
        candidate: dict[str, str],
        reflective_dataset: ReflectiveDataset,
        components_to_update: list[str],
    ) -> ProposalResult:
        """Propose mutated component text via LLM reflection.

        Args:
            candidate (dict[str, str]): Current candidate component texts.
                Keys are component names, values are component text.
                Example: {"instruction": "Be helpful and concise"}
            reflective_dataset (ReflectiveDataset): Trials per component name.
                Each trial contains input, output, feedback, and optional
                trajectory.
                Example: {"instruction": [{
                    "input": "Hello",
                    "output": "Hi there!",
                    "feedback": {"score": 0.75, "feedback_text": "Could be more formal"},
                    "trajectory": {...}
                }]}
            components_to_update (list[str]): Component names to generate
                proposals for. Components missing from the candidate or the
                reflective dataset (or with empty trials) are silently
                skipped. Example: ["instruction"]

        Returns:
            ProposalResult: Dictionary mapping component names to proposed
                component text, or None if the reflective dataset is empty
                or has no entries for the requested components.

        Raises:
            EmptyProposalError: If ADK reflection returns an empty or
                whitespace-only response twice for the same component.
            ReflectionError: If the reflection function raises. A retryable
                error is raised after the retry fails too (``attempts=2``);
                a non-retryable one is raised at once (``attempts=1``).
            EvolutionError: If ADK reflection returns a non-string response.
                ``EvolutionError`` subclasses the reflection function
                raises, such as ``ReflectionTimeoutError`` and
                ``IncompleteProposalError``, propagate unwrapped and are not
                retried.

        Examples:
            ```python
            result = await proposer.propose(
                candidate={"instruction": "Be helpful"},
                reflective_dataset={
                    "instruction": [
                        {
                            "input": "I am the King",
                            "output": "Hey!",
                            "feedback": {"score": 0.3, "feedback_text": "Too casual"},
                            "trajectory": {...},
                        }
                    ]
                },
                components_to_update=["instruction"],
            )
            # result: {"instruction": "Greet users formally..."}
            ```

        Notes:
            Calls the reflection function directly with
            ``(component_text, trials, component_name)``. The function
            returns ``(proposed_text, reasoning)`` or
            ``(proposed_text, reasoning, token_usage)``; reasoning is stored
            in ``self.last_reasoning`` (last non-None value wins) and each
            call that returned adds one row to ``self.last_token_usage``,
            which is reset to a new zero rollup object at the start of the
            call. When
            ``max_trials`` or ``max_trial_chars`` is set, the trials are
            capped before the call and ``proposer.trials_capped`` is logged
            whenever a trial was dropped or a string truncated. An empty
            response is retried once (logged as ``proposer.empty_retry``);
            a second empty response raises EmptyProposalError, which the
            engine records as a skipped iteration. A retryable exception
            from the reflection function is retried once after
            ``retry_backoff_seconds`` (logged as ``proposer.error_retry``);
            the two retries share the same two attempts. Non-string
            responses raise EvolutionError.
        """
        # Reset reasoning and reflection usage at start of each propose() call
        self.last_reasoning = None
        self.last_token_usage = _zero_usage()

        # Early return for empty dataset (no LLM calls)
        if not reflective_dataset:
            return None

        proposals = {}

        for component in components_to_update:
            # Skip if component not in reflective_dataset or has empty feedback
            if component not in reflective_dataset:
                continue

            trials: list[dict[str, Any]] = [
                dict(t) for t in reflective_dataset[component]
            ]
            if not trials:
                continue

            # Skip component not in candidate
            if component not in candidate:
                continue

            component_text = candidate[component]
            trials = self._cap_trials(component, trials)

            logger.debug(
                "proposer.reflection_path",
                method="adk",
                component=component,
            )

            proposals[component] = await self._reflect_with_retry(
                component_text, trials, component
            )

        # Return None if no valid proposals generated
        if not proposals:
            return None

        return proposals

    def _cap_trials(
        self, component: str, trials: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Apply the trial-count and string-length caps for one component.

        Args:
            component: Component name, used in the log event.
            trials: Trial records in batch order. They are not mutated.

        Returns:
            The trials to hand to the reflection function.
        """
        total = len(trials)
        if self.max_trials is not None:
            trials = _select_trials(trials, self.max_trials)
        truncated_fields = dropped_chars = 0
        if self.max_trial_chars is not None:
            trials, truncated_fields, dropped_chars = _truncate_trials(
                trials, self.max_trial_chars
            )
        dropped_trials = total - len(trials)
        if dropped_trials or truncated_fields:
            logger.info(
                "proposer.trials_capped",
                component=component,
                kept=len(trials),
                dropped_trials=dropped_trials,
                truncated_fields=truncated_fields,
                dropped_chars=dropped_chars,
            )
        return trials

    async def _reflect_with_retry(
        self,
        component_text: str,
        trials: list[dict[str, Any]],
        component: str,
    ) -> str:
        """Call reflection, retrying once on an empty response or a transient error.

        Args:
            component_text (str): Current text of the component.
            trials (list[dict[str, Any]]): Trial records for reflection.
            component (str): Name of the component being evolved.

        Returns:
            The stripped, non-empty proposed component text.

        Raises:
            EmptyProposalError: If the last of the two attempts returns
                empty or whitespace-only text.
            ReflectionError: If the reflection function raises a
                non-retryable error on any attempt, or a retryable error on
                the second attempt.
            EvolutionError: Any other error from ``_reflect_once``
                propagates without a retry, including
                ``ReflectionTimeoutError`` and ``IncompleteProposalError``.

        Notes:
            The empty retry and the error retry share one budget of two
            attempts. A retryable error on the first attempt is logged as
            ``proposer.error_retry`` and followed by
            ``asyncio.sleep(retry_backoff_seconds)``. The logged ``error``
            text is cut to ``_ERROR_TEXT_MAX_CHARS`` characters plus a
            ``…[truncated, N chars omitted]`` marker.
        """
        for attempt in (1, 2):
            try:
                proposed = await self._reflect_once(
                    component_text, trials, component, attempt=attempt
                )
            except ReflectionError as error:
                if attempt == 2 or not error.retryable:  # noqa: PLR2004  # second and last attempt of the two-attempt budget
                    raise
                logger.warning(
                    "proposer.error_retry",
                    component=component,
                    error_type=type(error.cause).__name__,
                    error=_truncate_value(
                        str(error.cause), _ERROR_TEXT_MAX_CHARS, [0, 0]
                    ),
                    attempt=attempt,
                    backoff_seconds=self.retry_backoff_seconds,
                )
                await asyncio.sleep(self.retry_backoff_seconds)
                continue
            if proposed:
                return proposed
            if attempt == 1:
                logger.warning(
                    "proposer.empty_retry",
                    component=component,
                    attempt=attempt,
                )
        raise EmptyProposalError(component)

    async def _reflect_once(
        self,
        component_text: str,
        trials: list[dict[str, Any]],
        component: str,
        *,
        attempt: int = 1,
    ) -> str:
        """Call the reflection function once and validate its output type.

        Args:
            component_text: Current text of the component.
            trials: Trial records for reflection.
            component: Name of the component being evolved.
            attempt: One-based number of this reflection call, recorded on
                a raised ``ReflectionError``.

        Returns:
            The stripped proposed text, which may be empty.

        Raises:
            ReflectionError: If the reflection function raises anything
                other than an ``EvolutionError``; ``retryable`` comes from
                ``is_retryable_reflection_error``.
            EvolutionError: If the response is not a string, or if the
                reflection function raises an ``EvolutionError`` (passed
                through unwrapped), such as ``IncompleteProposalError`` for a
                truncated proposal.

        Notes:
            A call that returned adds one row to ``self.last_token_usage``:
            counted when its third element is a ``TokenUsage``, unknown for
            a 2-tuple or a None usage. A call that raised adds nothing.
        """
        try:
            result = await self.adk_reflection_fn(component_text, trials, component)
        except EvolutionError:
            raise
        except Exception as e:
            raise ReflectionError(
                component,
                cause=e,
                retryable=is_retryable_reflection_error(e),
                attempts=attempt,
            ) from e

        proposed_component_text, reasoning, *rest = result
        usage = rest[0] if rest else None
        self.last_token_usage = self.last_token_usage.combine(_call_usage(usage))

        # Store the last non-None reasoning
        if reasoning is not None:
            self.last_reasoning = reasoning

        if not isinstance(proposed_component_text, str):
            raise EvolutionError(
                "Reflection agent must return a string, got "
                f"{type(proposed_component_text).__name__}."
            )
        return proposed_component_text.strip()