afnio.cognitive.modules

class Add(Module):
    """
    Performs element-wise addition of two input [`Variable`][afnio.Variable]s.

    This module utilizes the [`Add`][afnio.autodiff.basic_ops.Add] operation from
    `afnio.autodiff.basic_ops`. The inputs must be instances of the
    [`Variable`][afnio.Variable] class. The `forward` method applies the addition
    operation to the [`data`][afnio.Variable.data] field of the inputs and returns
    the resulting [`Variable`][afnio.Variable].

    Note:
        This module does not have any trainable parameters.

    Examples:
        >>> from afnio import cognitive as cog
        >>> class Addition(cog.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         self.add = cog.Add()
        >>>     def forward(self, x, y):
        ...         return self.add(x, y)
        >>> input1 = afnio.Variable(data="abc", role="input1")
        >>> input2 = afnio.Variable(data="def", role="input2")
        >>> addition = Addition()
        >>> result = addition(input1, input2)
        >>> print(result.data)
        'abcdef'
        >>> print(result.role)
        'input1 and input2'

    Raises:
        TypeError: If either input is not an instance of [`Variable`][afnio.Variable].
        TypeError: If addition between the input types is not allowed.
        ValueError: If a scalar [`data`][afnio.Variable.data] is added
            to a list[`data`][afnio.Variable.data].
        ValueError: If list [`data`][afnio.Variable.data] fields
            have mismatched lengths.

    See Also:
        [`afnio.autodiff.basic_ops.Add`][afnio.autodiff.basic_ops.Add]
        for the underlying operation.
    """

    def __init__(self):
        super().__init__()

    def forward(self, x: Variable, y: Variable) -> Variable:
        """
        Forward pass for element-wise addition.

        Warning:
            Users should not call this method directly. Instead, they should call the
            module instance itself, which will internally invoke this `forward` method.

        Args:
            x: The first input `Variable`.
            y: The second input `Variable`.

        Returns:
            A new `Variable` instance representing the result of the addition, \
            with appropriately aggregated [`data`][afnio.Variable.data], \
            [`role`][afnio.Variable.role], and \
            [`requires_grad`][afnio.Variable.requires_grad] attributes.

        Raises:
            TypeError: If either input is not an instance
                of [`Variable`][afnio.Variable].
            TypeError: If addition between the input types is not allowed.
            ValueError: If a scalar [`data`][afnio.Variable.data] is added
                to a list[`data`][afnio.Variable.data].
            ValueError: If list [`data`][afnio.Variable.data] fields
                have mismatched lengths.
        """
        return AddOp.apply(x, y)

class ChatCompletion(Module):
    """
    Generates a chat-based completion using a language model.

    This module leverages the [`ChatCompletion`][afnio.autodiff.lm_ops.ChatCompletion]
    operation from `afnio.autodiff.lm_ops` to perform model inference. The `forward`
    method accepts a list of `messages` representing the conversation history, with
    optional dynamic `inputs` for filling placeholders within the messages. The
    `forward_model_client` is responsible for interfacing with the language model
    (e.g., gpt-4.1), while `completion_args` allows customization of generation
    parameters such as temperature, maximum tokens, and seed.

    Examples:
        >>> from afnio import cognitive as cog
        >>> from afnio.models.openai import OpenAI
        >>> from afnio import set_backward_model_client
        >>> fwd_model_client = OpenAI()
        >>> fwd_model_args = {"model": "gpt-4o", "temperature": 0.7}
        >>> set_backward_model_client("openai/gpt-4o")
        >>> class Assistant(cog.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         self.chat = cog.ChatCompletion()
        ...     def forward(self, fwd_model, messages, inputs, **completion_args):
        ...         return self.chat(fwd_model, messages, inputs, **completion_args)
        >>> system = Variable(
        ...     "You are a helpful assistant.",
        ...     role="system instruction",
        ...     requires_grad=True
        ... )
        >>> user = Variable("Translate 'Hello' to {language}.", role="user query")
        >>> language = afnio.Variable("Italian", role="language")
        >>> messages = [
        ...     {"role": "system", "content": [system]},
        ...     {"role": "user", "content": [user]},
        ... ]
        >>> agent = Assistant()
        >>> response = agent(
        ...     fwd_model_client,
        ...     messages,
        ...     inputs={"language": language},
        ...     **fwd_model_args
        ... )
        >>> print(response.data)
        'Ciao'
        >>> feedback = Variable("Use only capital letters.", role="feedback")
        >>> response.backward(feedback)
        >>> system.grad[0].data
        'The system instruction should enforce the use of capital letters only.'

    Raises:
        TypeError: If the types of `forward_model_client`, `messages`, or `inputs`
            are not as expected.

    See Also:
        [`afnio.autodiff.lm_ops.ChatCompletion`][afnio.autodiff.lm_ops.ChatCompletion]
        for the underlying operation.
    """

    forward_model_client: Optional[ChatCompletionModel]
    messages: MultiTurnMessages
    completion_args: Dict[str, Any]

    def __init__(self):
        super().__init__()

        self.register_model("forward_model_client", None)
        self.register_chat("messages", None)
        self.register_completion_config("completion_args", None)

    def forward(
        self,
        forward_model_client: Optional[ChatCompletionModel],
        messages: MultiTurnMessages,
        inputs: Optional[Dict[str, Union[str, Variable]]] = None,
        **completion_args,
    ) -> Variable:
        """
        Forward pass for the chat completion function.

        Warning:
            Users should not call this method directly. Instead, they should call the
            module instance itself, which will internally invoke this `forward` method.

        Args:
            forward_model_client: The LM model client used for generating
                chat completions.
            messages: A list of messages that compose the prompt/context for the LM.
                Each message is a dictionary with a `"role"` (e.g., `"system"`,
                `"user"`, `"assistant"`) and a `"content"` field, which is a list of
                `Variable` objects. The `Variable` objects in the `"content"` can
                contain placeholders (e.g., `{prediction}`, `{target}`) that will be
                populated with the corresponding values from the `inputs` dictionary.
            inputs: A dictionary mapping placeholder names to their corresponding
                values, which can be strings or `Variable` instances. These values
                will be used to populate the placeholders in the `messages` content
                before sending the prompt to the LM. For example, if a message
                `"content"` field contains the placeholder `{color}`, the `inputs`
                dictionary should have a key `"color"` with the value to substitute
                in the prompt. Optional if there are no placeholders in the messages or
                if all placeholders are directly related to `prediction` and `target`.
            **completion_args: Additional keyword arguments to pass to the LM model
                client's `chat` method, such as temperature, max tokens, or seed values,
                to customize the LLM's behavior during the evaluation.

        Returns:
            response: A `Variable` containing the LM's response. \
                The [`data`][afnio.Variable.data] field of the returned `Variable` \
                will be a string if all inputs are scalar, or a list of strings if \
                any input is a list. The `role` field will indicate that this is a \
                response to the input messages, and the `requires_grad` field will \
                be set to `True` if any of the input `Variable` objects in `messages` \
                require gradients, otherwise `False`.

        Raises:
            TypeError: If the types of `forward_model_client`, `messages`,
                or `inputs` are not as expected.
        """
        self.forward_model_client = forward_model_client
        self.messages = messages
        self.completion_args = completion_args
        return ChatCompletionOp.apply(
            self.forward_model_client,
            self.messages,
            inputs,
            **self.completion_args,
        )

class DeterministicEvaluator(Module):
    """
    Evaluates predictions deterministically using a user-defined evaluation function.

    This module utilizes the [`DeterministicEvaluator`][afnio.autodiff.evaluator.DeterministicEvaluator]
    operation from `afnio.autodiff.evaluator` to compute evaluation scores and
    explanations. The `forward` method takes in a `prediction`, a `target`, an
    evaluation function (`eval_fn`), and its purpose description (`eval_fn_purpose`).
    It also accepts a reduction function (`reduction_fn`) and its purpose description
    (`reduction_fn_purpose`) to aggregate scores if needed. The method outputs
    an evaluation `score` and an `explanation`, both as `Variable` instances. The
    `success_fn` checks if all evaluations are successful, allowing the `backward` pass
    to skip unnecessary gradient computations. The method outputs an evaluation
    `score` and an `explanation`, both as `Variable` instances.

    Examples:
        >>> from afnio import cognitive as cog
        >>> from afnio import set_backward_model_client
        >>> set_backward_model_client("openai/gpt-4o")
        >>> class ExactColor(cog.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         def exact_match_fn(pred: str, tgt: str) -> int:
        ...             return 1 if pred == tgt_data else 0
        ...         self.exact_match_fn = exact_match_fn
        ...         self.fn_purpose = "exact match"
        ...         self.reduction_fn = sum
        ...         self.reduction_fn_purpose = "summation"
        ...         self.exact_match = cog.DeterministicEvaluator()
        ...     def forward(self, prediction, target):
        ...         return self.exact_match(
        ...             prediction,
        ...             target,
        ...             self.exact_match_fn,
        ...             self.fn_purpose,
        ...             self.reduction_fn,
        ...             self.reduction_fn_purpose,
        ...         )
        >>> prediction = afnio.Variable(
        ...     data=["the color is green", "blue"],
        ...     role="color prediction",
        ...     requires_grad=True
        ... )
        >>> target = ["green", "blue"]
        >>> eval = ExactColor()
        >>> score, explanation = eval(prediction, target)
        >>> print(score.data)
        1
        >>> print(explanation.data)
        'The evaluation function, designed for 'exact match', compared the <DATA> fields of the predicted variable and the target variable across all samples in the batch, generating individual scores for each pair. These scores were then aggregated using the reduction function 'summation', resulting in a final aggregated score: 1.'
        >>> explanation.backward()
        >>> prediction.grad[0].data
        'Reassess the criteria that led to the initial prediction of 'green'.'

    Raises:
        TypeError: If the types of `prediction`, `target`, `eval_fn`, `eval_fn_purpose`,
            `success_fn`, `reduction_fn`, or `reduction_fn_purpose` are not as expected.
        ValueError: If the lengths of `prediction.data` and `target` (or `target.data`,
            when `target` is a `Variable`) do not match when both are lists, or if
            `eval_fn_purpose` (or `eval_fn_purpose.data`) is an empty string, or if
            `reduction_fn_purpose` (or `reduction_fn_purpose.data`) is an empty string,
            or if the number of scores returned by `eval_fn` does not match the number
            of samples in the batch.

    See Also:
        [`afnio.autodiff.evaluator.DeterministicEvaluator`][afnio.autodiff.evaluator.DeterministicEvaluator]
        for the underlying operation.
    """  # noqa: E501

    eval_fn: Callable[[Variable, Union[str, Variable]], List[Any]]
    eval_fn_purpose: Union[str, Variable]
    success_fn: Optional[Callable[[List[Any]], bool]]
    reduction_fn: Optional[Callable[[List[Any]], Any]]
    reduction_fn_purpose: Optional[Union[str, Variable]]

    def __init__(self):
        super().__init__()

        self.register_function("eval_fn", None)
        self.register_buffer("eval_fn_purpose", None)
        self.register_function("success_fn", None)
        self.register_function("reduction_fn", None)
        self.register_buffer("reduction_fn_purpose", None)

    def forward(
        self,
        prediction: Variable,
        target: Union[str, List[str], Variable],
        eval_fn: Callable[[Variable, Union[str, Variable]], List[Any]],
        eval_fn_purpose: Union[str, Variable],
        success_fn: Optional[Callable[[List[Any]], bool]],
        reduction_fn: Optional[Callable[[List[Any]], Any]],
        reduction_fn_purpose: Optional[Union[str, Variable]],
    ) -> Tuple[Variable, Variable]:
        """
        Forward pass for the deterministic evaluator function.

        Warning:
            Users should not call this method directly. Instead, they should call the
            module instance itself, which will internally invoke this `forward` method.

        Args:
            prediction: The predicted variable to evaluate, which can have scalar or
                list [`data`][afnio.Variable.data] (supporting both individual and
                batch processing).
            target: The target (ground truth) to compare against, which can be a string,
                a list of strings, or a `Variable`.
            eval_fn: A user-defined function that takes a prediction and a target
                and returns a list of scores for each sample. If `target` is a
                [`Variable`][afnio.Variable], the function should compare the
                [`data`][afnio.Variable.data] fields of `prediction` and `target`.
            eval_fn_purpose: A brief description of the purpose of `eval_fn`,
                used by the autodiff engine to generate the explanations.
            success_fn: A user-defined function that takes the list of scores returned
                by `eval_fn` and returns `True` if all predictions are considered
                successful, or `False` otherwise.
            reduction_fn: An optional function to aggregate scores across a batch of
                predictions and targets. If `None`, no aggregation is applied.
            reduction_fn_purpose: A brief description of the purpose of `reduction_fn`,
                used by the autodiff engine to generate explanations. Required if
                `reduction_fn` is provided.

        Returns:
            score: A variable containing the evaluation score(s),
                or their aggregation if `reduction_fn` is provided.
            explanation: A variable containing the explanation(s) of the evaluation,
                or their aggregation if `reduction_fn` is provided.

        Raises:
            TypeError: If the types of `prediction`, `target`, `eval_fn`,
                `eval_fn_purpose`, `success_fn`, `reduction_fn`,
                or `reduction_fn_purpose` are not as expected.
            ValueError: If the lengths of `prediction.data` and `target` (or
                `target.data`, when `target` is a `Variable`) do not match when
                both are lists, or if `eval_fn_purpose` (or `eval_fn_purpose.data`)
                is an empty string, or if `reduction_fn_purpose` (or
                `reduction_fn_purpose.data`) is an empty string,
                or if the number of scores returned by `eval_fn`
                does not match the number of samples in the batch.
        """
        self.eval_fn = eval_fn
        self.eval_fn_purpose = (
            None
            if eval_fn_purpose is None
            else (
                eval_fn_purpose
                if isinstance(eval_fn_purpose, Variable)
                else Variable(eval_fn_purpose)
            )
        )
        self.success_fn = success_fn
        self.reduction_fn = reduction_fn
        self.reduction_fn_purpose = (
            None
            if reduction_fn_purpose is None
            else (
                reduction_fn_purpose
                if isinstance(reduction_fn_purpose, Variable)
                else Variable(reduction_fn_purpose)
            )
        )
        return DeterministicEvaluatorOp.apply(
            prediction,
            target,
            self.eval_fn,
            self.eval_fn_purpose,
            self.success_fn,
            self.reduction_fn,
            self.reduction_fn_purpose,
        )

class ExactMatchEvaluator(Module):
    """
    Evaluates predictions using an exact match criterion.

    This module leverages the [`ExactMatchEvaluator`][afnio.autodiff.evaluator.ExactMatchEvaluator]
    operation from `afnio.autodiff.evaluator` and is a specialized version of the
    [`DeterministicEvaluator`][afnio.cognitive.modules.deterministic_evaluator.DeterministicEvaluator]
    that uses an exact matching function to compare the `prediction` and `target`.
    It returns an evaluation `score` (`1` for exact match, `0` otherwise)
    and an `explanation` describing the evaluation result.

    Examples:
        >>> from afnio import cognitive as cog
        >>> from afnio import set_backward_model_client
        >>> set_backward_model_client("openai/gpt-4o")
        >>> class ExactColor(cog.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         self.exact_match = cog.ExactMatchEvaluator()
        ...     def forward(self, prediction, target):
        ...         return self.exact_match(prediction, target)
        >>> prediction = afnio.Variable(
        ...     data="green",
        ...     role="color prediction",
        ...     requires_grad=True
        ... )
        >>> target = "red"
        >>> eval = ExactColor()
        >>> score, explanation = eval(prediction, target)
        >>> print(score.data)
        0
        >>> print(explanation.data)
        'The evaluation function, designed for 'exact match', compared the <DATA> fields of the predicted variable and the target variable, resulting in a score: 0.'
        >>> explanation.backward()
        >>> prediction.grad[0].data
        'Reassess the criteria that led to the initial prediction of 'green'.'

    Raises:
        TypeError: If the types of `prediction`, `target`, `reduction_fn`, or
            `reduction_fn_purpose` are not as expected.
        ValueError: If the lengths of `prediction.data` and `target` (or `target.data`,
            when `target` is a `Variable`) do not match when both are lists, or if
            `reduction_fn_purpose` (or `reduction_fn_purpose.data`) is an empty string.

    See Also:
        [`afnio.autodiff.evaluator.ExactMatchEvaluator`][afnio.autodiff.evaluator.ExactMatchEvaluator]
        for the underlying operation.
    """  # noqa: E501

    reduction_fn: Optional[Callable[[List[Any]], Any]]
    reduction_fn_purpose: Optional[Union[str, Variable]]

    def __init__(self):
        super().__init__()

        self.register_function("reduction_fn", None)
        self.register_buffer("reduction_fn_purpose", None)

    def forward(
        self,
        prediction: Variable,
        target: Union[str, List[str], Variable],
        reduction_fn: Optional[Callable[[List[Any]], Any]] = sum,
        reduction_fn_purpose: Optional[Union[str, Variable]] = "summation",
    ) -> Tuple[Variable, Variable]:
        """
        Forward pass for the exact match evaluator function.

        Warning:
            Users should not call this method directly. Instead, they should call the
            module instance itself, which will internally invoke this `forward` method.

        Args:
            prediction: The predicted variable to evaluate, which can have scalar or
                list [`data`][afnio.Variable.data] (supporting both individual and
                batch processing).
            target: The target (ground truth) to compare against, which can be a string,
                a list of strings, or a `Variable`.
            reduction_fn: An optional function to aggregate scores across a batch of
                predictions and targets. If `None`, no aggregation is applied.
            reduction_fn_purpose: A brief description of the purpose of `reduction_fn`,
                used by the autodiff engine to generate explanations. Required if
                `reduction_fn` is provided.

        Returns:
            score: A variable containing the evaluation score(s),
                or their aggregation if `reduction_fn` is provided.
            explanation: A variable containing the explanation(s) of the evaluation,
                or their aggregation if `reduction_fn` is provided.

        Raises:
            TypeError: If the types of `prediction`, `target`, `reduction_fn`,
                or `reduction_fn_purpose` are not as expected.
            ValueError: If the lengths of `prediction.data` and `target` (or
                `target.data`, when `target` is a `Variable`) do not match when
                both are lists, or if `reduction_fn_purpose` (or
                `reduction_fn_purpose.data`) is an empty string.
        """
        self.reduction_fn = reduction_fn
        self.reduction_fn_purpose = (
            None
            if reduction_fn_purpose is None
            else (
                reduction_fn_purpose
                if isinstance(reduction_fn_purpose, Variable)
                else Variable(reduction_fn_purpose)
            )
        )
        return ExactMatchEvaluatorOp.apply(
            prediction, target, self.reduction_fn, self.reduction_fn_purpose
        )

class LMJudgeEvaluator(Module):
    """
    Evaluates predictions using a language model (LM) as the judge.

    This module leverages the [`LMJudgeEvaluator`][afnio.autodiff.evaluator.LMJudgeEvaluator]
    operation from `afnio.autodiff.evaluator` to perform model-based evaluations.
    The `forward` method accepts a list of `messages` that construct the evaluation
    prompt, with optional `inputs` to dynamically fill placeholders within message
    templates. A `prediction` is compared against a `target` (optional) to generate
    a `score` and an `explanation`.

    When processing a batch of predictions and targets, `reduction_fn` function
    aggregates individual scores (e.g., using [`sum`][sum] to compute a total score). The
    `reduction_fn_purpose` parameter is a brief description of the aggregation's purpose
    (e.g., `"summation"`). If aggregation is not desired, set `reduction_fn` and
    `reduction_fn_purpose` to `None`. The `success_fn` checks if all evaluations are
    successful, allowing the `backward` pass to skip unnecessary gradient computations.

    This module supports both evaluation (`eval_mode=True`) and optimization
    (`eval_mode=False`) modes.

    The `forward_model_client` specifies the LM responsible for evaluation, while
    `completion_args` allows customization of generation parameters like temperature,
    max tokens, and seed.

    Examples:
        >>> from afnio import cognitive as cog
        >>> from afnio.models.openai import OpenAI
        >>> from afnio import set_backward_model_client
        >>> fwd_model_client = OpenAI()
        >>> fwd_model_args = {"model": "gpt-4o", "temperature": 0.5}
        >>> set_backward_model_client("openai/gpt-4o")
        >>> class Evaluator(cog.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         self.judge = cog.LMJudgeEvaluator()
        ...     def forward(self, fwd_model, messages, prediction, target, inputs, **completion_args):
        ...         return self.judge(fwd_model, messages, prediction, target, inputs, **completion_args)
        >>> task = afnio.Variable(
        ...     "Evaluate if the translation is {metric}.",
        ...     role="evaluation task",
        ...     requires_grad=True
        ... )
        >>> format = afnio.Variable(
        ...     "Provide 'score' (true/false) and 'explanation' in JSON.",
        ...     role="output format"
        ... )
        >>> metric = afnio.Variable(["accurate", "accurate"], role="metric")
        >>> user = afnio.Variable(
        ...     "<PREDICTION>{prediction}</PREDICTION><TARGET>{target}</TARGET>",
        ..      role="user query"
        ... )
        >>> prediction = afnio.Variable(
        ...     ["Hola Mundo", "Salve a tutti"],
        ...     role="translated text",
        ...     requires_grad=True
        ... )
        >>> target = ["Ciao Mondo", "Salve a tutti"]
        >>> messages = [
        ...     {"role": "system", "content": [task, format]},
        ...     {"role": "user", "content": [user]},
        ... ]
        >>> eval = Evaluator()
        >>> score, explanation = eval(
        ...     fwd_model_client,
        ...     messages,
        ...     prediction,
        ...     target,
        ...     inputs={"metric": metric},
        ...     reduction_fn=sum,
        ...     reduction_fn_purpose="summation",
        ...     **fwd_model_args
        ... )
        >>> print(score.data)
        1
        >>> print(explanation.data)
        'The evaluation function, designed using an LM as the judge, compared the <DATA> fields of the predicted variable and the target variable across all samples in the batch. These scores were then aggregated using the reduction function 'summation', resulting in a final aggregated score: 1.'
        >>> explanation.backward()
        >>> system.grad[0].data
        'The translated text should be in Italian.'

    Raises:
        RuntimeError: If the LM response to generate the evaluation `score` and
            `explanation` cannot be parsed as valid JSON.
        TypeError: If the types of `forward_model_client`, `messages`, `prediction`,
            `target`, `inputs`, `success_fn`, `reduction_fn`, `reduction_fn_purpose`,
            or `eval_mode` are not as expected.
        ValueError: If the lengths of `prediction.data` and `target` (or `target.data`,
            when `target` is a `Variable`) do not match when both are lists, or if
            `reduction_fn_purpose` (or `reduction_fn_purpose.data`) is an empty string,
            or if `inputs` contains keys that conflict with `prediction` or `target`.

    See Also:
        [`afnio.autodiff.evaluator.LMJudgeEvaluator`][afnio.autodiff.evaluator.LMJudgeEvaluator]
        for the underlying operation.
    """  # noqa: E501

    forward_model_client: Optional[ChatCompletionModel]
    messages: MultiTurnMessages
    success_fn: Optional[Callable[[List[Any]], bool]]
    reduction_fn: Optional[Callable[[List[Any]], Any]]
    reduction_fn_purpose: Optional[Union[str, Variable]]
    eval_mode: Union[bool, Variable]
    completion_args: Dict[str, Any]

    def __init__(self):
        super().__init__()

        self.register_model("forward_model_client", None)
        self.register_chat("messages", None)
        self.register_function("success_fn", None)
        self.register_function("reduction_fn", None)
        self.register_buffer("reduction_fn_purpose", None)
        self.register_buffer("eval_mode", None)
        self.register_completion_config("completion_args", None)

    def forward(
        self,
        forward_model_client: Optional[ChatCompletionModel],
        messages: MultiTurnMessages,
        prediction: Variable,
        target: Optional[Union[str, List[str], Variable]] = None,
        inputs: Optional[Dict[str, Union[str, Variable]]] = None,
        success_fn: Optional[Callable[[List[Any]], bool]] = None,
        reduction_fn: Optional[Callable[[List[Any]], Any]] = sum,
        reduction_fn_purpose: Optional[Union[str, Variable]] = "summation",
        eval_mode: Union[bool, Variable] = True,
        **completion_args,
    ) -> Tuple[Variable, Variable]:
        """
        Forward pass for the LM Judge evaluator function.

        Warning:
            Users should not call this method directly. Instead, they should call the
            module instance itself, which will internally invoke this `forward` method.

        Args:
            forward_model_client: The LM model client used for the forward
                pass evaluation.
            messages: A list of messages that compose the prompt/context for the LM.
                Each message is a dictionary with a `"role"` (e.g., `"system"`,
                `"user"`, `"assistant"`) and a `"content"` field, which is a list of
                `Variable` objects. The `Variable` objects in the `"content"` can
                contain placeholders (e.g., `{prediction}`, `{target}`) that will be
                populated with the corresponding values from the `inputs` dictionary.
            prediction: The predicted variable to evaluate, which can have scalar or
                list [`data`][afnio.Variable.data] (supporting both individual and
                batch processing).
            target: The target (ground truth) to compare against, which can be a string,
                a list of strings, or a `Variable`. Optional if the evaluation does not
                require a target and only relies on the correctness of the LM Judge's
                assessment of the `prediction`.
            inputs: A dictionary mapping placeholder names to their corresponding
                values, which can be strings or `Variable` instances. These values
                will be used to populate the placeholders in the `messages` content
                before sending the prompt to the LM. For example, if a message
                `"content"` field contains the placeholder `{color}`, the `inputs`
                dictionary should have a key `"color"` with the value to substitute
                in the prompt. Optional if there are no placeholders in the messages or
                if all placeholders are directly related to `prediction` and `target`.
            success_fn: A user-defined function that takes the list of scores returned
                by the LM Judge and returns `True` if all predictions are considered
                successful, or `False` otherwise.
            reduction_fn: An optional function to aggregate scores across a batch of
                predictions and targets. If `None`, no aggregation is applied.
            reduction_fn_purpose: A brief description of the purpose of `reduction_fn`,
                used by the autodiff engine to generate explanations. Required if
                `reduction_fn` is provided.
            eval_mode: Indicates the evaluation mode. If `True`, the `backward` pass
                will compute gradients for the `prediction` variable only. If `False`,
                the `backward` pass will compute gradients for the `messages` and
                `inputs`, allowing optimization of the evaluator itself or alignment
                with human evaluation datasets.
            **completion_args: Additional keyword arguments to pass to the LM model
                client's `chat` method, such as temperature, max tokens, or seed values,
                to customize the LLM's behavior during the evaluation.

        Returns:
            score: A variable containing the evaluation score(s),
                or their aggregation if `reduction_fn` is provided.
            explanation: A variable containing the explanation(s) of the evaluation,
                or their aggregation if `reduction_fn` is provided.

        Raises:
            RuntimeError: If the LM response to generate the evaluation `score` and
                `explanation` cannot be parsed as valid JSON.
            TypeError: If the types of `forward_model_client`, `messages`, `prediction`,
                `target`, `inputs`, `success_fn`, `reduction_fn`,
                `reduction_fn_purpose`, or `eval_mode` are not as expected.
            ValueError: If the lengths of `prediction.data` and `target` (or
                `target.data`, when `target` is a `Variable`) do not match when both are
                lists, or if `reduction_fn_purpose` (or `reduction_fn_purpose.data`) is
                an empty string, or if `inputs` contains keys that conflict with
                `prediction` or `target`.
        """
        self.forward_model_client = forward_model_client
        self.messages = messages
        self.success_fn = success_fn
        self.reduction_fn = reduction_fn
        self.reduction_fn_purpose = (
            None
            if reduction_fn_purpose is None
            else (
                reduction_fn_purpose
                if isinstance(reduction_fn_purpose, Variable)
                else Variable(reduction_fn_purpose)
            )
        )
        self.eval_mode = (
            eval_mode if isinstance(eval_mode, Variable) else Variable(eval_mode)
        )
        self.completion_args = completion_args
        return LMJudgeEvaluatorOp.apply(
            self.forward_model_client,
            self.messages,
            prediction,
            target,
            inputs,
            self.success_fn,
            self.reduction_fn,
            self.reduction_fn_purpose,
            self.eval_mode,
            **self.completion_args,
        )