"""Self-Refine structured output module."""

from typing import List

from pydantic import BaseModel, Field

from agential.core.llm import Response
from agential.prompting.base.output import BasePromptingOutput


class SelfRefineStepOutput(BaseModel):
    """Self-Refine Pydantic output class.

    Attributes:
        answer (str): The answer generated by the llm.
        critique (str): The critique of the answer generated by the llm.
        answer_response (Response): The response of the answer generated by the llm.
        critique_response (Response): The response of the critique generated by the llm.
    """

    answer: str = Field(..., description="The answer generated by the llm.")
    critique: str = Field(..., description="The answer's critique.")
    answer_response: Response = Field(..., description="The answer's response.")
    critique_response: Response = Field(..., description="The critique's response.")


class SelfRefineOutput(BasePromptingOutput):
    """Self-Refine Pydantic output class.

    Attributes:
        additional_info (List[SelfRefineStepOutput]): Additional information about the steps.
    """

    additional_info: List[SelfRefineStepOutput] = Field(
        ..., description="Additional information about the steps."
    )
