afnio.cognitive.modules.chat_completion

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,
        )