afnio.optim

class TGD(Optimizer):
    """
    Textual Gradient Descent (TGD) optimizer.

    TGD is an optimization algorithm for language-model–based systems where
    gradients are represented and propagated as *natural language feedback*
    rather than numerical tensors. Instead of computing numerical derivatives,
    TGD relies on a language model to generate textual critiques (gradients)
    that are used to iteratively refine prompt-based parameters.

    This implementation follows the ideas introduced in the *TextGrad* paper,
    which proposes treating language-model feedback as a differentiable signal
    for optimizing textual variables and prompt programs.

    TGD operates over [`Variable`][afnio.Variable] objects and consumes textual
    gradients produced by the automatic differentiation process. These gradients
    are used to update the optimized variables, with optional momentum applied
    to recent gradient history to stabilize and accelerate optimization.

    Parameters are organized into parameter groups, similar to optimizers in
    PyTorch. This allows different optimization settings—such as optimization
    meta-prompts (`messages`), `constraints`, and `momentum`—to be applied
    consistently across groups.

    **References:**

    - *TextGrad: Automatic Differentiation via Large Language Models*
        [https://arxiv.org/abs/2406.07496](https://arxiv.org/abs/2406.07496)
    """

    def __init__(
        self,
        params: ParamsT,
        model_client: Optional[ChatCompletionModel],
        messages: MultiTurnMessages = TGD_MESSAGES,
        inputs: Optional[Dict[str, Union[str, Variable]]] = None,
        constraints: Optional[List[Union[str, Variable]]] = None,
        momentum: int = 0,
        **completion_args,
    ):
        """Initialize the Textual Gradient Descent (TGD) optimizer.

        Args:
            params (iterable): Iterable of parameters to optimize or dicts defining
                parameter groups.
            model_client: LM model client used for optimization.
            messages: Messages for multi-turn interactions. It typically defines
                the optimizer system prompt and user instruction. In-context
                examples (shots) can be added as well.
            inputs: Dynamic values to fill placeholders within message templates
            constraints: A list of natural language constraints for optimization.
            momentum (int, optional): Momentum window size. Tracks the last `momentum`
                gradients, which helps accelerate updates in the right direction and
                dampen oscillations. Defaults to 0.
            completion_args (Dict[str, Any], optional): Additional arguments to pass to
                the model client when generating text completions. Defaults to an
                empty dictionary.
        """
        # Workaround to trigger TGD_MESSAGES registration with the server
        # and store related variable_ids on the client side
        if messages is TGD_MESSAGES:
            messages = [
                {
                    "role": "system",
                    "content": [
                        Variable(
                            data="Placeholder for Textual Gradient Descent optimizer system prompt",  # noqa: E501
                            role="Textual Gradient Descent optimizer system prompt",
                        )
                    ],
                },
                {
                    "role": "user",
                    "content": [
                        Variable(
                            data="Placeholder for Textual Gradient Descent optimizer user prompt",  # noqa: E501
                            role="Textual Gradient Descent optimizer user prompt",
                        )
                    ],
                },
            ]

        defaults = dict(
            model_client=model_client,
            messages=messages,
            inputs=inputs or {},
            constraints=constraints or [],
            momentum=momentum,
            completion_args=completion_args,
        )
        super().__init__(params, defaults)

    def step(
        self, closure: Optional[Callable] = None
    ) -> Optional[Tuple[Variable, Variable]]:
        """Performs a single optimization step.

        Args:
            closure: A closure that reevaluates the model and returns the loss.

        Returns:
            The loss if `closure` is provided, otherwise None. The loss should \
            return a numerical or textual score and a textual explanation, both \
            wrapped as [`Variable`][afnio.Variable] objects
        """
        loss = closure() if closure else (None, None)
        super().step()
        return loss

    def _extract_variable_ids_from_state(self, state):
        """
        Extract only the variable_ids of deepcopied parameters (i.e., those generated
        on the server) from the optimizer state.

        Args:
            state (list): The serialized optimizer state as returned by the server.

        Returns:
            Set[str]: Set of variable_ids for deepcopied parameters.
        """
        var_ids = set()
        for entry in state:
            momentum_buffer = entry.get("value", {}).get("momentum_buffer", [])
            for buf_entry in momentum_buffer:
                if (
                    isinstance(buf_entry, list)
                    and len(buf_entry) > 0
                    and isinstance(buf_entry[0], dict)
                    and "variable_id" in buf_entry[0]
                ):
                    var_ids.add(buf_entry[0]["variable_id"])
        return var_ids