Adk reflection

def is_length_stop(finish_reason: Any) -> bool:
    """Report whether a finish reason marks output cut off at the token limit.

    Args:
        finish_reason: A ``google.genai.types.FinishReason`` member, its name
            as a string in any case, litellm's ``"length"``, or ``None``.

    Returns:
        True when the reason is ``MAX_TOKENS`` or ``LENGTH`` (case-insensitive);
        False for ``None``, an empty value or any other reason.

    Examples:
        ```python
        from google.genai import types

        is_length_stop(types.FinishReason.MAX_TOKENS)  # True
        is_length_stop("length")  # True
        is_length_stop("STOP")  # False
        is_length_stop(None)  # False
        ```
    """
    if finish_reason is None or finish_reason == "":
        return False
    name = getattr(finish_reason, "name", None)
    text = name if isinstance(name, str) else str(finish_reason)
    return text.upper() in _LENGTH_STOP_NAMES

def create_adk_reflection_fn(
    reflection_agent: Any,  # LlmAgent from google.adk.agents
    executor: AgentExecutorProtocol,
    output_key: str = "proposed_component_text",
    timeout_seconds: int | None = None,
    reasoning_tags: Sequence[str] = _DEFAULT_REASONING_TAGS,
) -> ADKReflectionFn:
    """Create a reflection function from an ADK LlmAgent.

    This factory function creates an async callable that uses the Google ADK
    framework for reflection. The returned function can be passed to
    AsyncReflectiveMutationProposer as the adk_reflection_fn parameter.

    The caller is responsible for selecting the appropriate reflection agent.
    See ``gepa_adk.api.evolve()`` for the standard wiring pattern.

    Args:
        reflection_agent: ADK LlmAgent configured with instruction containing
            `{component_text}` and `{trials}` placeholders. The agent's
            instruction should include logic for improving text based on
            trial results.
        executor: AgentExecutorProtocol implementation for unified agent
            execution. Handles session management and execution, enabling
            feature parity across all agent types.
        output_key: Key in session state where ADK stores the agent's output.
            Defaults to "proposed_component_text". When set, the agent's output_key
            is configured to this value, and output is retrieved from session
            state after execution. Falls back to event-based extraction if
            the output_key is not found in session state.
        timeout_seconds: Seconds the reflection agent may run per call. Passed
            to ``executor.execute_agent`` as ``timeout_seconds`` only when not
            ``None``; ``None`` (default) keeps the executor's own default.
        reasoning_tags: Reasoning tag names (without angle brackets) that
            mark an incomplete proposal when the stripped text opens with one
            of them and never closes it. Matched case-insensitively. Defaults
            to ``("think", "thinking", "reasoning")``.

    Returns:
        Async callable matching ReflectionFn signature that generates proposed
        component text via the ADK agent and returns
        ``(proposed_text, reasoning, token_usage)``. Typed as
        ``ADKReflectionFn`` so its ``component_name`` parameter may also be
        passed by keyword.

    Raises:
        RuntimeError: If ADK agent execution fails (propagated from executor).
        ReflectionTimeoutError: From the returned callable, when the executor
            reports ``ExecutionStatus.TIMEOUT``.
        IncompleteProposalError: From the returned callable, when the last
            captured event with a finish reason reports a length stop, or the
            proposal opens one of ``reasoning_tags`` without closing it.

    Examples:
        Basic usage with executor:

        ```python
        from google.adk.agents import LlmAgent
        from google.adk.models.lite_llm import LiteLlm
        from gepa_adk.adapters.execution.agent_executor import AgentExecutor
        from gepa_adk.engine.adk_reflection import create_adk_reflection_fn

        agent = LlmAgent(
            name="InstructionReflector",
            model=LiteLlm(model="ollama_chat/gpt-oss:20b"),
            instruction=\"\"\"Improve this component text:
            {component_text}

            Based on these trials:
            {trials}

            Return proposed component text only.\"\"\"
        )

        executor = AgentExecutor()
        reflection_fn = create_adk_reflection_fn(agent, executor=executor)
        trials = [{"input": "Hi", "output": "Hey", "feedback": {"score": 0.5}}]
        proposed, reasoning, usage = await reflection_fn(
            "Be helpful", trials, "instruction"
        )
        ```

    See Also:
        - [`gepa_adk.engine.proposer`][gepa_adk.engine.proposer]: Module containing
          ReflectionFn type alias and AsyncReflectiveMutationProposer class.
        - [`gepa_adk.api.evolve`][gepa_adk.api.evolve]: Standard wiring pattern
          for constructing the reflection chain.

    Notes:
        Opens a fresh ADK session for each invocation via AgentExecutor, ensuring
        complete isolation between reflection operations. State is initialized with
        component_text (str) and trials (JSON-serialized list of trial records).
        A timed-out execution logs ``reflection.timeout`` and raises
        ``ReflectionTimeoutError``; ``reflection.empty_response`` is logged only
        for a completed run that produced no text. An incomplete proposal logs
        ``reflection.incomplete`` at warning and raises
        ``IncompleteProposalError``; the length check runs before the
        empty-response path, so a length stop with no text is incomplete.
    """
    timeout_source = "config" if timeout_seconds is not None else "executor_default"
    timeout_kwargs: dict[str, Any] = (
        {"timeout_seconds": timeout_seconds} if timeout_seconds is not None else {}
    )

    # Configure output_key on agent if not already set
    # This enables ADK's automatic output storage to session.state
    if output_key and (
        not hasattr(reflection_agent, "output_key") or not reflection_agent.output_key
    ):
        reflection_agent.output_key = output_key
        logger.debug(
            "reflection.output_key.configured",
            output_key=output_key,
            agent_name=getattr(reflection_agent, "name", "unknown"),
        )

    async def reflect(
        component_text: str,
        trials: list[dict[str, Any]],
        component_name: str = "",
    ) -> tuple[str, str | None, TokenUsage | None]:
        """Reflect on component text using ADK agent to generate proposed version.

        Uses the configured ADK reflection agent to analyze the current component
        text and trials, then generates proposed component text based on the
        performance results.

        Args:
            component_text (str): The current component text to improve.
            trials (list[dict[str, Any]]): List of trial records from evaluation. Each
                trial contains input, output, feedback, and optional trajectory.
            component_name (str): Component name passed by the proposer. Logged
                for observability but not used for agent selection (caller
                pre-selects the agent). Defaults to empty string, which the
                logs and ``ReflectionTimeoutError`` report as ``"unknown"``.

        Returns:
            Tuple of (proposed_component_text, reasoning, token_usage). The
            proposed text is empty string if the agent produces no output.
            Reasoning is extracted from thought-tagged parts of captured
            events, or None if no thinking content is available. Token usage
            is read from the captured events, or None when the executor
            captured none or they carried no usage.

        Raises:
            RuntimeError: If the executor reports a FAILED execution. It and
                any other unexpected error are logged as ``reflection.error``
                and re-raised for upstream handling.
            ReflectionTimeoutError: If the executor reports a timeout. It is
                logged as ``reflection.timeout`` and not retried here.
            IncompleteProposalError: If the last model response reports a
                length stop or the text opens a reasoning tag it never
                closes. It is logged as ``reflection.incomplete``.

        Notes:
            Opens a unique session with fresh state for each invocation via
            AgentExecutor, ensuring isolation between reflection operations.
            The configured timeout, if any, is passed to the executor, and the
            start log records it with its source. Completeness is checked
            before the empty-response handling, and only for a string value;
            a non-string value is left to the proposer's type check.
        """
        # Generate unique session ID for this reflection
        session_id = f"reflect_{uuid4()}"

        # Log reflection start
        logger.info(
            "reflection.start",
            session_id=session_id,
            component_text_length=len(component_text),
            trial_count=len(trials),
            component_name=component_name or "unknown",
            timeout_seconds=timeout_seconds,
            timeout_source=timeout_source,
        )

        # Prepare session state for template substitution
        session_state: dict[str, Any] = {
            "component_text": component_text,
            "trials": json.dumps(trials),
        }

        # Simple trigger message - data is in session state via template placeholders
        user_message = "Please improve the component text based on the trial results."

        try:
            result = await executor.execute_agent(
                agent=reflection_agent,
                input_text=user_message,
                session_state=session_state,
                **timeout_kwargs,
            )

            if result.status == ExecutionStatus.TIMEOUT:
                component = component_name or "unknown"
                logger.warning(
                    "reflection.timeout",
                    session_id=result.session_id,
                    component=component,
                    timeout_seconds=timeout_seconds,
                    timeout_source=timeout_source,
                )
                raise ReflectionTimeoutError(component, timeout_seconds)  # noqa: TRY301  # outer handler logs and re-raises it unchanged

            if result.status == ExecutionStatus.FAILED:
                logger.error(
                    "reflection.error",
                    session_id=result.session_id,
                    error=result.error_message,
                )
                raise RuntimeError(result.error_message or "Executor returned FAILED")  # noqa: TRY301  # outer handler logs and re-raises it unchanged

            proposed_component_text = result.extracted_value or ""
            captured = getattr(result, "captured_events", None)

            # A non-string value is left for the proposer's type check
            incomplete = (
                _incomplete_reason(proposed_component_text, captured, reasoning_tags)
                if isinstance(proposed_component_text, str)
                else None
            )
            if incomplete is not None:
                component = component_name or "unknown"
                logger.warning(
                    "reflection.incomplete",
                    session_id=result.session_id,
                    component=component,
                    finish_reason=incomplete,
                    response_length=len(proposed_component_text),
                )
                raise IncompleteProposalError(  # noqa: TRY301  # outer handler logs and re-raises it unchanged
                    component,
                    finish_reason=incomplete,
                    raw_text=proposed_component_text,
                )

            # Extract reasoning from captured events
            reasoning = extract_reasoning_from_events(captured)
            if reasoning:
                logger.debug(
                    "reflection.reasoning_captured",
                    reasoning_length=len(reasoning),
                )

            usage = _extract_token_usage(captured) if captured else None

            # Log reflection complete
            logger.info(
                "reflection.complete",
                session_id=result.session_id,
                response_length=len(proposed_component_text),
            )

        except (ReflectionTimeoutError, IncompleteProposalError):
            raise
        except Exception as e:
            logger.error(  # noqa: TRY400  # re-raised below; the caller owns the traceback
                "reflection.error",
                session_id=session_id,
                error=str(e),
                error_type=type(e).__name__,
            )
            raise
        else:
            # Handle empty response
            if not proposed_component_text:
                logger.warning(
                    "reflection.empty_response",
                    session_id=result.session_id,
                )
                return ("", reasoning, usage)

            return (proposed_component_text, reasoning, usage)

    return reflect